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

[zhongbozhu]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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 improves the device-init grouped linear module by adding full support for a single grouped weight/bias parameter (single_grouped_weight / single_grouped_bias) through the native GroupedTensor GEMM backend. It also introduces an explicit grouped_gemm_backend argument to GroupedLinear, deprecates the NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM environment variable, and fixes an MXFP8 scale-address bug in the group_quantize_mxfp8 CUDA kernel for single-tensor uniform-group layouts.

• Single grouped weight/bias: _GroupedLinear._prepare_weights_for_grouped_tensor_gemm and _prepare_bias_for_grouped_tensor_gemm handle both discrete per-expert tensors and a single stacked GroupedTensor parameter; the backward mirrors this by recovering per-group gradients via rowwise_data views.
• MXFP8 weight caching: tex.group_quantize gains an output= reuse argument (with quantizer-identity and shape validation in C++) and a noop_flag path so CUDA-graph-captured weight quantization can be replayed in-place without pointer changes.
• Bug fix – out_features = weights[0].size(-2) replaces .size(0), which incorrectly returned the expert count for a 3-D grouped-weight tensor.

Confidence Score: 4/5

The changes are well-structured and covered by targeted tests; the main risks are the unguarded quantizer-internal flag mutation and test helpers that would silently skip on future GPU generations.

The core forward/backward plumbing for single grouped parameters is logically consistent and the new tests exercise the key paths (BF16, MXFP8, fuse_wgrad, delay_wgrad, return_bias). The weight_quantizer.internal flag is mutated without a try/finally guard — an exception during quantization would leave it in the wrong state for all subsequent calls. The device-capability upper bound <= (11, 0) in both test helpers is tighter than the production gating, meaning tests would silently skip on SM >= 12.x hardware. The CUDA kernel integer division first_logical_dim / num_tensors lacks an assertion, though the invariant is enforced upstream.

transformer_engine/pytorch/module/grouped_linear.py (quantizer-internal exception safety) and the two test helper functions in test_grouped_linear.py and test_grouped_mlp.py (device capability upper bound).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Core module refactored to support a single grouped weight/bias parameter alongside discrete per-GEMM parameters; adds grouped_gemm_backend selection, new weight/bias preparation helpers, corrects out_features extraction (.size(-2) vs .size(0)), and wires single-grouped-weight wgrad through GroupedTensorStorage in the backward.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Ops-level GroupedLinear replaces split_into_quantized_tensors() usage with direct parameter access, adds _get_packed_bias_tensor() helper, and guards the legacy path against single-grouped parameters.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Fixes MXFP8 scale-address computation for single-tensor (SAME_BOTH_DIMS) grouped quantization by deriving per-member row count and base offset rather than treating the whole payload as one block.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Extends group_quantize() to accept an optional output GroupedTensor for in-place reuse; validates quantizer identity and shape compatibility before reusing, enabling CUDA-graph-safe weight caching.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Adds output=py::none() default argument to group_quantize pybind11 binding.
transformer_engine/pytorch/csrc/extensions.h	Adds const py::object &output parameter to group_quantize declaration to match the new implementation signature.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Replaces the inline _get_bias_tensors()/stack() pair with the new _get_packed_bias_tensor() helper for FC2 dbias computation.
tests/pytorch/test_grouped_linear.py	Adds a helper to gate tests on native grouped-tensor GEMM hardware, and new tests covering single-grouped-weight/bias forward+backward correctness, MXFP8 workspace caching, and return_bias semantics.
tests/pytorch/test_grouped_mlp.py	Adds grouped_tensor_path_supported() helper and a test verifying _get_packed_bias_tensor() returns a view of the registered parameter, with a skip guard for unsupported single-grouped-parameter configurations.
tests/pytorch/test_grouped_tensor.py	Adds a comprehensive test covering group_quantize in-place reuse and noop_flag behaviour, including CUDA graph capture/replay to validate the weight-caching contract.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant GroupedLinear_Module
    participant _GroupedLinear_Autograd
    participant PrepareWeights
    participant tex_group_quantize
    participant GroupedGEMM

    Caller->>GroupedLinear_Module: forward(x, m_splits, is_first_microbatch)
    GroupedLinear_Module->>GroupedLinear_Module: resolve grouped_gemm_backend
    GroupedLinear_Module->>GroupedLinear_Module: get weight/bias tensors (single or discrete)
    GroupedLinear_Module->>_GroupedLinear_Autograd: "apply(inp, m_splits, non_tensor_args, *weights, *biases)"

    _GroupedLinear_Autograd->>PrepareWeights: "_prepare_weights_for_grouped_tensor_gemm(single_grouped_weight=True)"
    alt BF16 weight
        PrepareWeights-->>_GroupedLinear_Autograd: GroupedTensorStorage (view)
    else MXFP8 weight, workspace cached
        PrepareWeights->>tex_group_quantize: "group_quantize(src, quantizer, N, noop_flag, output=workspace)"
        tex_group_quantize-->>PrepareWeights: workspace (in-place, same pointers)
        PrepareWeights-->>_GroupedLinear_Autograd: cached GroupedTensorStorage
    else MXFP8 weight, no workspace
        PrepareWeights->>tex_group_quantize: group_quantize(src, quantizer, N)
        tex_group_quantize-->>PrepareWeights: new GroupedTensorStorage
        PrepareWeights-->>_GroupedLinear_Autograd: new workspace (saved to _fp8_workspaces)
    end

    _GroupedLinear_Autograd->>GroupedGEMM: general_grouped_gemm_for_grouped_tensor(weight, x, out)
    GroupedGEMM-->>_GroupedLinear_Autograd: output tensor

    _GroupedLinear_Autograd-->>GroupedLinear_Module: (output, new_workspaces)
    GroupedLinear_Module->>GroupedLinear_Module: update _fp8_workspaces["weight"]

    alt return_bias and single_grouped_bias
        GroupedLinear_Module-->>Caller: (output, bias_GroupedTensor)
    else
        GroupedLinear_Module-->>Caller: output
    end

Loading

Comments Outside Diff (1)

1. transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh, line 786-791 (link)

   Implicit divisibility assumption with no assertion. first_logical_dim / num_tensors uses integer truncation. For the SAME_BOTH_DIMS shape representation all tensors are required to have equal row counts, so the division should always be exact, but there is no compile-time or runtime guard. A silent truncation would produce incorrect per-member scale bases and corrupt the MXFP8 quantization without any visible error. Consider adding a debug-mode assertion such as NVTE_CHECK(first_logical_dim % num_tensors == 0, "...").

_{Reviews (1): Last reviewed commit: "lint" | Re-trigger Greptile}

[greptile-apps]

Upper-bound capability check blocks future architectures. The <= (11, 0) cap means any GPU with SM ≥ 12.x will be unconditionally skipped even if it supports the native grouped-tensor GEMM. The production path in _is_grouped_tensor_path_supported checks (9, 0) <= capability with no upper bound, so the test guard is more restrictive than the implementation. Consider removing or raising the upper bound.

Suggested change

	def _require_native_grouped_tensor_gemm(*, mxfp8: bool = False) -> None:
	"""Skip unless the native GroupedTensor grouped GEMM is available."""
	device_capability = torch.cuda.get_device_capability()
	if not (9, 0) <= device_capability <= (11, 0):
	pytest.skip("Native GroupedTensor grouped GEMM requires Hopper or Blackwell.")
	def _require_native_grouped_tensor_gemm(*, mxfp8: bool = False) -> None:
	"""Skip unless the native GroupedTensor grouped GEMM is available."""
	device_capability = torch.cuda.get_device_capability()
	if not device_capability >= (9, 0):
	pytest.skip("Native GroupedTensor grouped GEMM requires Hopper or newer.")

[greptile-apps]

Same upper-bound issue as in test_grouped_linear.py. The check (9, 0) <= compute_capability <= (11, 0) will return False for any SM ≥ 12.0, causing single_grouped_weight and single_grouped_bias test cases to be silently skipped on future GPU generations even if the backend supports them.

Suggested change

	compute_capability = get_device_compute_capability()
	if not (9, 0) <= compute_capability <= (11, 0):
	return False
	compute_capability = get_device_compute_capability()
	if not compute_capability >= (9, 0):
	return False

[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.

[timmoon10]

These are misleading names.

• Legacy implies we are deprecating it, but in fact it is the default backend and it is the fallback when the grouped tensor impl is not supported.
• As mentioned in https://github.com/NVIDIA/TransformerEngine/pull/3224/changes#r3618050729, it's confusing that we might not actually perform the backend specified by the user. The user might put grouped_gemm_backend="grouped_tensor" and we'll actually do split-quantize because the cuBLAS version is old or something.

Really, this is not a backend config. It's a hint where the user promises to provide m_splits with a certain alignment. I'd suggest changing this to a bool like enable_grouped_tensor_backend. Alternatively if we want to generalize in case we add more backends in the future, we could have a string like enabled_backends (enabled_backends="grouped_tensor,future_backend").