Improve device-init grouped linear module with single grouped weight support

[zhongbozhu]

Description

Fixes numerical issues when using single weight for TE module grouped linear. Limit the single weight feature to the grouped tensor API instead of the legacy path.

TODO: test E2E convergence, unit test directly from Mcore.

Note: needs to pay extra attention to whether bias grad and weight grad are generated properly.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[zhongbozhu]

Note: this code block was deleted because it was duplicated

[zhongbozhu]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR wires single_grouped_weight/single_grouped_bias end-to-end through the native grouped-tensor GEMM path instead of splitting the packed GroupedTensorStorage back into discrete tensors before each forward. It also promotes path-selection logic to a public API (is_module_grouped_tensor_path_supported), replaces the deprecated NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM env variable with a use_grouped_tensor constructor parameter, and fixes the previously flagged meta-device crash in _apply_delay_wgrad_param_hooks.

• Single grouped weight/bias dispatch: _get_weight_tensors() and _get_bias_tensors() now return the packed GroupedTensorStorage directly; new helpers _prepare_weights_for_grouped_tensor_gemm and _prepare_bias_for_grouped_tensor_gemm convert them for cuBLASLt grouped GEMM. The backward assembles a single GroupedTensor wgrad output and returns one gradient per packed parameter.
• CUDA/C++ layer: setup_grouped_gemm_kernel replaces null-pointer sentinels with explicit a_is_discrete/c_is_discrete/d_is_discrete flags; group_quantize_mxfp8_kernel computes per-member scale bases for SAME_BOTH_DIMS tensors; the group_quantize C++ binding gains an output workspace argument for weight caching.
• FP8 current scaling unblocked: make_grouped_weights no longer early-returns for Float8CurrentScaling, fixing the previously identified silent miss of weight registration under quantized_model_init.

Confidence Score: 4/5

The core dispatch and CUDA changes look structurally sound, but the gradient type contract for packed parameters is unverified and the PR description flags end-to-end convergence testing as a TODO.

The backward returns plain tensors (3D view for wgrad, 2D tensor for bias grad) against GroupedTensorStorage parameters. PyTorch is permissive about .grad types, but Megatron/NeMo distributed training hooks may not be. The two near-identical path-support predicates have a subtle behavioral difference that could silently diverge.

Files Needing Attention: transformer_engine/pytorch/module/grouped_linear.py backward path for single_grouped_weight/bias; transformer_engine/pytorch/ops/basic/grouped_linear.py path-support predicate vs. module version.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Core module change: adds is_module_grouped_tensor_path_supported, reworks _get_weight_tensors/_get_bias_tensors to return the packed GroupedTensorStorage directly, rewires forward/backward for single_grouped_weight/single_grouped_bias, adds use_grouped_tensor init-time parameter.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Mirrors the module changes for the op-fuser path: adds is_op_fuser_grouped_tensor_path_supported, replaces _is_graph_safe_path_supported, adds _get_discrete_bias_tensors/_get_packed_bias_tensor, fixes meta-device _apply_delay_wgrad_param_hooks. Logic is near-identical to the module version but missing the backward_override guard.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Adds tensor_rows_for_scales for SAME_BOTH_DIMS shape representation so GEMM-swizzled scales are computed per member rather than for the whole tall tensor; also fixes the VARYING_FIRST_DIM offset path.
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu	Replaces has_*_multi_tensor null-pointer sentinels with explicit a_is_discrete/c_is_discrete/d_is_discrete bool flags; handles null-data zero-sized allocations and moves scale-null guard.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds output parameter to group_quantize for workspace reuse with quantizer identity validation; also wires noop_flag through for the FP8 block-scaling path.

_{Reviews (21): Last reviewed commit: "fix cublas setup meta data issue for emp..." | Re-trigger Greptile}

[timmoon10]

The biggest change in this PR is that TE is abandoning any attempt to make single_grouped_weight=True a general feature. Things must be exactly right, or we crash. Given how delicate and experimental this feature has been, I'm not opposed.

The second change is that users must opt-in to access the grouped GEMM kernel. This is also reasonable, since it has alignment requirements for m_splits and it's helpful having a way for users to accept that stricter contract.

We are experiencing many test failures. Given that single_grouped_weight is no longer a general feature, I think it's reasonable we move the corresponding tests to test_grouped_linear.py and test_grouped_mlp.py.

[timmoon10]

I get that the h2d memcpy is suboptimal, but it's trivially easy to handle. Erroring out seems excessively rigid.

Suggested change

	raise ValueError(
	"The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 "
	"tensor, or select grouped_gemm_backend='legacy'."
	)
	m_splits = m_splits.to(device=device)

We need to handle the d2h case anyways when the user has specified grouped_gemm_backend="grouped_tensor", but it's not supported and we fallback to split-quantize.

[zhongbozhu]

But if the alignment is not provided in the first place, converting it to a device tensor also wouldn't work right, I am okay with another alignment check before adding this H2D copy.

[zhongbozhu]

/te-ci pytorch

[zhongbozhu]

/te-ci pytorch

[zhongbozhu]

Note: this is for single weight quantize for mxfp8.

Before this change, the weight quantizer didn't have the first_dims because moe weights are uniform shape for both dimension. This will then trigger a CUDA illegal access because offsets_ptr=nullptr

[zhongbozhu]

/te-ci pytorch

[zhongbozhu]

/te-ci pytorch

[zhongbozhu]

/te-ci pytorch