Numerical accuracy issue with Triton MXFP4 kernel with respect to quantization

[amital-amd]

Summary

The Triton MXFP4 cast-transpose path (te_cast_transpose_mxfp4_triton, selected by
NVTE_USE_CAST_TRANSPOSE_TRITON=1) does not honor the quantizer's shuffle_columnwise_data /
shuffle_rowwise_data flags. It applies the scale swizzle (with_gemm_swizzled_scales) but
stores the packed FP4 data in the plain transposed layout, ignoring the data-preshuffle flags
entirely.

The native/HIP path (tex.quantize → te_mxfp4::cast_transpose_mxfp4_shuffled) does apply the data
preshuffle. So when a recipe sets shuffle_columnwise_data=True (the real MXFP4 training layout),
the Triton kernel's columnwise operand is unshuffled and does not match the preshuffled layout
the AITER a4w4 GEMM consumes (bpreshuffle=True).

Because the columnwise operand feeds the backward weight-gradient (wgrad) GEMM, the FP4 values
are correct but their ordering is wrong. This produces structurally-wrong weight gradients that
compound over many steps into a slow grad-norm blow-up (not an immediate NaN).

Reproduction

Attaching a script which helps in reproduction of bug. triton_mxfp4_ignores_columnwise_shuffle.py

The attached script reproduces the bug on a single GPU (random bf16 tensor, no training). It
quantizes the same tensor with both kernels and compares the columnwise FP4 data with the shuffle
flag on and off: the native output changes when the flag flips (it preshuffles), while the Triton
output stays byte-identical (it ignores the flag).

Expected vs. actual

• Expected: with shuffle_columnwise_data=True, the Triton kernel's columnwise FP4 data is
  preshuffled identically to the native kernel.
• Actual: the Triton kernel emits identical columnwise data whether the flag is True or False —
  i.e. it ignores the flag and leaves the data unshuffled.

Environment

• GPU: AMD Instinct MI355X (gfx950), ROCm 7.x
• Image (reproduces the bug): rocm/primus:pytorch-2.12.0-rocm7.14.0a20260608_te-2.14.0.dev0-e6ede467_v26.4

Root cause (code pointers)

In transformer_engine/pytorch/triton_kernels/cast_transpose.py, the launcher only ever reads the
scale-swizzle flag; the data-shuffle flags are never consulted:

# te_cast_transpose_mxfp4_triton
shuffle_scales = out_metadata["quantizer"].with_gemm_swizzled_scales
SHUFFLE_ROWWISE_SCALING = shuffle_scales and USE_ROWWISE_SCALING
SHUFFLE_COLWISE_SCALING = shuffle_scales and USE_COLWISE_SCALING
# quantizer.shuffle_columnwise_data / shuffle_rowwise_data are never read.

Inside the kernel (_cast_transpose_triton_mxfp4), the columnwise branch stores FP4 data with a
plain transposed store and then swizzles only the scales:

if USE_COLWISE:
    packed_col, bs_col = _mxfp4_quantize_32x32_block(tl.trans(x_chunk), MXFP4_BLOCK_SIZE)
    tl.store(colwise_fp4_ptr + ..., packed_col, mask=col_fp4_mask)  # plain transpose, no data preshuffle
    if SHUFFLE_COLWISE:
        ...  # shuffles SCALES only

The native C++/HIP path reads the same flags and applies the data preshuffle in-kernel, so the two
kernels diverge only when a data-shuffle flag is set.

Impact

With NVTE_USE_CAST_TRANSPOSE_TRITON=1 and a recipe that sets shuffle_columnwise_data=True, wgrad
reads a misordered FP4 operand. In a Flux-12B MXFP4 run (SR off) this turns a converging run
(loss ~0.63, grad norm ~0.2) into a diverging one (grad norm climbing to ~58 over ~1000 iters),
while the native kernel converges under identical settings.

Additional: the Triton path is also slower

Independently of the correctness bug above, the Triton cast-transpose is structurally slower than the
native kernel, for these reasons:

1. Two quantizes plus an in-kernel transpose per tile. The columnwise operand is produced by
   physically transposing each tile (tl.trans) and re-running the full quantizer, reusing none of
   the rowwise work. The native kernel emits both the rowwise and columnwise layouts in a single
   fused pass.
2. Expensive per-element encoding. The E8M0 scale is computed with transcendentals
   (log2/floor/exp2) instead of reading the exponent bits directly, and the E2M1 data uses a
   dependent chain of tl.where comparisons instead of the gfx950 hardware packed-FP4 cast.
3. Not autotuned. The kernel uses a fixed tile size and default launch parameters with generic
   codegen, so it under-uses the GPU and cannot adapt to different tensor shapes. The native kernel
   is hand-tuned for gfx950.