fix: work around the Triton buffer-store miscompile per kernel - #962
Open
WangLingxun wants to merge 8 commits into
Open
fix: work around the Triton buffer-store miscompile per kernel#962WangLingxun wants to merge 8 commits into
WangLingxun wants to merge 8 commits into
Conversation
WangLingxun
force-pushed
the
fix/gfx942-triton-bufops-war
branch
3 times, most recently
from
August 12, 2026 11:33
f791a77 to
18071e1
Compare
WangLingxun
force-pushed
the
fix/gfx942-triton-bufops-war
branch
from
August 13, 2026 07:18
18071e1 to
cedd6b4
Compare
WangLingxun
marked this pull request as ready for review
August 13, 2026 07:21
WangLingxun
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
August 13, 2026 07:21
WangLingxun
force-pushed
the
fix/gfx942-triton-bufops-war
branch
2 times, most recently
from
August 13, 2026 07:43
a0d5682 to
9269d33
Compare
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
force-pushed
the
fix/gfx942-triton-bufops-war
branch
from
August 13, 2026 08:27
9269d33 to
b74383c
Compare
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
The AMD Triton backend can emit a
buffer_store_dwordx4whose data VGPRs are redefined by a later instruction with no intervenings_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. Wrapstriton.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 thebuild_argsphase, 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=0is 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 ondeepseek_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 callstriton.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.pynow 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.pymarks 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.mdand one row inenvironment-variables.md.Test plan