Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

Numerical accuracy issue with Triton MXFP4 kernel with respect to quantization

Aperta
  • Qualcosa non va?
#665 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
55/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Tranquilla
Stack tecnologico
python
Ambito
machine-learning, performance

Direzione di ricerca

Start in transformer_engine/pytorch/triton_kernels/cast_transpose.py, reading te_cast_transpose_mxfp4_triton and _cast_transpose_triton_mxfp4 alongside the attached reproduction script. Compare Triton and native outputs with shuffle_columnwise_data enabled and disabled. Done means the Triton columnwise FP4 data changes with the flag and matches the native preshuffled layout, while the existing scale-swizzle behavior remains correct.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

bug

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.
Lingua principale
Python
Stelle
77
Fork
36
Merge medio
3g 16h
PR unite (30g)
10

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di ROCm/TransformerEngine

Tutte le issue di ROCm/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.