[None][fix] Kimi K3 KDA prefill: K1 warps finish reading sPartialLast before the next chunk

[vsabavat]

Description

What goes wrong. fused_kernel123
(tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py) has a race between chunks:

• its K1 warps read their scanned prefix row and the chunk total from sPartialLast after the scan barrier;
• the next chunk's partial sums overwrite sPartialLast before any K1 barrier.

A K1 warp that reaches its reads late uses the next chunk's sums.

• Racecheck reports 8 error records per run.
• With one K1 warp stalled 20 us after the scan barrier, every case of
  test_fused_prefill_beta_sigmoid_matches_unfused_kernel changes, by up to 7e14 in the output.
• Under racecheck, test_fused_prefill_beta_sigmoid_matches_unfused_kernel[padded_eqlen] gives a NaN output.
• Natively the existing tests pass.

Fix. One K1 barrier after the reads orders them before the next chunk's writes. Outputs are bit-identical to
before.

Test Coverage

tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py::test_fused_prefill_does_not_depend_on_k1_warp_timing
(SM 100):

• it runs a copy of the kernel with one K1 warp stalled after the scan barrier;
• it requires the unstalled kernel's bits;
• it fails without the barrier.

The file is already in l0_b200.yml.

Results on GB200:

• The new test fails 2 / 2 on main (padded_eqlen, varlen).
• With this PR test_kda_prefill_op.py passes 26 / 26.
• modules/kimi_kda/ has no new failures: 119 passed and 1 skipped with this PR, 116 passed and 2 skipped on
  main. The new test adds 2 passes, and one test that skipped in main's run ran and passed in this PR's run.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title. (No API change.)

• Any new dependencies have been scanned for license and vulnerabilities (No new dependencies.)

• CODEOWNERS updated if ownership changes (No ownership change.)

• Documentation updated as needed (No user-facing change.)

• Update tava architecture diagram if there is a significant design change in PR. (No design change.)

• The reviewers assigned automatically/manually are appropriate for the PR. (To check once the PR is open.)

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

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

Dev Engineer Review

• fused_kernel123 adds a K1 barrier after each warp reads its prefix and chunk total from sPartialLast. This prevents the next chunk from overwriting those values before all K1 warps finish reading them.
• The synchronization change adds a barrier on the kernel path. No performance results were supplied.
• Review severity counts are unavailable because no current review findings were supplied.

QA Engineer Review

• tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py adds the parameterized test_fused_prefill_does_not_depend_on_k1_warp_timing test for padded equal-length and variable-length inputs. It injects a delay into a K1 warp and checks bit-exact output and state parity against the unmodified kernel.
• The test file is listed in tests/integration/test_lists/test-db/l0_b200.yml and tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml. No changed test-list files were reported.
• The PR objectives report 26/26 tests passing in test_kda_prefill_op.py on GB200. These results were not independently verified here.
• Coverage verdict: sufficient for the reported race and the two tested input modes.

Per-File QA Perspective

• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py: Verify K1 synchronization prevents reads of overwritten sPartialLast values and preserves prefill results. The barrier also adds synchronization overhead.
• tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py: The timing-stall test checks bit-exact output and state parity for padded equal-length and variable-length inputs. The test file appears in the GB200 CI test lists noted above.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

🧰 Additional context used

📚 Code guidelines (1)

tests/AGENTS.md — auto-discovered

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: ae49fa16-97ec-4ad0-b855-8ee4f7a7f08a

📥 Commits

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

📒 Files selected for processing (2)

• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py
• tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py

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

---

Walkthrough

The K1 kernel now synchronizes after reading prefix and chunk-total values, before Pass 2a. A parameterized test delays a K1 warp and checks output and state parity for padded and variable-length prefill inputs.

Changes

K1 synchronization

Layer / File(s)	Summary
K1 read barrier and timing validation tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py, tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py	The kernel adds a barrier after K1 reads prefix and chunk-total values. The test delays a K1 warp and compares output and state with the unmodified kernel for padded and variable-length inputs.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: brnguyen2

Merge Risk: ⚪ Minimal · up to c5a4e

The barrier and timing test have no identified issue that needs resolution before merge.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the K1 barrier fix and the race it prevents.
Description check	✅ Passed	The description explains the race, the fix, and the regression test results. It includes the required sections and checklist; the reviewer assignment remains unchecked because it can be confirmed afte…
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.

• 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.}