[matthiasdiener]

Description

Followup of #723.

Extends the MXFP4 recipe to route te.GroupedLinear through the grouped GEMM MXFP4/a8w4 Triton kernels (#723) on gfx950, and adds a8w4 support (fwd-only) for K3.

Performance results: https://github.com/AMD-ROCm-Internal/frameworks-internal/issues/371#issuecomment-5879699372

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

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

[Copilot]

Copilot review overview

🟡 Changes recommended

The fast-path dispatch breaks supported buffer and gradient-handling modes and lacks sufficient layout and shape validation.

Review effort: Balanced
Findings: 2 · 2

Open (4)

• Dispatch ignores shuffled, scaled, and Hadamard quantizer layouts · New
• Raw Triton call lacks input shape and dtype validation · New
• New file uses an invalid copyright header · New
• Docstring incorrectly describes the helper as forward-only · New

What changed in this PR

Extends ROCm gfx950 MXFP4 support to GroupedLinear, including a8w4 forward inference and a4w4 training paths.

Changes:

• Adds prequantized grouped Triton GEMM dispatch.
• Adds MXFP8 E4M3 reference quantization and recipe factories.
• Adds Linear and GroupedLinear coverage for a8w4/a4w4.

File	Description
grouped_gemm_mxfp4_impl.py	Adds prequantized grouped GEMM entry point.
grouped_linear.py	Integrates MXFP4/a8w4 fast paths and backward dispatch.
quantization_mxfp4_grouped.py	Defines reference quantizers and recipe factories.
test_a8w4_recipe.py	Tests a8w4 Linear forward and backward.
test_a8w4_grouped_fast.py	Tests grouped kernels and GroupedLinear integration.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Copilot review overview

🟡 Changes recommended

The optimized path mishandles supported output, gradient, cache, empty-input, and quantization-parity cases.

Review effort: Balanced
Findings: 1 · 2

Open (3)

• Handle zero-sized grids for empty MoE routing batches · New
• Validate cached workspace type and shape before reuse · New
• Align zero-block MXFP8 encoding with the reference · New

Previously missed (2)

In code that hasn't changed since last review

Honor preallocated out and dgrad_out buffers

transformer_engine/​pytorch/​module/​grouped_linear.py:1216

The optimized path returns before handling the public out and dgrad_out arguments, so callers that supply preallocated buffers no longer receive aliased output or dgrad as documented. Either plumb both buffers through the custom kernels or fall back whenever either is provided.

Support fused and delayed wgrad modes in the optimized path

transformer_engine/​pytorch/​module/​grouped_linear.py:1216

This early path ignores fuse_wgrad_accumulation and delayed wgrad (wgrad_store), although both are supported GroupedLinear modes. Its backward returns ordinary parameter gradients immediately, so callers requesting fused accumulation will not get main_grad updates and delayed-wgrad callers will not have work queued. Gate this route for those modes or implement the same handling as _backward_blockwise_fp8_triton.

[Copilot]

Copilot review overview

🟡 Changes recommended

Scale-layout ambiguity, skipped validation, and unsupported GroupedLinear options can produce incorrect behavior.

Review effort: Balanced
Findings: 1 · 2

Open (3)

• Honor buffer and gradient semantics in early-return path · New
• Disambiguate scale layout when dimensions are equal · New
• Validate splits and operands before empty-batch return · New

Resolved since last review (3)

• Handle zero-sized grids for empty MoE routing batches
• Align zero-block MXFP8 encoding with the reference
• Validate cached workspace type and shape before reuse

[Copilot]

Copilot review overview

🟡 Changes recommended

Cache namespace collisions, device-insensitive architecture gating, and fused-quantization underflow can cause failures or incorrect results.

Review effort: Balanced
Findings: 1

Open (1)

• Check operand device instead of current CUDA device · New

Resolved since last review (1)

• Shared weight0 cache leaks incompatible raw tuples across recipe calls

Previously missed (1)

In code that hasn't changed since last review

Clamp tiny block scales to preserve nonzero values

transformer_engine/​pytorch/​triton_kernels/​cast_transpose.py:909

The reused cast kernel does not match the documented reference for every nonzero block. When amax / 448 underflows into FP32's subnormal range, float_to_e8m0_triton emits scale byte 0 and exp2f_rcp_triton treats that byte as an inverse scale of 1, so tiny BF16 values quantize to zero. The torch reference clamps the exponent to -127 and rescales first (for example, amax=2^-126 becomes E4M3 value 2), preserving the value. This makes the new fused path numerically differ for nonzero low-magnitude blocks; use a downcast path that applies the clamped -127 inverse scale for this case.

[Micky774]

There are a couple of tests that fail to collect -- also the new test files still need to be added to the CI script right?

[matthiasdiener]

There are a couple of tests that fail to collect

Thanks, this is fixed in c4c0e64

-- also the new test files still need to be added to the CI script right?

The CI test script already runs all test files in the mxfp4/ folder, so it should be fine:

TransformerEngine/ci/pytorch.sh

Line 129 in 423347c

run_default_fa 1 mxfp4/

[aris134]

Since we now have the fp8 blocked as well as mxfp4 implementations sharing this func, can we move it into a shared utility file like common.py?

[matthiasdiener]

Thanks, done in b4e9e19