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 improves the device-init grouped linear module by adding full single grouped weight/bias support restricted to the grouped tensor API, fixing numerical correctness issues where the old code allowed the single-weight path to silently fall through to the incompatible split-quantize path.

• Grouped tensor path gating: introduces is_module_grouped_tensor_path_supported / is_op_fuser_grouped_tensor_path_supported predicates (now exported) that subsume the old private static method and add a backward_override check, with a RuntimeError guard in fuser_forward blocking single_grouped_weight/bias when use_grouped_tensor_path is false.
• MXFP8 kernel fix: group_quantize_mxfp8.cuh corrects column-wise scale computation for SAME_BOTH_DIMS shape representation by dividing first_logical_dim by num_tensors to get per-member row strides; cublaslt_grouped_gemm.cu adds a zero-row early-return to launch_grouped_bias_add before accessing data pointers.
• group_quantize output reuse: the C++ extension gains an output parameter enabling in-place update of pre-allocated GroupedTensor storage, making CUDA-graph-safe weight caching possible across microbatches.

Confidence Score: 4/5

Safe to merge with one minor concern about an implicit divisibility assumption in the MXFP8 kernel.

The PR is a well-scoped correctness fix with extensive new test coverage (8+ new tests). The numerical correctness issues it addresses — silent fall-through to the split-quantize path, zero-row null-pointer crash, and incorrect MXFP8 column-wise scale strides — are all properly fixed. The only remaining concern is the integer division first_logical_dim / num_tensors in the MXFP8 kernel which lacks a divisibility assertion; in practice the grouped linear module always creates same-shape experts so this is unlikely to trigger, but it is a latent silent-correctness hazard.

Files Needing Attention: transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh — the new tensor_rows_for_scales calculation lacks a divisibility guard.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Major refactor adding single grouped weight/bias support to the module-level GroupedLinear; introduces is_module_grouped_tensor_path_supported, _prepare_weights/bias_for_grouped_tensor_gemm, make_grouped_weights, backward_dw single-weight path, and checkpoint remap helpers. Logic is well-structured with good test coverage.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Adds single grouped weight/bias support to the ops-level GroupedLinear fuser; exports is_op_fuser_grouped_tensor_path_supported; fixes _apply_delay_wgrad_param_hooks meta-device crash via _parameters.get(); grad_params ordering intentionally matches PyTorch ordered _parameters dict.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Fixes column-wise scale computation for SAME_BOTH_DIMS / GEMM-swizzled single grouped weight; uses first_logical_dim/num_tensors as per-member row stride with no divisibility assertion — silent truncation is possible if row counts are not exact multiples.
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu	Adds total_rows==0 early-return guard to launch_grouped_bias_add before any data-pointer dereference, correctly fixing the null-pointer crash for zero-token batches.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds output-reuse path to group_quantize: validates num_tensors and logical_shape match, propagates noop_flag for FP8 block-scaling path, enabling CUDA-graph-safe in-place weight caching.
transformer_engine/pytorch/csrc/extensions.h	Adds const py::object &output parameter to group_quantize declaration; straightforward signature update matching the implementation change.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers py::arg("output") = py::none() for group_quantize binding; backward-compatible default keeps all existing call sites working.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Minor updates to adopt is_op_fuser_grouped_tensor_path_supported; no substantive logic changes to the MLP fuser itself.
transformer_engine/pytorch/ops/basic/init.py	Exports the new is_op_fuser_grouped_tensor_path_supported predicate; trivial change.
transformer_engine/pytorch/module/init.py	Exports the new is_module_grouped_tensor_path_supported predicate; trivial change.
tests/pytorch/test_grouped_linear.py	Adds 8 new tests covering zero-work bias, host m_splits rejection, numerical equivalence of single vs discrete grouped tensor path, workspace caching, swizzle bypass, primary MxFP8 bypass, parameter layout, and grouped bias return; replaces SM-cap guards with is_module_grouped_tensor_path_supported.
tests/pytorch/test_grouped_mlp.py	Adds test for single grouped bias packed storage; replaces SM capability guards with is_op_fuser_grouped_tensor_path_supported predicate.
tests/pytorch/test_grouped_tensor.py	Adds two tests for group_quantize output-reuse and noop semantics; loosens one tolerance from exact equality to rtol=1e-5/atol=4e-3 for a numerical stability fix.
tests/pytorch/test_sanity.py	Adds shape alignment checks, use_grouped_tensor flag, cuda synchronize, and empty-split+FP8 skip guards to test_sanity_grouped_linear; also adds synchronize to test_sanity_linear_with_zero_tokens.

Sequence Diagram

sequenceDiagram
    participant User as Caller
    participant GL_mod as GroupedLinear (module)
    participant GL_ops as GroupedLinear (ops)
    participant PrepW as _prepare_weights_for_grouped_tensor_gemm
    participant GQ as group_quantize (C++ ext)
    participant GEMM as grouped_gemm kernel
    participant BiasAdd as launch_grouped_bias_add

    User->>GL_mod: forward(inp, m_splits, single_grouped_weight)
    GL_mod->>GL_mod: validate single_grouped_weight requires use_grouped_tensor
    GL_mod->>GL_ops: fuser_forward(...)
    GL_ops->>GL_ops: check is_op_fuser_grouped_tensor_path_supported()
    alt "single_grouped_weight=True"
        GL_ops->>PrepW: "_get_grouped_weight_for_gemm(single=True)"
        PrepW->>GQ: "group_quantize(..., output=cached_workspace)"
        GQ-->>PrepW: GroupedTensor (quantized, in-place reuse)
    else discrete weights
        GL_ops->>PrepW: "_prepare_weights_for_grouped_tensor_gemm(single=False)"
        PrepW->>GQ: group_quantize(weight_i...)
        GQ-->>PrepW: GroupedTensor
    end
    PrepW-->>GL_ops: grouped weight tensor
    GL_ops->>GEMM: general_grouped_gemm_for_grouped_tensor(...)
    GEMM-->>GL_ops: output activations
    GL_ops->>BiasAdd: launch_grouped_bias_add(...)
    Note over BiasAdd: early-return if total_rows==0
    BiasAdd-->>GL_ops: biased output
    GL_ops-->>GL_mod: output
    GL_mod-->>User: output

    User->>GL_mod: backward(grad_output)
    GL_mod->>GL_ops: fuser_backward(...)
    alt "single_grouped_weight=True"
        GL_ops->>GL_ops: compute wgrad (single grouped)
        GL_ops->>GL_mod: grad_weights (stacked or rowwise_data view)
        GL_mod->>GL_mod: set self.weight.grad
    else discrete weights
        GL_ops->>GL_mod: grad_weights list
        GL_mod->>GL_mod: set weight_i.grad for each expert
    end
    GL_ops-->>GL_mod: (dinp, dweights, dbiases)
    GL_mod-->>User: dinput

Loading

_{Reviews (20): Last reviewed commit: "add another specific case guard" | 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