[https://nvbugs/5970614][fix] Sync CTA before PDL trigger in quantize_with_block_size - #14668
Conversation
📝 WalkthroughWalkthroughThe PR addresses a CUDA graph execution issue in NVIDIA's TensorRT-LLM by adding memory-ordering synchronization to the quantization kernel and adjusting the DeepSeek-R1 NVFP4 test configuration, then removing the corresponding test waiver. ChangesNVFP4 DeepSeek-R1 Throughput Stabilization
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/kernels/quantization.cuh (1)
2-2:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUpdate the copyright header year on this modified source file.
This file was modified, but the NVIDIA header still ends at 2023. Please update it to include the latest modification year.
Suggested fix
- * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved.As per coding guidelines:
**/*.{cpp,cc,h,hpp,py,cu,cuh}: Include NVIDIA copyright header on ALL new files; update year on modified files.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/quantization.cuh` at line 2, Update the NVIDIA copyright header at the top of the modified source file by changing the year range from "2019-2023" to include the current modification year (e.g. "2019-2026"); edit the top-of-file header comment in quantization.cuh (the initial copyright comment block) so the year range reflects the latest modification year.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/tensorrt_llm/kernels/quantization.cuh`:
- Line 2: Update the NVIDIA copyright header at the top of the modified source
file by changing the year range from "2019-2023" to include the current
modification year (e.g. "2019-2026"); edit the top-of-file header comment in
quantization.cuh (the initial copyright comment block) so the year range
reflects the latest modification year.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 04887fc1-1f9e-4ec2-8b75-a6af13d6921c
📒 Files selected for processing (3)
cpp/tensorrt_llm/kernels/quantization.cuhtests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
/bot run --disable-fail-fast |
|
PR_Github #53276 [ run ] triggered by Bot. Commit: |
|
PR_Github #53276 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53316 [ run ] triggered by Bot. Commit: |
|
PR_Github #53316 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53440 [ run ] triggered by Bot. Commit: |
|
PR_Github #53440 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53533 [ run ] triggered by Bot. Commit: |
|
PR_Github #53533 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53601 [ run ] triggered by Bot. Commit: |
|
PR_Github #53601 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53784 [ run ] triggered by Bot. Commit: |
|
PR_Github #53784 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #57996 [ run ] triggered by Bot. Commit: |
|
PR_Github #57996 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58175 [ run ] triggered by Bot. Commit: |
|
PR_Github #58175 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58221 [ run ] triggered by Bot. Commit: |
|
PR_Github #58221 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #58384 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #58500 [ run ] triggered by Bot. Commit: |
|
PR_Github #58384 [ run ] completed with state
|
|
PR_Github #58500 [ run ] completed with state
|
…_with_block_size Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
…e_with_block_size Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
… SF group The CTA-wide __syncthreads() added before the PDL trigger in quantize_with_block_size deadlocks against the full-warp __shfl_xor_sync(0xffffffff) inside the cvt helpers when the thread-strided column loop gives warp lanes different trip counts (e.g. MXFP8 k=1568 alignment=64 -> blockDim=200, GPT-OSS hidden=2880 -> blockDim=360): lanes still in the loop wait on the shuffle for lanes already parked at the barrier. Narrow the shuffle membermask to the 2- or 4-lane scale-factor group, which the xor-1/xor-2 butterfly never crosses and which is always trip-uniform because the block size and the column-thread count are multiples of the group size. Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
…4_mtp test The PDL race behind this waive is fixed by the producer-side CTA barrier in quantize_with_block_size; remove the waive so CI validates the fix end to end. Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
…ad of the PDL trigger Address review: with the previous __syncthreads(); __threadfence() order, the first thread past the barrier could reach the trigger before its peers had executed their own fences, so peer stores were not yet guaranteed device-visible. Fencing before the barrier makes every thread flush its own stores first; barrier release then implies all stores are device-visible before any thread triggers. Also bump the copyright year. Signed-off-by: Tianyu Xiong <117647511+tianyuxbear@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #58551 [ run ] triggered by Bot. Commit: |
|
PR_Github #58551 [ run ] completed with state |
4 participants
Summary
Fixes a PDL (Programmatic Dependent Launch) race in
quantize_with_block_sizethat intermittently corrupts NVFP4 GEMM outputs and degrades GSM8K accuracy for DeepSeek-R1 NVFP4 on GB300 + PP=4 + MTP (nvbug 5970614).cudaTriggerProgrammaticLaunchCompletion()only signals that the CTA has reached the trigger point — it does not flush prior stores to global memory. Memory visibility for the secondary kernel must be provided either by the producer (a fence before the trigger) or by the consumer (wait_on_dependent_grids()before its first dependent load). In the current NVFP4 path neither side does so: the producer lacks a fence, and the sm103 blockscaled GEMM'smain_sf_loadwarp branch is missing the correspondingwait_on_dependent_grids()(tracked separately in NVIDIA/cutlass#3279).Compounding this, PDL completion is reported per-CTA at-least-once: a single warp reaching the trigger marks its whole CTA as "trigger reached", even if peer warps in the same CTA are still writing
sf_out/out. Once every CTA has been marked, the driver launches the secondary kernel, which TMA-loads partial data — NaNs then propagate through the DeepSeek-R1 forward and corrupt output tokens.Fix: insert
__threadfence(); __syncthreads();immediately before the trigger. Each thread first makes its own stores device-visible; the barrier then guarantees every thread has done so before any thread can signal PDL completion.In-tree precedent
This mirrors an existing pattern in
cpp/tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.cuh:865-867, which uses a barrier plusmembar.gl(__threadfence()is the CUDA intrinsic formembar.gl) before the samecudaTriggerProgrammaticLaunchCompletion()call. Per review, this PR places the fence before the barrier — the strictly stronger order: barrier release then already implies every thread's stores are device-visible, without relying on post-barrier fence timing.Evidence
The race was characterized on
release/1.2, where the timing window reproduces it deterministically enough to measure. A cache-bypassld.global.cv.u8probe on the SF buffer detects0x7fpoison fill (from thecudaMallocAsyncpool) as a direct race indicator:On
mainthe race is latent: the GSM8K test no longer fails even without this fix (10/10 PASS, mean 94.95 ± 0.25 with autotuner off; with the fix mean 95.14 ± 0.19). Code/scheduler changes sincerelease/1.2appear to have narrowed the producer-to-consumer window enough that the race no longer trips the GSM8K threshold — but the underlying defect is unchanged, and any future change that widens the window (cutlass bump, scheduler reordering, kernel fusion) can re-expose it.The full investigation branch — probes, race-rate measurement, per-run logs — lives on my fork at
tianyuxbear/TensorRT-LLM:fix/5970614-bak, which is based onrelease/1.2(where the race is deterministically reproducible).The probe code used to measure the race is not part of this PR; it was a one-off diagnostic inside the cutlass consumer.
Test-case waive removed by this PR
On
mainthe affected test case (TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_pp4_mtp]) is waived intests/integration/test_lists/waives.txtunder this very nvbug (5970614). Since this PR fixes the underlying race, it also removes the waive so CI validates the fix end to end on the real reproducer.Historical note: the deterministic accuracy gap in the Evidence section was measured on
release/1.2, where the producer-to-consumer timing window reproduces the race reliably; onmainthe race is latent for this shape, so the unwaived run primarily guards against future changes re-widening the window.Companion cutlass fix
The consumer-side fix at NVIDIA/cutlass#3279 addresses the missing
wait_on_dependent_grids()in the sm103 blockscaled GEMMmain_sf_loadwarp branch (the other half of the race described in Summary). Either fix alone closes the race; both are correct individually:quantize_with_block_size, regardless of which cutlass revision trtllm pulls in.Follow-up fix: shuffle-mask deadlock exposed by the barrier
The first full CI run of this PR hung 19 test cases on Blackwell stages (B200/B300) — all hangs, no accuracy failures:
test_fp8_quantize.py::test_mxfp8_quantize_alignment_torch_device[*-1568-*]a_fp8.cpu()after 2400 s; B200 xdist workers mass-crashed)TestGPTOSS::test_w4_1gpu[*]/test_dummy_load_formatLLM(...)init during the first warmup forwardtest_indexer_topk.py::test_indexer_topk_prefill[*]Root cause. The cvt helpers called from
quantize_with_block_size's thread-strided column loop reduce the per-SF absmax with a full-warp__shfl_xor_sync(0xffffffff, ...)— but the shuffle is only reached on the real-data branch of the loop body; lanes that draw padding columns just store zeros and skip it. When the data/padding boundary cuts through a warp, only part of that warp reaches the shuffle: MXFP8 k=1568 / alignment=64 gives blockDim=200 with 196 data-column threads (196 = 6x32+4, splitting warp 6 at lane 4); the GPT-OSS MoE activation quant (hidden=2880, weight alignment 256) gives blockDim=384 with 360 data-column threads (360 = 11x32+8, splitting warp 11 at lane 8). Before this PR that was benign: the padding lanes exited the kernel after their stores, and exited lanes satisfy sync collectives on sm_70+, so the data lanes' full-warp shuffle completed. With the new__syncthreads(), the padding lanes are instead parked alive at the barrier: the shuffle waits for lanes held at the barrier, the barrier waits for lanes held at the shuffle — a deterministic circular wait. The barrier did not create the bug; it removed the exit-semantics cover from a latent over-wide membermask.Fix. Narrow the shuffle membermask to the 2-/4-lane scale-factor group (new
cvt_sf_group_shfl_mask()helper; 5 call sites acrosscvt_warp_fp16_to_fp4/cvt_warp_fp8_to_fp4/cvt_warp_fp16_to_mxfp8). The xor-1/xor-2 butterfly never exchanges data outside the aligned group, so the mask now names exactly the lanes that must converge. Group lanes always share a trip count because blockDim and the column-thread count are both multiples of the group size (guaranteed by thek % sfVecSize == 0checks at the op layer). This is a strictly weaker convergence precondition with bit-identical numerics — the membermask never selects shuffle sources, so callers of these helpers in other kernels are unaffected.Local verification (single B200). All 19 failed CI cases rerun after the fix: 19/19 pass, 0 timeouts. The deterministic hang case
[64-True-dtype0-1568-1]completes in 0.08 s (previously 2400 s timeout). GPT-OSS 120B GSM8K scores 89.23 vs the 85.0 reference across all three w4_1gpu variants — the mask change preserves quantization numerics end to end.Why the DeepSeek-R1 e2e never hit this. The deadlock requires the boundary of the shuffle-participation set (which lanes reach the cvt call) to cut through the middle of a warp. DeepSeek-R1 (hidden=7168) has no padding columns at all (7168 is already SF-layout-aligned), and its only boundary is the loop-stride one: blockDim=512 with 896 data-column threads, so the second stride covers threads 0..383 — 384 is an exact multiple of 32, whole warps diverge together, and every warp's full-mask shuffle stayed convergent. Neither the release/1.2 race experiments nor the DeepSeek CI case could ever expose the hazard; it takes shapes like k=1568 (alignment unit tests) or hidden=2880 (GPT-OSS MoE) to land a boundary mid-warp.
Risk
quantize_with_block_size, NVFP4 / FP8 / MXFP8 paths) plus the shuffle membermasks in the shared cvt helpers (strictly weaker convergence requirement, identical data movement).__syncthreads()+ one__threadfence()at the very end of the kernel (cold path). No effect on PDL launch overlap — the trigger still happens; it just happens after the CTA has drained.