[Common] Fix Grouped MXFP8 work mapping and TMA synchronization

[Oleg-Goncharov]

Description

This PR fixes three correctness issues in the Grouped MXFP8 quantization kernels introduced by the optimized work mapper:

1. the direct VARYING_FIRST_DIM path assumed that every tensor's first dimension was aligned to the 128-row kernel tile, but only the combined logical dimension was validated. This could allow a CTA to cross an expert boundary and use incorrect per-expert metadata.

2. device-modified TMA descriptors were copied from shared to global memory without the release/acquire protocol required by the TMA memory model. The descriptor updater now follows the recommended warp-collective publication sequence.

3. the direct mapper stored per-expert metadata in dynamic shared memory that was also used as the destination of asynchronous TMA loads. A TMA load could therefore overwrite the metadata before every warp had consumed it, resulting in corrupted tensor metadata and invalid global-memory addresses. The mapper metadata is now kept in dedicated static shared memory, and the row offset is computed relative to the current expert.

Addresses #3474

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:

• Restore validation that every VARYING_FIRST_DIM tensor has a first dimension divisible by the 128-row kernel tile.
• Store direct-mapper metadata in dedicated static shared memory so that it cannot be overwritten by asynchronous TMA loads.
• Compute the mapper row offset relative to the current expert.
• Align shared-memory CUtensorMap storage to 128 bytes.
• Publish modified TMA descriptors with a warp-collective tensormap.cp_fenceproxy operation using GPU-scope release semantics.
• Use the matching GPU-scope acquire fence before consuming descriptors from another CTA.
• Launch the descriptor updater with a full warp and update the grouped quantize and dequantize paths accordingly.

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

[greptile-apps]

Greptile Summary

This PR corrects Grouped MXFP8 direct work mapping and TMA descriptor synchronization.

• Restores per-tensor row-alignment validation for the direct VARYING_FIRST_DIM mapper.
• Moves mapper metadata out of dynamic shared memory used by asynchronous TMA transfers.
• Computes colwise swizzled row offsets relative to the current expert.
• Publishes modified TMA descriptors collectively from a full warp with GPU-scope release/acquire synchronization.

Confidence Score: 5/5

The PR appears safe to merge; the previously identified incomplete-warp descriptor publication is corrected and no new actionable failure remains.

The updater now executes descriptor publication with a full warp, pairs GPU-scope release and acquire operations across CTAs, and prevents asynchronous TMA traffic from corrupting mapper metadata. The only relevant change since the previous review is formatting of the corrected launch.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/core/grouped_tma.cuh	Publishes aligned, device-modified TMA descriptors with a full-warp GPU-scope release/acquire protocol.
transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh	Makes descriptor modification warp-collective and launches the updater with a full warp.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Isolates direct-mapper metadata, restores per-tensor validation, fixes expert-relative row indexing, and launches descriptor updates with a full warp.

Sequence Diagram

sequenceDiagram
    participant Host
    participant U as Descriptor updater CTA
    participant G as Global CUtensorMap
    participant Q as Quantize/dequantize CTA

    Host->>U: Launch one full warp per tensor
    U->>U: Copy and modify aligned shared descriptor
    U->>G: Warp-collective copy with release.gpu
    Host->>Q: Launch grouped kernel in same stream
    Q->>G: acquire.gpu tensor-map fence
    G-->>Q: Published descriptor
    Q->>Q: Resolve expert and expert-relative row
    Q->>Q: Execute TMA operation

Loading

_{Reviews (5): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[Oleg-Goncharov]

/te-ci

[janbernloehr]

This PR fixes the reported MXFP8 MoE illegal-memory-access regression in #3474.

We reran the reproducer using the current PR head (be88f0c), applied to the exact first-bad TransformerEngine revision (c3066f8), and confirmed that the patched sources were rebuilt.
The workload was Qwen3-30B-A3B Megatron Bridge/PyTorch FP8 training on eight GB200 GPUs (2 nodes × 4 GPUs), with sequence length 4096, batch size 4, TP=1, PP=1, CP=1, EP=8, ETP=1, and hybrid expert-parallel dispatch.
With the patch, the workload:

• completed three MXFP8 MoE training iterations at 107.7 model TFLOP/s/GPU;
• completed full-iteration CUDA graph capture; and
• produced no illegal-address or other CUDA-error hits across any of the eight ranks.