[https://nvbugs/6535874][fix] Size DFlash/DSpark worker slot buffers from num_seq_slots - #17154
[https://nvbugs/6535874][fix] Size DFlash/DSpark worker slot buffers from num_seq_slots#17154trtllm-agent wants to merge 6 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughDFlash and DSpark now receive ChangesSpeculative decoding slot-pool sizing
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py (1)
195-200: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise the
num_seq_slotsallocation 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.ymlthroughunittest/_torch/speculative/hw_agnostic.- Coverage verdict: insufficient.
_make_metadata()does not exposenum_seq_slots, so this test covers only themax_num_requestsfallback.- Add
num_seq_slots=4withmax_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
📒 Files selected for processing (5)
tensorrt_llm/_torch/speculative/dflash.pytensorrt_llm/_torch/speculative/dspark.pytensorrt_llm/_torch/speculative/utils.pytests/integration/test_lists/waives.txttests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
…slot range The DFlash and DSpark workers allocate their per-request context pools once, lazily, on the first drafting forward and then reuse them for every later batch shape. Both sized the pool from spec_metadata.max_num_requests, which create_cuda_graph_metadata overwrites with the captured graph bucket size. Whichever bucket drafts first therefore pinned the pool to that bucket, leaving the scratch slot (_dummy_slot / _scratch_slot) inside the range real requests draw from, so CUDA-graph padding writes could overwrite a live request's context. Read num_seq_slots instead: it spans the full seq-slot pool and is left untouched by the graph copy, which is the same reason prepare_rejection_sampling_buffers derives its slot capacity from it. The DFlash and DSpark metadata constructions did not forward num_seq_slots, so the field defaulted to 0 and any bound derived from it silently collapsed back to max_num_requests; plumb it through as the MTP-Eagle path already does. Both halves are load-bearing. This does not change the pool size on either path today: num_seq_slots is non-None only under should_enable_dsv4_overlap_headroom, which requires is_mtp_eagle_one_model(), so it resolves to exactly max_num_requests here. test_seed_context_windows_preserves_state_across_prefill_chunks built its metadata as a bare SimpleNamespace carrying only three attributes, so it could not survive _lazy_init reading any further sizing field. Switch it to the real DSparkSpecMetadata via the file's own _make_metadata helper, as the other ten _lazy_init call sites in that file already do. The two test_dflash_qwen3_5_4b waivers are stale: the SSM-pool ValueError they were filed against was fixed by 7f7dccf, which landed after the waivers, so they were never re-evaluated. Remove them. Leaves batch_indices_cuda (metadata-owned, reallocated by the graph copy's __post_init__) and the eagle3 padding computation (must track the current batch shape) unchanged. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
2749a0a to
69f7361
Compare
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py (1)
40-75: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the new functions.
The coding guidelines require annotations for every function.
tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py#L40-L75: Add-> Noneto the test and precisetorch.Tensorannotations to theDraftModelmethods.tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py#L138-L138: Add-> Nonetotest_worker_graph_bucket_uses_full_seq_slot_pool.Proposed change
-def test_dflash_graph_bucket_uses_full_seq_slot_pool(): +def test_dflash_graph_bucket_uses_full_seq_slot_pool() -> None: ... - def _build_fused_kv_buffers(self): + def _build_fused_kv_buffers(self) -> None: pass - def project_target_hidden(self, hidden_states): + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states - def precompute_context_kv(self, hidden_states, position_ids): + def precompute_context_kv( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: ... -def test_worker_graph_bucket_uses_full_seq_slot_pool(): +def test_worker_graph_bucket_uses_full_seq_slot_pool() -> None:As per coding guidelines, “Annotate every function.” Based on learnings, this repository supports Python 3.10+ built-in generic annotations.
🤖 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_dflash.py` around lines 40 - 75, Annotate test_dflash_graph_bucket_uses_full_seq_slot_pool with -> None, and add precise torch.Tensor parameter and return annotations to DraftModel.project_target_hidden and DraftModel.precompute_context_kv in tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py:40-75. Also annotate test_worker_graph_bucket_uses_full_seq_slot_pool with -> None in tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py:138.Sources: Coding guidelines, Learnings
🤖 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_dflash.py`:
- Around line 40-75: Annotate test_dflash_graph_bucket_uses_full_seq_slot_pool
with -> None, and add precise torch.Tensor parameter and return annotations to
DraftModel.project_target_hidden and DraftModel.precompute_context_kv in
tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py:40-75. Also
annotate test_worker_graph_bucket_uses_full_seq_slot_pool with -> None in
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py:138.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 32c3e59a-8a3d-4d03-ae35-bd1f28af906a
📒 Files selected for processing (2)
tests/unittest/_torch/speculative/hw_agnostic/test_dflash.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py
|
PR_Github #63721 [ run ] triggered by Bot. Commit: |
|
PR_Github #63721 [ run ] completed with state
|
|
/bot run |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
PR_Github #64031 [ run ] triggered by Bot. Commit: |
|
PR_Github #64031 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The core fix looks right and the tests target the actual failure mode.
Main question is the unwaive. The PR says the reported ValueError was fixed upstream by 7f7dccf991, i.e. these two waivers are being removed for a reason unrelated to this PR's fix. Two things I'd want before merging that part:
- Evidence of
test_dflash_qwen3_5_4b[True]/[False]actually passing on the CI stage that runs them (they need a GPU +LLM_MODELS_ROOT; the new unit tests don't cover them). "Automated fix generated by repair-bot" plus a checked "verified on the same GPU type" box isn't enough on its own — please paste the run. - https://nvbugs/6535767 should be updated/closed. Waiver lines get regenerated from open NVBugs during triage, so an unwaive with a live bug comes back.
If the unwaive can't be substantiated right now, split it out — the sizing fix stands on its own and is worth landing separately.
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
Review follow-up: restored both Qwen3.5 DFlash waivers for NVBug 6535767. This PR now keeps the slot-buffer sizing fix and focused unit coverage; unwaiving remains gated on actual GPU CI evidence and the NVBug status. |
|
/bot run |
|
PR_Github #64287 [ run ] triggered by Bot. Commit: |
|
PR_Github #64287 [ run ] completed with state
|
|
CI #52189 completed. The only hard failure was the unrelated |
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
|
PR_Github #64465 [ run ] triggered by Bot. Commit: |
|
PR_Github #64465 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64543 [ run ] triggered by Bot. Commit: |
|
PR_Github #64543 [ run ] completed with state
|
| # shrunk to the captured graph bucket by create_cuda_graph_metadata, | ||
| # which would pin the pool to whichever bucket drafts first and leave | ||
| # _dummy_slot aliasing a live request's row. | ||
| max_batch = spec_metadata.num_seq_slots |
There was a problem hiding this comment.
num_seq_slots defaults to zero and documents max_num_requests as its fallback. Using the fallback inside _lazy_init is too late for CUDA-graph metadata because max_num_requests has already been reduced to the graph bucket. Could we normalize num_seq_slots in both metadata classes’ initial post_init calls instead? That preserves the original capacity through the shallow graph copy and makes both workers consume the same invariant. A direct-construction regression test should verify that metadata created with max_num_requests=5 and default num_seq_slots=0 retains num_seq_slots == 5 after create_cuda_graph_metadata(max_batch_size=2).
5 participants
Summary
DFlashWorker._lazy_init_ctx_buffersandDSparkWorker._lazy_initallocated worker-owned context/window buffers once fromspec_metadata.max_num_requests.create_cuda_graph_metadatashrinks that field to the captured graph bucket, so whichever bucket drafted first could permanently undersize the slot pool and make the dummy/scratch row overlap a live request slot.num_seq_slotsis threaded throughget_spec_metadatafor DFlash and DSpark. DFlash sizes its persistent buffers directly from the resolvedspec_metadata.num_seq_slots; DSpark uses that value with its existing compatibility fallback for metadata constructed outside the normal factory.main.Test plan
unittest/_torch/speculative/hw_agnostic.ci: full pre-merge approvedlabel is granted.Links