[Bug] Illegal memory access in grouped MXFP8 quantize on GB200 (sm_100) introduced by #3337; TMA store completion handoff missing in group_quantize_mxfp8.cuh

[janbernloehr]

Describe the bug

Since #3337 ("[Common] Improved performance of Group MXFP8 kernels"), MXFP8 MoE pretraining that uses the TransformerEngine fused grouped-MLP path aborts with CUDA error: an illegal memory access was encountered (cudaErrorIllegalAddress) during the first backward pass on GB200 (sm_100, aarch64, 8 GPUs over 2 nodes).

We bracketed this to an exact commit pair by building TransformerEngine from source at each side and running the identical workload:

• last known good: 3231a150ef56dd7e6600f59c6a3ce19d64936b7c (the first parent of the bad commit) — no illegal memory access on any rank; two full training iterations complete, including CUDA graph capture.
• first bad: c3066f8522613bb3077d8bd676ca3b8cb345f49c (the release_v2.19 cherry-pick of [Common] Improved performance of Group MXFP8 kernels #3337; main counterpart 89300522a3be973c7f2e67272ce12c61dcfae80e) — reproduces the illegal memory access inside iteration 0, zero completed steps.

Both builds were produced from the same container base and the same pinned versions of every other component. Between the two runs the only material difference was transformer_engine (2.19.0+3231a15 vs 2.19.0+c3066f8): torch and torchvision were built from identical source commits, and CUDA (13.4.1.012), cuDNN (9.26.0.37), NCCL (2.31.2), CUTLASS-DSL (4.6.1) and the driver (580.178.04) were byte-identical. A per-package/per-library environment diff of the two jobs reported no differing shared libraries and no differing system information.

Steps/Code to reproduce bug

We do not yet have a standalone single-kernel reproducer; the failure was reproduced through a full MoE pretraining step. Minimal steps in public terms:

1. Build TransformerEngine from source at c3066f8522613bb3077d8bd676ca3b8cb345f49c (or any release_v2.19 revision at or after it) against a recent PyTorch + CUDA 13.4 aarch64 toolchain.
2. Run a Megatron-Bridge / Megatron-Core pretraining job for a Qwen3-30B-A3B-shaped MoE with mock data on 8 GB200 GPUs across 2 nodes, with:
   • MXFP8 compute dtype (fp8_mx),
   • expert parallelism 8, tensor/pipeline/context parallel 1, expert-tensor parallel 1,
   • sequence length 4096, micro-batch size 4,
   • the TransformerEngine fused grouped MLP path enabled: NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 together with Megatron-Core's use_transformer_engine_op_fuser=True (both are set for us by Megatron-Bridge's GB200 preset helper _enable_hybridep_full_iteration_mxfp8).
3. The job aborts in the first backward pass with the trace shown below. Rebuilding only TransformerEngine at 3231a150ef56dd7e6600f59c6a3ce19d64936b7c and changing nothing else makes the illegal memory access disappear.

Because the reuse race is timing- and backpressure-dependent, a reproducer is more likely to fire with many expert groups and with device memory close to fully reserved.

Expected behavior

MXFP8 MoE training with the fused grouped-MLP path should run without a CUDA illegal memory access, as it does at 3231a150ef56dd7e6600f59c6a3ce19d64936b7c, where the same job completes iterations with step times of ~162 s.

Actual behavior

The job dies inside the first training iteration's backward pass with a CUDA illegal memory access; zero training steps complete (no step-time line is ever emitted on any rank). Concretely, on the first-bad commit:

• The last framework activity before the abort is inside Variable._execution_engine.run_backward(...), immediately preceded by transformer_engine/pytorch/quantized_tensor.py:161: UserWarning: Quantizer is being updated.
• The fault is reported asynchronously — it is not raised at the offending kernel launch, but later by the NCCL watchdog polling a CUDA event, so the reported frame is not the faulting kernel:

[rank1] [PG ID 16 PG GUID 88(EXPERT_TENSOR_AND_MODEL_PARALLEL_GROUP) Rank 1] Process group watchdog thread
  terminated with exception: CUDA error: an illegal memory access was encountered
Exception raised from query at c10/cuda/CUDAEvent.h:111
frame #3: c10d::ProcessGroupNCCL::WorkNCCL::isCompleted()
frame #4: Watchdog::runLoop()
frame #5: Watchdog::run()
terminate called after throwing an instance of 'c10::DistBackendError'

• The illegal memory access is reported on a subset of ranks (2 of the 8 ranks report it, twice each; the remaining ranks report none and are torn down by the same watchdog), which is consistent with a timing-dependent race rather than a deterministic out-of-bounds index.
• The process then aborts with SIGABRT (signal 6, core dumped), and the job exits with status 134.
• On the last-known-good commit the same run is clean: zero illegal memory accesses on every rank, two completed iterations, CUDA graph capture succeeds. (It later hits a plain torch.OutOfMemoryError trying to allocate 12 MiB in the distributed optimizer — a different, non-memory-safety failure mode that we are tracking separately.)

Likely cause, and relation to open PR #3417

Open PR #3417 ("[Common] Fix TMA synchronization in quantization kernels") describes exactly this defect class: bulk async-groups are maintained per issuing thread, so calling cp_async_bulk_wait_group_read() from only the issuing thread without a subsequent CTA-wide __syncthreads() handoff lets cooperative writers begin overwriting a shared-memory output ring buffer while the TMA engine is still reading it — described there as producing "rare, timing-dependent output corruption, particularly under global-memory backpressure". That matches both the partial-rank, non-deterministic signature above and the fact that this workload runs with its device memory almost fully reserved.

However, #3417's 13 changed files cover the FP8, MXFP8 quantize/dequantize, specialized MXFP8, NVFP4 and fused-group-requantize paths plus common/util/ptx.cuh, and do not include transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh. #3417 branches from main at abf0ebb079b666ea4a37af9f5b9a36adb854d0db, which predates #3337, so the rewritten grouped MXFP8 quantize kernel was never in its audit scope. That file still performs an issuer-scoped ptx::cp_async_bulk_wait_group_read<PREFETCH_STAGES>() inside an if (leading_thread) { ... } block.

This is also consistent with #3337's own commit history, which contains the messages "Changed kernel launch configs from 3D grid to 1D grid" and "Fixed TMA synchronization. Restored the number of tensors check": the grouped MXFP8 quantize kernel's launch geometry and TMA store/reuse pipeline both changed substantially in the same commit that introduces the fault, and its diff is concentrated in group_quantize_mxfp8.cuh (+463/-182) and cast/core/grouped_tma.cuh (+6).

Request: please consider extending #3417's issuer-to-CTA completion handoff audit (and the accompanying compile-time ring/wait-depth invariants) to group_quantize_mxfp8.cuh as rewritten by #3337, and/or add a grouped-MXFP8 quantize test that runs under memory pressure with many groups so the reuse race is exercised.

Environment overview

• Environment location: Docker, on-premise GB200 NVL cluster (Slurm, 2 nodes x 4 GPUs).
• Method of Transformer Engine install: from source, at an exact commit. TransformerEngine is checked out at the commit above and built inside the container image; no wheel was used, because no per-commit or commit-stamped wheels are published for release_v2.19 revisions.
• Docker: a PyTorch base container image with TransformerEngine rebuilt from source on top. Exact docker pull/docker run commands are not applicable — the image is built and launched by our CI, so they would not be reproducible outside it.

Environment details

• PyTorch: 2.14.0a0 (single source commit b2c75dd062, identical on both sides of the bisect)
• Transformer Engine: 2.19.0+c3066f8 (bad) and 2.19.0+3231a15 (good)
• CUDA: 13.4.1.012
• cuDNN: 9.26.0.37
• NCCL: 2.31.2
• Driver: 580.178.04
• Architecture: linux/arm64 (aarch64)
• OS and Python versions: not captured in the job environment record we have; both sides ran the same container base, and the environment diff reported no differing system information.

Device details

• GPU model: NVIDIA GB200 NVL (sm_100), 8 GPUs across 2 nodes, ~184 GiB usable memory per GPU.

Additional context

• The workload runs with device memory nearly fully reserved (peak reserved ~191.5 GiB against a ~184.3 GiB per-GPU budget, with allocator retries and repeated expandable-segment mapping failures that are survived). We mention this only because [Common] Fix TMA synchronization in quantization kernels #3417 identifies global-memory backpressure as the condition that makes this race observable.
• Related but distinct existing issues we checked and ruled out: [Bug] Wrong gradients with fused grouped MLP under Format.HYBRID (MXFP8) #3342 (wrong gradients with fused grouped MLP under Format.HYBRID — a non-fatal numerical defect predating this bracket), [Bug] CUDA graph backward replay illegal memory access after validation with mixed NVFP4/MXFP8 Mamba #3316, Multi-GPU FP8 (accelerate.dispatch_model + DelayedScaling HYBRID) — illegal memory access on H200 sm_90 #3124, [Bug] FusedAttention THD + learnable softmax backward IMA (cuDNN err 700) on Hopper #3249.

---

This issue was drafted with assistance from the opus AI model.

[ptrendx]

@Oleg-Goncharov Please take a look.

[Oleg-Goncharov]

Hi @janbernloehr, I found another difference between the grouped TMA kernels and the other quantization kernels that may explain the illegal memory access.
The grouped path was still aligning its dynamic shared-memory pointer by converting it to an integer, rounding the address, and converting it back to a pointer. This pattern was changed by #3291 because it can make the compiler lose its knowledge of the shared-memory address-space and emit incorrect LD.E/ST.E instructions instead of LDS/STS.
I prepared PR#3482 [Common] Preserve shared-memory pointer provenance in TMA kernels that fixes it.
Since the reported failure is memory pressure sensitive and depends on the original GB200 multi-node environment, could you please test this branch using the same reproducer?

[kainzhong]

@Oleg-Goncharov

hum I wonder why that could be the cause. I thought align_smem_ptr_per_TMA_requirements / align_up only differ in term of how the address is computed but to the caller they should have the same semantic. Even if you don't use align_up the compiler still emits LD.E / ST.E which should still be legal for accessing SMEM (at least that's what happens in #3291) so I'm unsure if that's the root cause and #3482 can fix it.

Nevertheless I think that PR by itself is correct and should be a net improvement