Support row-only MXFP8 distributed master-weight casts

[xiuhu17]

Description

Support distributed master-weight writeback into MXFP8 primary weights whose columnwise data and scales are absent.

Optimize memory usage for: radixark/miles#615
Related initialization change: #3468
Related Megatron integration: NVIDIA/Megatron-LM#7095

The initialization change permits row-only primary storage for high-precision/dequantized backward overrides. The existing distributed cast path nevertheless dereferences columnwise scales and writes columnwise data. This PR makes the native cast path handle that layout directly, independently of how the primary was initialized.

Changes

• Pack/reduce amax only for storage directions that exist, including batches mixing row-only and bidirectional weights.
• Skip absent columnwise scale updates and output writes in the Python helper and native CUDA partial-amax/partial-cast kernels.
• Keep the existing single packed-amax MAX all-reduce. No high-precision parameter/tile reconstruction collective is introduced.
• Preserve existing rowwise storage, rounding through the model dtype, and the bidirectional path.
• Support both in-place primary writes and the helper's FSDP fragment-output interface. This is not a claim of complete FSDP training integration coverage.
• Keep native function signatures unchanged: empty [0, 0] columnwise amax/scales and [0] output represent omission. Detect this by shape, since an empty view may still have a backing pointer.
• Test both public entry points, two consecutive updates, mixed layouts, a shard boundary inside a 32-value block, ranks without master shards, storage pointer stability, and direct partial kernels with empty input/views.

Scope

MXFP8 distributed cast only. No initialization policy changes, TMS/backup integration, NVFP4 changes, or Megatron-specific fallback are included. Kernel support can replace the high-precision tile-reconstruction fallback discussed in Megatron PR #7095; Megatron's separate gather copy-back layout handling remains relevant.

Validation

All checks below ran on training_gb200_dev (NVIDIA GB200), using a source-built TE wheel in an isolated environment. The Python package and native extension paths were verified to point to that build, not a preinstalled TE. Build: CUDA 13.3; runtime Torch: 2.13.0+cu130; Python 3.14.6.

• Single GPU: 9 passed across tests/pytorch/mxfp8/test_rowwise_master_cast.py, tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py, and tests/pytorch/test_partial_cast.py.
• Four GPUs: torchrun --standalone --nproc_per_node=4 -m pytest tests/pytorch/mxfp8/test_rowwise_master_cast.py: 7 passed on each rank, process exit code 0.
• Exact native data/active-scale comparisons pass, including mixed layouts, fragment outputs, unaligned/empty shards and consecutive writes through both helper entry points.
• Python syntax checks, new-test Black check, and cpplint on both modified C++ files pass.

The distributed test fixture shares one process group across its module, avoiding repeated NCCL teardown/reinitialization between parameterized cases. Environment warnings were present for Python 3.14 TorchScript deprecation and unavailable optional FlashAttention 4/CUTLASS; these tests do not use FlashAttention.

Full Megatron training, checkpoint/resharding, end-to-end FSDP, and performance measurements are outside the validation performed for this PR.

Type of change

• Bug fix

Checklist

• I have read the contributing guidelines
• I have commented non-obvious behavior and updated native API documentation
• I have added regression tests
• New and relevant existing GPU tests pass

[greptile-apps]

Greptile Summary

This PR extends distributed MXFP8 master-weight casting to support row-only primary storage while preserving the existing bidirectional path.

• Packs and reduces amax values only for available storage directions.
• Uses explicit empty tensor shapes to communicate omitted columnwise work across the Python/native boundary.
• Guards columnwise scale loads, updates, and output writes in the CUDA partial kernels.
• Adds numerical coverage for repeated updates, mixed layouts, fragment outputs, unaligned shards, empty shards, and direct empty-view kernel calls.
• The new test is collected for single-GPU validation, but its multi-rank branches are not currently wired into the repository’s distributed CI suite.

Confidence Score: 4/5

The implementation appears safe to merge, with a non-blocking gap in automated multi-rank regression coverage.

The row-only omission contract is consistently implemented across amax packing, scale updates, Python fragment selection, native validation, and guarded CUDA accesses. The remaining concern is that repository CI runs the new test only as a single process, leaving the changed cross-rank behavior unprotected.

Files Needing Attention: tests/pytorch/mxfp8/test_rowwise_master_cast.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/utils.py	Adds direction-aware amax packing, scale updates, and fragment selection for row-only and mixed MXFP8 master-weight casts.
transformer_engine/common/recipe/mxfp8_scaling.cu	Detects shape-based columnwise omission and guards all omitted scale accesses and output writes in the partial CUDA kernels.
transformer_engine/common/include/transformer_engine/recipe.h	Documents the empty-shape convention used to omit columnwise computation without changing native signatures.
tests/pytorch/mxfp8/test_rowwise_master_cast.py	Provides strong focused coverage, but its distributed branches are not registered with the multi-rank CI suite.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  M[Distributed master-weight shard] --> A[Compute partial rowwise amax]
  M --> C{Columnwise storage present?}
  C -->|Yes| CA[Compute partial columnwise amax]
  C -->|No| O[Use empty 0x0 omission marker]
  A --> P[Pack available amax values]
  CA --> P
  O --> P
  P --> R[Single MAX all-reduce]
  R --> S[Update available inverse scales]
  S --> RW[Cast rowwise fragment]
  C -->|Yes| CW[Cast columnwise fragment]
  C -->|No| SKIP[Skip columnwise output]

Loading

_{Reviews (1): Last reviewed commit: "Share the NCCL process group across row-..." | Re-trigger Greptile}

[greptile-apps]

Distributed Paths Lack CI

The L0 MXFP8 suite collects this test using single-process pytest, while the L1 distributed suite does not select or launch it. As a result, the rank-partitioned branches—including the unaligned shard boundary, empty tail ranks, and mixed row-only/bidirectional layout—are not exercised in repository CI. This is non-blocking, but future regressions in the distributed behavior changed here could go undetected. Please register this test with the distributed suite or add equivalent self-launched multi-rank coverage.

Knowledge Base Used: Verification and CI matrix

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[ptrendx]

Checking for empty 2D shape is not great. Just check for has_data being false (empty NVTETensor).

[ptrendx]

Same comment as above.

[ptrendx]

If there is an issue of the pointer passed from the PyTorch side, then that should be handled in the PyTorch extension rather than WARing it here.

[ptrendx]

Please include some benchmark results on the old vs new time for the case still using columnwise output to see whether those additional conditions didn't introduce overhead. If they did then they would need to be handled via templated kernels.

[ptrendx]

Why do we need a separate test file for that? It should be part of one of the larger suites like test_sanity or test_numerics. Also, for the separate test files, they would not actually be run unless they are part of the test.sh in one of the directories in qa/.