[vsabavat]

Description

CUDA-graph padding rows and warmup dummies decode on DFlash's dummy slot (worker._dummy_slot,
tensorrt_llm/_torch/speculative/dflash.py).

• Nothing resets it. DFlashSpecMetadata.prepare resets only evicted slots, and _assign_slot never runs for the
  dummy slot, so the dummy's context length grows by every padded step's accepted tokens.
• Its page-table rows point at a live page. The rows are the padding request's pages, with every other entry
  mapped to page 0, which belongs to another, live request.
• Result: once the length passes the padding request's own pages, the padding rows' context K / V lands on page 0.
  A live request's drafter context is overwritten on every padded step. This costs acceptance only, because the
  target verifies.

It is reachable in any configuration whose CUDA graphs pad a batch.

The drafting step empties the dummy slot right after it adds the step's accepted tokens (a zero_() on the fixed
slot, captured with the step), so prepare() keeps writing only evicted slots.

Test Coverage

tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py (new; needs a GPU, the slot lengths live on
it):

• the drafting-step update (_advance_ctx_len) keeps the dummy slot at 0 and leaves the real slots untouched;
• the same update captured in a CUDA graph and replayed over a dirtied dummy slot;
• prepare() writes only evicted slots.

On GB200 all 3 pass; the two dummy-slot tests fail without the zero_(). speculative/hw_agnostic/
has no new failures. l0_h100.yml collects the file through unittest/_torch/speculative/hw_agnostic; on
l0_cpu it skips (no GPU).

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

DFlashWorker._advance_ctx_len increments and clamps context lengths, then resets the dummy slot to zero. prepare_1st_drafter_inputs uses this helper, keeping the reset in the drafting step and avoiding dummy-slot writes in prepare(). The change also applies after the Kimi K3 fused context-K/V kernel. No public API changes are reported.

QA Engineer Review

Added three CUDA-gated tests in test_dflash_dummy_slot.py. They cover padded-slot mapping and repeated advances, dummy-slot reset during captured CUDA-graph replay, and prepare() writing only evicted slots. The file is included by the hw_agnostic test-list entry in l0_h100.yml; l0_cpu.yml also selects the suite, but these tests skip when CUDA is unavailable. Coverage verdict: sufficient for the dummy-slot reset behavior. Test execution status is not established by the supplied evidence.

Per-File QA Perspective

• tensorrt_llm/_torch/speculative/dflash.py: Verify that real request context lengths still advance and clamp correctly, while the dummy slot resets after the drafting-step update, including the fused context-K/V path.
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py: Covers padding, captured-graph replay, and eviction-only writes in prepare(). The hw_agnostic suite is selected by l0_h100.yml and l0_cpu.yml; CUDA-unavailable runs skip these tests.

[coderabbitai]

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: 7f46f5a1-23aa-4461-999d-a5318629ca7b

📥 Commits

Reviewing files that changed from the base of the PR and between 1e0dceb and b12bfad.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/speculative/dflash.py
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py

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

---

Walkthrough

DFlash now uses _advance_ctx_len to update accepted-token context lengths, clamp them to _max_ctx, and reset the dummy slot. CUDA tests cover padding rows, graph replay, and cleanup of an evicted request slot.

Changes

DFlash dummy-slot preparation

Layer / File(s)	Summary
Reset and validate dummy-slot context length tensorrt_llm/_torch/speculative/dflash.py, tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py	_advance_ctx_len updates and clamps accepted-token context lengths, then resets the dummy slot. CUDA tests cover padded rows, graph replay, and evicted-slot cleanup.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Suggested reviewers: brnguyen2

Merge Risk: ⚪ Minimal · up to b12bf

The dummy-slot reset is in place for the inspected drafting paths. No issue requiring a fix before merge was identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.36% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 2 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’s [None][fix] format and clearly identifies the DFlash dummy-slot context-length fix.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the issue and fix, identifies the tests, and reports their coverage and results.

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

[longlee0622]

Thanks for tracking this down. The root cause matches what we have seen. For performance, though, I'd suggest doing the reset inside the drafting step rather than in prepare(): zero the dummy slot right after the in-graph accumulation in DFlashWorker:

                self._ctx_len[slots] += gen_num_accepted_long
                self._ctx_len.clamp_(max=self._max_ctx)
                # Padding and ADP dummy rows share a slot that is never freed; keep its
                # context empty. zero_() is graph-capturable, a scalar store is not.
                self._ctx_len[self._dummy_slot].zero_()

Why this is better for perf:

• No per-step host work. With the reset in prepare(), evicted is never empty, so _write_ctx_len builds pinned tensors and issues the H2D copies and index_copy_ on every step. Before this change it returned early on steps with no eviction. The zero_() is captured into the CUDA graph, so the host does nothing extra.
• Dummy rows attend no context in the current step. With the prepare() reset, the dummy slot starts at 0 but is incremented by this step's accepted tokens before num_ctx_per_req_t is read, so padding / ADP dummy rows still attend a short context whose drafts are thrown away. Zeroing after the += makes it 0, so the padded rows cost the least.
• Same correctness. This += is the only place the dummy slot's length grows. Zeroing it there keeps the length bounded, so the padding rows' context K / V write starts from column 0 every step and stays out of page 0.

We validated this variant end to end on an internal branch: padded-step drafter attention dropped from ~520 to ~45 us/layer, and tokens/s/user improved by +2.8% (helix16, bs16) and +17% (DEP16).

The unit tests would then need to drive the drafting-step update (or a small helper factored out of it) instead of DFlashSpecMetadata.prepare.

[vsabavat]

@juney-nvidia could you please /bot run on the current head (1e0dceb)? It has never run CI and #19841 waits on it. Thanks!

[vsabavat]

Thanks, agreed. It is pushed as a follow-up commit, b12bfad22c, on this branch (a fast-forward of the previous
head).

• The drafting step now calls _advance_ctx_len, which does the +=, the clamp, and then
  self._ctx_len[self._dummy_slot].zero_().
• prepare() again writes slot lengths only for evicted requests.

I checked the three points against the code:

• Only growth. On main the += in the drafting step is the only place the dummy slot grows:

  • _assign_slot hands out only range(max_batch);
  • prefill writes only assigned slots;
  • the warm-up and failed-forward rollbacks copy a snapshot back.

  On our Kimi K3 stack the fused context-K/V kernel (k3_ctx_kv) also adds the accepted tokens. That path gets the
  same zero_() right after the kernel, on the K3 side.

• Capture. _dummy_slot is a Python int and _ctx_len is allocated once, before any capture. So the zero_() is
  a fixed-address store that the graph replays. The new test captures the update and replays it over a dirtied slot.

• Host work. Confirmed: with the reset in prepare(), _write_ctx_len ran on every step. It now returns early
  unless a request was evicted.

The tests in test_dflash_dummy_slot.py now drive _advance_ctx_len, eager and under a captured graph, and check
that prepare() writes only evicted slots. They fail on the current head and pass with the change. The rest of
speculative/hw_agnostic gives the same results before and after.

Kimi K3 numbers: 16 GB200, DSpark with 7 drafts, acceptance forced to 6, ISL / OSL 1024, 5 reps per cell. Median TPOT
vs the current head:

• No padding (every batch size has its own graph, so dummy rows appear only at warm-up; A B B A):
  -0.27 / -0.75 / -0.18 / -0.73 % at batch 1 / 2 / 4 / 8. This is the removed per-step write.
• With enable_padding (graphs 1 / 2 / 4 / 8; concurrency 3 and 6, so 1-2 padding rows per step): -0.25 / -0.41 %.
  With no reset at all, the same cells are only +0.41 / +0.27 % slower than the current head, so we could not
  reproduce the ~520 us / layer here. Which configuration was it measured on (attention DP, block tables, drafter)?
• Outputs: the 8 fixed-prompt outputs are identical before and after. On GSM8K (all 1319, real acceptance, batch
  8), all 1319 responses are identical, 96.21 / 96.13 both times.

So the change is correct and a small win here.

[vsabavat]

@juney-nvidia could you please /bot run on the current head (b12bfad)? It addresses @longlee0622's review and #19841 waits on it. Thanks!

[vsabavat]

/bot run

[tensorrt-cicd]

PR_Github #76602 [ run ] triggered by Bot. Commit: b12bfad Link to invocation

[tensorrt-cicd]

PR_Github #76602 [ run ] completed with state SUCCESS. Commit: b12bfad
/LLM/main/L0_MergeRequest_PR pipeline #63165 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Either:

• Wait for the PR to be fully approved — the label is added automatically once approval is complete. Having unresolved open comments is fine, or
• If needed, ask a member of NVIDIA/trt-llm-ci-approvers to add the label manually.
  Then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[longlee0622]

/bot run

[tensorrt-cicd]

PR_Github #76625 [ run ] triggered by Bot. Commit: b12bfad Link to invocation

[tensorrt-cicd]

PR_Github #76625 [ run ] completed with state SUCCESS. Commit: b12bfad
/LLM/main/L0_MergeRequest_PR pipeline #63186 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[vsabavat]

Superseded by the consolidated draft PR: #19988. Source commits and follow-ups are retained; prior reviews remain historical evidence, not approval of the replacement heads.