[None][fix] DFlash: start the dummy slot at context length 0 every step

[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.

prepare() now writes 0 to the dummy slot's context length every step, with the 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):

• padded steps keep the dummy slot at 0 and leave the real slots untouched;
• an evicted request and the dummy reset in one step.

Both fail without the reset: on GB200, 2 failed on main and 2 passed with this PR. 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

DFlashSpecMetadata.prepare now resets the dummy slot’s context length to zero on each step, alongside evicted slots. This prevents CUDA-graph padding and warmup dummy rows from extending context K/V writes into page 0, which may belong to a live request. The change does not alter public APIs.

QA Engineer Review

Added two CUDA-gated tests. They cover repeated padded steps with real-slot lengths preserved, and simultaneous reset of the dummy slot and an evicted request. The tests simulate slot updates; they do not run a model forward. Coverage verdict: sufficient for the reset behavior. The supplied results do not establish test execution status.

Per-File QA Perspective

• tensorrt_llm/_torch/speculative/dflash.py: Verify that prepare resets the dummy context length without changing real request slots. The change also applies when a step has no evicted requests.
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py: Covers padded slot mappings and resets alongside eviction. The hw_agnostic suite is included in the l0_h100.yml CI list; the test file skips when CUDA is unavailable.

[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: 08e1e43c-d324-455d-8456-c93a6b1d352a

📥 Commits

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

📒 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; 7 remain after this review.

---

Walkthrough

DFlash preparation now resets the dummy slot’s context length on each step. CUDA tests cover padding rows across acceptance steps 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	prepare includes the dummy slot in the context-length reset. CUDA tests check that its length remains zero across four acceptance steps and after an evicted request slot is cleared.

Priority: ➖ Normal

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

Change: Bug fix

Suggested reviewers: brnguyen2

Merge Risk: ⚪ Minimal · up to 1e0dc

The dummy-slot reset addresses the reported padding-row behavior, and no merge-blocking issue is established.

🚥 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 5 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 follows the repository format and clearly states the main change: resetting the DFlash dummy slot’s context length each step.
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 reported results, and marks the reviewer-assignment…
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.}