[https://nvbugs/6535874][fix] Size DFlash/DSpark worker slot buffers from num_seq_slots

[trtllm-agent]

Summary

• Root cause: DFlashWorker._lazy_init and DSparkWorker._lazy_init sized their once-allocated, worker-owned context/window buffers from spec_metadata.max_num_requests, but create_cuda_graph_metadata shrinks that field to the captured graph bucket. Because those buffers are allocated once and then reused for every subsequent batch shape, whichever CUDA-graph bucket drafts first would pin the slot pool to its own (smaller) size, leaving _dummy_slot aliasing a live request's row instead of spanning the full seq-slot pool. The reported ValueError from the SSM state-pool floor was independently fixed upstream by 7f7dccf991, so the target test already passes on HEAD.
• Fix: num_seq_slots is now threaded through get_spec_metadata into both DFlash and DSpark metadata, and each worker sizes its buffer from spec_metadata.num_seq_slots or spec_metadata.max_num_requests — an inline read rather than a class-level slot_capacity property, which would have broken two unwaived l0_h100.yml tests (a duck-typed stub fixture, and test_rejection_buffers_guard.py, which calls prepare_rejection_sampling_buffers unbound on a SimpleNamespace where a property is unreachable). The DSpark worker unit test was switched from a sparse SimpleNamespace to the real metadata factory so newly added sizing fields cannot be silently omitted, and the two test_dflash_qwen3_5_4b waiver lines are removed now that both params pass.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6535874

Dev Engineer Review

• DFlash and DSpark now size slot pools from num_seq_slots, with fallback to max_num_requests.
• Metadata creation now preserves num_seq_slots across CUDA graph metadata copies.
• The change prevents buffer overlap with live request rows, the context window, and the dummy slot.
• Existing paths retain the current pool size when num_seq_slots is unavailable.
• The removed DFlash waivers target tests already fixed upstream.
• No public API declarations changed.
• No configuration issues or unintended scope changes were identified.

QA Engineer Review

• Modified test: the DSpark context-window seeding test now uses DSparkSpecMetadata from the real metadata factory.
• The test now provides capture buffers through metadata and removes the per-chunk get_hidden_states override.
• The modified test is covered by the existing integration test configuration.
• Verdict: sufficient.

[coderabbitai]

Walkthrough

DFlash and DSpark now receive sequence-slot metadata and size internal storage from that capacity, with max_num_requests as fallback. The DSpark test now uses real metadata, and two obsolete test waivers were removed.

Changes

Speculative decoding slot-pool sizing

Layer / File(s)	Summary
Sequence-slot metadata propagation tensorrt_llm/_torch/speculative/utils.py	DFlash and DSpark metadata now receives num_seq_slots.
Worker storage sizing and validation tensorrt_llm/_torch/speculative/dflash.py, tensorrt_llm/_torch/speculative/dspark.py, tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py, tests/integration/test_lists/waives.txt	DFlash context buffers and DSpark rolling-window storage use the configured slot pool, with max_num_requests as fallback. The DSpark test uses real metadata, and two obsolete waivers were removed.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#15775: Both changes modify DFlash speculative decoding and metadata handling.
• NVIDIA/TensorRT-LLM#17041: Both changes modify waiver entries for test_dflash_qwen3_5_4b.

Suggested reviewers: qijune, bowenfu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug, affected workers, and buffer-sizing fix.
Description check	✅ Passed	The description explains the root cause, fix, affected tests, regression checks, and linked bug.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py (1)

195-200: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Exercise the num_seq_slots allocation path.

• Modified test: test_seed_context_windows_preserves_state_across_prefill_chunks.
• Test-list membership: covered by tests/integration/test_lists/test-db/l0_h100.yml through unittest/_torch/speculative/hw_agnostic.
• Coverage verdict: insufficient. _make_metadata() does not expose num_seq_slots, so this test covers only the max_num_requests fallback.
• Add num_seq_slots=4 with max_num_requests=1. Assert _kv_windows.shape[0] == 5, _batch_to_slot.shape[0] == 4, and _scratch_slot == 4.

🤖 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 `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py` around
lines 195 - 200, Update
test_seed_context_windows_preserves_state_across_prefill_chunks to construct
metadata with num_seq_slots=4 alongside max_num_requests=1, exercising the
slot-allocation path instead of its fallback. Add assertions that
_kv_windows.shape[0] is 5, _batch_to_slot.shape[0] is 4, and _scratch_slot
equals 4.

Source: Path instructions

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

Nitpick comments:
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py`:
- Around line 195-200: Update
test_seed_context_windows_preserves_state_across_prefill_chunks to construct
metadata with num_seq_slots=4 alongside max_num_requests=1, exercising the
slot-allocation path instead of its fallback. Add assertions that
_kv_windows.shape[0] is 5, _batch_to_slot.shape[0] is 4, and _scratch_slot
equals 4.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0d3549f4-1254-41d0-a83e-773695ac6aa0

📥 Commits

Reviewing files that changed from the base of the PR and between 09b6bea and 2749a0a.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/speculative/dflash.py
• tensorrt_llm/_torch/speculative/dspark.py
• tensorrt_llm/_torch/speculative/utils.py
• tests/integration/test_lists/waives.txt
• tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt