[None][fix] kdaDecode legacy: __syncwarp before the lane-0 block-reduce stores by vsabavat · Pull Request #19830 · NVIDIA/TensorRT-LLM

vsabavat · 2026-10-02T23:57:58Z

Description

In block_reduce_sum2_for and block_reduce_sum2_active_for
(cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu, used by kda_decode_fusion_compact_heads_kernel):

the lanes of warp 0 read the per-warp partials scratch[lane] / scratch[kReduceWarps + lane] and reduce them with
shuffles;
lane 0 then stores the block totals into scratch[0] and scratch[1], which lane 1 read.

Shuffles do not order shared memory, so lane 0's stores and lane 1's reads were unordered. compute-sanitizer racecheck
reports 32 warnings per run. Lane 0's totals depend on lane 1's loaded value, so results were not affected.
__syncwarp() before the stores makes the order explicit.

Test Coverage

No new test: the outputs do not change, and the order is a racecheck finding.

Racecheck: on extension builds of the original and the fixed .cu it goes from 3 records to 0.
Bits: 0 of 99 runs differ.
Existing tests: tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py covers the kernel's outputs.

PR Checklist

Please review the following before submitting your PR:

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

__syncwarp() was added after the shuffle reductions in both block_reduce_sum2_for and block_reduce_sum2_active_for. It makes warp synchronization explicit before lane 0 stores the block totals in shared memory. No public API changes are reported.

QA Engineer Review

No test changes.

Per-File QA Perspective

cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu: Verify both reduction paths produce correct outputs and that racecheck no longer reports the described shared-memory race. The PR reports no output-bit differences in 99 runs, but this result was not independently verified here.

…ce stores In block_reduce_sum2_for and block_reduce_sum2_active_for, lanes of warp 0 read the per-warp partials scratch[lane] / scratch[kReduceWarps + lane], reduce them with shuffles, and lane 0 then stores the block totals into scratch[0] and scratch[1], which lane 1 read. Shuffles do not order shared memory, so the store and lane 1's read were unordered (compute-sanitizer racecheck: 32 warnings per run of test_k3_kda_decode_attn, all in kda_decode_fusion_compact_heads_kernel). Lane 0's totals depend on lane 1's loaded value, so the results were not affected; __syncwarp() makes the order explicit. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>

coderabbitai · 2026-10-03T05:56:00Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 87a3bc85-3aef-44c4-a550-33255316330a

📥 Commits

Reviewing files that changed from the base of the PR and between f388b7c and af3e384.

📒 Files selected for processing (1)

cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 4 remain after this review.

Walkthrough

Both block_reduce_sum2_for and block_reduce_sum2_active_for now synchronize the first warp before lane 0 overwrites shared-memory partial sums.

Changes

Warp reduction synchronization

Layer / File(s)	Summary
Synchronize reduction helpers `cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu`	Both helpers add `__syncwarp()` after shuffle reductions and before lane 0 overwrites shared-memory partial sums.

Priority: ➖ Normal

Estimated code review effort: 1 (Trivial) | ~3 minutes

Change: Bug fix

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to af3e3

The added barriers order the shared-memory accesses, and the inspected call paths show no remaining actionable merge risk.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title follows the repository format and clearly identifies the fix: adding __syncwarp() before lane-0 block-reduce stores.
Description check	✅ Passed	The description explains the issue and solution, reports test coverage and racecheck results, and includes the required PR checklist sections.

Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create a new PR

Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

github-actions · 2026-10-06T06:10:31Z

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.

longlee0622 · 2026-10-06T06:11:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-06T06:18:33Z

PR_Github #76326 [ run ] triggered by Bot. Commit: af3e384 Link to invocation

tensorrt-cicd · 2026-10-06T10:18:38Z

PR_Github #76326 [ run ] completed with state SUCCESS. Commit: af3e384
/LLM/main/L0_MergeRequest_PR pipeline #62931 completed with status: 'SUCCESS'

CI Report

Link to invocation

Conflicts: - tests/integration/test_lists/test-db/l0_b200.yml: both sides appended to the same B200 pre-merge block (main a MiniMax-H3 visual_gen case, this branch the Kimi K3 KDA / MLA decode kernel tests). Both are kept, main's line first. Also: - cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu: main already has this branch's two __syncwarp() calls (NVIDIA#19830), so the file now equals main's. Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>

This branch carries NVIDIA#19816's content at 546e567. The conflicts and semantic fixes below are the ones NVIDIA#19816 resolved in its published merges 0c9ab14 (main 70193ec) and 44a06db (main a1f562a), applied to this branch's copies of those files. Conflicts: - model_engine.py: main's NVIDIA#19696 moved the engine decoder buffers and state into _allocate_decoder_buffers() and _init_decoder_state(). Kept main's layout with this branch's __init__ changes ported into it (_staged_previous_batch_indices / _staged_previous_pos_indices next to _encoder_decoder_staged_request_ids, _step_input_stage next to _prepare_inputs_event, no num_accepted_draft_tokens_cuda). The gather_ids upload stays copy_to_device_if_changed, under main's enable_spec_decode argument. - cuda_graph_runner.py: main's NVIDIA#19306 strict-buffer snapshot now follows the captured forward directly; the post-capture postprocess_fn call stays removed. - test_cuda_graph_capture_replay.py: both sides' test classes kept. - l0_b200.yml: both sides appended to the same B200 pre-merge block; main's MiniMax-H3 line first, then the Kimi K3 kernel tests. Semantic fixes: - model_engine.py: the step-copy selection reads the enable_spec_decode argument that NVIDIA#19696 passes to _prepare_tp_inputs. - kv_cache_manager_v2.py: stage_batch_block_offsets builds its copy index as main's copy_batch_block_offsets now does (NVIDIA#17306: beam width up to max_copy_beam_width, beam 0 replicated for cross KV). - test_pytorch_model_engine.py: the staged-step test passes _prepare_tp_inputs' keyword-only enable_spec_decode, runtime_draft_len and is_dummy arguments. - test_cuda_graph_capture_replay.py: the postprocess_fn case of main's test_replay_rejects_rebound_graph_attr needs the post-capture call and is dropped; test_postprocess_fn_follows_only_the_warmup_forwards pins that postprocess_fn runs only after the warmup forward. Also: cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu now equals main's (NVIDIA#19830 is this branch's two __syncwarp() calls). Signed-off-by: Vasanth Sabavat <vsabavat@nvidia.com>

github-actions Bot assigned vsabavat Oct 2, 2026

svc-trtllm-gh-bot added the Community want to contribute PRs initiated from Community label Oct 3, 2026

This was referenced Oct 3, 2026

[None][feat] modeling_v2 catalog: Kimi K3 KDA and MLA decode kernels as stateful entries (sm_100) #19832

Closed

[None][feat] modeling_v2: Kimi K3 MXFP4 target, sm_100, tp16 attention + moe tp4 x ep4 #19839

Closed

vsabavat marked this pull request as ready for review October 3, 2026 05:49

vsabavat requested a review from a team as a code owner October 3, 2026 05:49

vsabavat requested a review from zongfeijing October 3, 2026 05:49

vsabavat mentioned this pull request Oct 3, 2026

[None][feat] Kimi K3 on modeling_v2: stateful catalog entries, MNNVL and collective decode kernels, drafter attention, MoE route B #19841

Closed

10 tasks done

longlee0622 approved these changes Oct 6, 2026

View reviewed changes

kris1025 approved these changes Oct 6, 2026

View reviewed changes

trtllm-agent added the ci: full pre-merge approved label Oct 6, 2026

longlee0622 enabled auto-merge (squash) October 6, 2026 07:42

longlee0622 merged commit 089d586 into NVIDIA:main Oct 6, 2026
25 of 27 checks passed

karljang removed the Community want to contribute PRs initiated from Community label Oct 6, 2026

Navigation Menu

Repository navigation

[None][fix] kdaDecode legacy: __syncwarp before the lane-0 block-reduce stores - #19830

[None][fix] kdaDecode legacy: __syncwarp before the lane-0 block-reduce stores#19830
longlee0622 merged 1 commit into
NVIDIA:mainfrom
vsabavat:k3-up/p6-kda-decode-legacy-syncwarp

vsabavat commented Oct 2, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Oct 3, 2026 •

edited

Loading

❌ Failed checks (1 warning)

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

longlee0622 commented Oct 6, 2026

Uh oh!

tensorrt-cicd commented Oct 6, 2026

Uh oh!

tensorrt-cicd commented Oct 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

vsabavat commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Dev Engineer Review

QA Engineer Review

Per-File QA Perspective

Uh oh!

coderabbitai Bot commented Oct 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

longlee0622 commented Oct 6, 2026

Uh oh!

tensorrt-cicd commented Oct 6, 2026

Uh oh!

tensorrt-cicd commented Oct 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

vsabavat commented Oct 2, 2026 •

edited

Loading

coderabbitai Bot commented Oct 3, 2026 •

edited

Loading