Skip to content

fix: work around the Triton buffer-store miscompile per kernel - #962

Open
WangLingxun wants to merge 8 commits into
mainfrom
fix/gfx942-triton-bufops-war
Open

fix: work around the Triton buffer-store miscompile per kernel#962
WangLingxun wants to merge 8 commits into
mainfrom
fix/gfx942-triton-bufops-war

Conversation

@WangLingxun

@WangLingxun WangLingxun commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Why

The AMD Triton backend can emit a buffer_store_dwordx4 whose data VGPRs are redefined by a later instruction with no intervening s_waitcnt vmcnt, so the store writes whatever the clobbering instruction left in those registers. Nothing faults and nothing warns; a few percent of the kernel's output elements are simply garbage. Inductor's SwiGLU backward fusion trips this and poisons the MLP weight gradients, which diverges TorchTitan training. Megatron runs the same inductor-generated kernels, but there the corruption tends to stay inside the normal numeric range, so a finite loss curve does not prove a clean run.

What changed

New patch: primus/core/patches/triton_bufops_war_patches.py. Wraps triton.compile, scans each kernel's emitted AMDGCN for the hazard, and recompiles only the affected kernels with buffer ops disabled. Keying on the machine code means there is no kernel or architecture allowlist to maintain. Registered for both torch backends at the build_args phase, gated only on being a ROCm GPU; non-HIP kernels are skipped per compile. It carries no off switch of its own, since the recompile is a correctness fix and a kernel that does not need one is left alone anyway.

AMDGCN_USE_BUFFER_OPS=0 is deliberately not set globally: the hazardous kernels lose nothing without buffer addressing, but the Primus-Turbo grouped-GEMM kernels spill without it and carry no hazard, so the global switch costs ~18% on MoE recipes for kernels that never needed fixing. Per-kernel selection measured 299.9 vs 299.0 TFLOP/s per GPU on deepseek_v2_lite-BF16, against 244.6 with the global switch.

The patch also appends to TORCH_COMPILE_CACHE_KEY_TAG, since on an FX graph cache hit inductor never calls triton.compile — without invalidation a cache filled before this lands would keep serving hazardous binaries. This is the only environment variable the PR touches, and it is PyTorch's own.

Divergence detection. tests/utils.py now requires every logged loss and grad norm to be finite, because a numerically diverged run still exits 0 and still prints the completion marker. tools/auto_benchmark/metrics.py marks a run with non-finite metrics invalid instead of publishing its throughput.

Tests. 18 tests for the hazard criterion and the ROCm gate, 14 for the divergence parser on both backends' log formats. All static text analysis: no GPU, no Triton.

Docs. A workarounds section in performance-tuning.md and one row in environment-variables.md.

Test plan

  • 32 unit tests, ~0.2 s
  • MI300X: hazardous kernel fixed and clean kernel untouched, residual 0
  • MI355X: E2E green on both backends; the pattern also appears there, though an isolated numeric A/B of the matching kernel came back bit-identical
  • CI E2E on both backends

@WangLingxun
WangLingxun force-pushed the fix/gfx942-triton-bufops-war branch 3 times, most recently from f791a77 to 18071e1 Compare August 12, 2026 11:33
@WangLingxun WangLingxun changed the title fix: work around the gfx942 Triton buffer-store miscompile per kernel fix: work around the Triton buffer-store miscompile per kernel Aug 12, 2026
Comment thread tests/unit_tests/core/patches/test_triton_bufops_war.py Fixed
@WangLingxun
WangLingxun force-pushed the fix/gfx942-triton-bufops-war branch from 18071e1 to cedd6b4 Compare August 13, 2026 07:18
@WangLingxun
WangLingxun marked this pull request as ready for review August 13, 2026 07:21
@WangLingxun
WangLingxun force-pushed the fix/gfx942-triton-bufops-war branch 2 times, most recently from a0d5682 to 9269d33 Compare August 13, 2026 07:43
The AMD Triton backend can redefine a buffer_store_dwordx4's data VGPRs before
the store has read them, with no s_waitcnt vmcnt in between. The store then
writes whatever the clobbering instruction left there. Nothing faults and
nothing warns; a few percent of the kernel's output elements are simply garbage.
Inductor's SwiGLU backward fusion trips this and poisons the w1/w2/w3 weight
gradients, so TorchTitan training NaNs out on the second step. Megatron runs the
same inductor-generated kernels -- megatron/core/jit.py sets
jit_fuser = torch.compile -- but there the corruption tends to stay inside the
normal numeric range, so a finite loss curve does not prove a clean run.

AMDGCN_USE_BUFFER_OPS=0 fixes it globally and costs ~18% on MoE recipes, which
is why this does not use it. The broken kernels and the expensive ones turn out
to be disjoint: the hazard lives in the inductor pointwise kernels and in
_grouped_gemm_output_tail_kernel, none of which lose anything without buffer
addressing, while the cost lives in the Primus-Turbo grouped-GEMM kernels,
which carry no hazard but spill once they lose it (_grouped_variable_k_gemm_kernel
scratch grows 128 -> 524 bytes). Turning the feature off everywhere would pay
for kernels that never needed fixing.

So compile every kernel normally, scan the emitted AMDGCN, and recompile only
the affected kernels with buffer ops off. On deepseek_v2_lite-BF16 that gives
299.9 TFLOP/s/GPU against 299.0 with buffer ops left on and 244.6 with them off
globally; a microbenchmark whose output held 201,021 wrong values out of 4.2M
elements drops to zero, and dense recipes are unaffected either way.

The criterion is the machine code actually produced, so neither a kernel nor an
architecture allowlist has to be maintained. The wrong values above were measured
on gfx942. gfx950 emits the same pattern -- 9 of 141 kernels scanned in a real
TorchTitan run, triton_poi_fused__to_copy_view_3/4 -- but an isolated A/B of the
matching kernel there produced bit-identical output with and without buffer ops
across 1.3e9 elements, so that chip appears to tolerate it. The scan stays on for
both, because the omitted s_waitcnt is a compiler defect regardless of which
chips currently tolerate it, and because an allowlist leaves every new
architecture unprotected until somebody remembers to widen it. Per-kernel
selection measured -0.17% on a dense llama3.1_8B recipe over 8xMI355X, within
run-to-run noise.

The patch also appends to TORCH_COMPILE_CACHE_KEY_TAG. On an FX graph cache hit
inductor never calls triton.compile, so a cache filled before this landed would
keep serving hazardous binaries -- invisibly, because the residual counter only
sees kernels that reach the compiler.

This is not fixed by moving to a newer Triton: upstream 3.7.0, which the CI
image builds from source, and the ROCm 3.7.0, 3.7.1 and 3.8.0 builds all still
emit the pattern on at least one production kernel.

Also stop letting a diverged run look like a good one. E2E tests now require
every logged loss and grad norm to be finite, since a numerically diverged run
still exits 0 and still prints the completion marker, and auto_benchmark marks
a run with non-finite metrics invalid rather than publishing its throughput.
@WangLingxun
WangLingxun force-pushed the fix/gfx942-triton-bufops-war branch from 9269d33 to b74383c Compare August 13, 2026 08:27
Enable Inductor precision-cast emulation before compilation so FSDP backward graphs retain eager numerical semantics instead of silently producing NaN gradients.
Compile only dense DeepSeek blocks to avoid Inductor corruption and v0.2.2 fragmentation, and honor the Turbo attention switch so MI300X FP8 recipes can avoid the invalid aiter hd128 backward path.
Distinguish the ROCm BF16/FSDP precision failure from the resolved upstream float8 issue, and automatically bypass the current gfx942 DeepSeek FP8 Turbo-attention failure without changing recipe configuration.
Keep Turbo attention enabled while bypassing the invalid gfx942 fmha_v3 backward dispatch for DeepSeek FP8, matching the working CK path from the original v26.5 image without sacrificing steady-state throughput.
Apply the repository's Black and isort rules so the new TorchTitan fixes pass pre-commit in CI.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant