[#16767][fix] Fix DSpark rolling-window slot collision in disaggregated serving#16772
[#16767][fix] Fix DSpark rolling-window slot collision in disaggregated serving#16772longlee0622 wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughDSpark adds a dedicated scratch rolling-window row, expands lazy-init buffers, classifies dummy and unknown request IDs, and assigns persistent slots to disaggregated generation requests. Tests cover scratch routing, slot persistence, and dummy handling. ChangesDSpark slot mapping
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
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 (2)
tensorrt_llm/_torch/speculative/dspark.py (1)
284-293: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNit: comment says "runtime_draft_len" but code uses
self.max_draft_len.Harmless here since
block_size == max_draft_lenis enforced for DSpark, but worth aligning the wording to avoid confusion for future readers.🤖 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 `@tensorrt_llm/_torch/speculative/dspark.py` around lines 284 - 293, Update the comment above `_graph_dummy_id_floor` to describe the padding range using `self.max_draft_len`/the configured maximum draft length instead of “runtime_draft_len,” matching the calculation and preserving the existing routing explanation.tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py (1)
115-357: 🎯 Functional Correctness | 🔵 TrivialQA test coverage summary.
Changed/added test functions (all in
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py, outsidetests/integration/test_lists/):
test_worker_lazy_init_window_buffers(modified) — asserts the new scratch-row buffer shapes,_scratch_slot, and_graph_dummy_id_floor.test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(new) — unknown/padded ids route to the scratch row, not slot 0.test_prepare_assigns_slots_to_disagg_generation_requests(new) — disagg gen requests get distinct, stable persistent slots.test_prepare_keeps_dummy_generation_requests_on_scratch_row(new) — ADP-idle and CUDA-graph-padding dummy ids never consume a real slot (boundary-tests_graph_dummy_id_floorexactly).These are hardware-agnostic unit tests under
tests/unittest/, not undertests/integration/test_lists/, so notest-db//qa/list entries apply; they run via standard pytest CI collection.Coverage verdict: sufficient for the four areas the PR objectives explicitly claim (scratch-row routing, disaggregated slot assignment, dummy-request handling, buffer initialization). Gaps not covered here (acceptable given stated scope): a mixed batch with simultaneous context and disagg-generation requests in one
prepare()call, and the_assign_slotslot-exhaustion (RuntimeError) path. Per PR objectives, end-to-end DEP8 validation is explicitly still pending and should be tracked separately.🤖 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 115 - 357, The review identifies the added DSpark worker tests as sufficient for scratch-row routing, disaggregated slot assignment, dummy handling, and buffer initialization. Preserve these tests and their current scope; no test-list entries are needed, and do not add coverage for the explicitly accepted mixed-batch or slot-exhaustion gaps.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 `@tensorrt_llm/_torch/speculative/dspark.py`:
- Around line 284-293: Update the comment above `_graph_dummy_id_floor` to
describe the padding range using `self.max_draft_len`/the configured maximum
draft length instead of “runtime_draft_len,” matching the calculation and
preserving the existing routing explanation.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py`:
- Around line 115-357: The review identifies the added DSpark worker tests as
sufficient for scratch-row routing, disaggregated slot assignment, dummy
handling, and buffer initialization. Preserve these tests and their current
scope; no test-list entries are needed, and do not add coverage for the
explicitly accepted mixed-batch or slot-exhaustion gaps.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0b2c5d73-ec05-493f-8991-5936ef074520
📒 Files selected for processing (2)
tensorrt_llm/_torch/speculative/dspark.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py
…ow row DSparkSpecMetadata.prepare() mapped unknown request IDs (CUDA-graph padding, ADP idle requests, or disagg seed forwards without a real id) to slot 0, letting them overwrite the rolling KV window of whichever real request owned slot 0. Allocate one extra scratch window row at index max_batch (never handed out via _free_slots) and map unknown IDs to it. Adds a regression test asserting the live slot's window and ctx_len stay untouched (validated on B200). Related to GitHub NVIDIA#16767; disagg real-request-id propagation into the DSpark seed forward remains open. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
…equests DSpark assigns per-request rolling-window slots only in _seed_context_windows, which runs only when num_contexts>0. On the disaggregated generation server the prompt is prefilled (and the window seeded) on the context server, so _seed_context_windows never runs there, _req_to_slot stays empty, and every real generation request falls through to the single shared scratch/slot-0 row. At batch 1 that row is effectively private (accept length fine), but at batch >1 all concurrent requests collide on it and corrupt each other's draft window, collapsing accept length (~1.4 vs ~4.5 in a DEP8 repro). This is GitHub NVIDIA#16767; the reporter's "request_id=0 at seed" was the incidental ADP-idle dummy id, not the mechanism. Fix: DSparkSpecMetadata.prepare() now assigns each real generation-tail request its own persistent slot when seeding never mapped it, skipping the ADP-idle (id 0) and CUDA-graph padding dummies (which stay on the scratch row). Context-prefix entries are still left to _seed_context_windows, so aggregated mode is unchanged. Adds hw-agnostic regression tests (validated on B200). Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
d7c8c7f to
ee5f4b5
Compare
|
/bot run |
|
PR_Github #61209 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
The slot-map maintenance in prepare() is a near line-by-line duplicate of dflash.py:86-113. Consider extracting this shared logic into SpecWorkerBase as a follow up work.
Consolidated dev-branch notes for the DSpark implementation: current status, algorithm/implementation, validated behavior, and the review backlog. Records the rolling-window slot fixes (scratch-row aliasing + disagg gen-request slot assignment, GitHub NVIDIA#16767 / PR NVIDIA#16772), confirmed end-to-end at draft length 7 on disaggregated serving. Dev-branch only; not part of any PR. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
|
PR_Github #61209 [ run ] completed with state
|
3 participants
Dev Engineer Review
DSparkWorker._scratch_slotand changed lazy allocation of rolling-window KV/context tensors to include one extra scratch row (max_batch + 1), with CUDA-graph dummy/padding IDs routed below a computed dummy-ID floor to the scratch row.DSparkSpecMetadata.prepare()host-side slot assignment:_batch_to_slotno longer defaults missing/unknown request IDs to slot 0; it maps them to the scratch slot.worker._assign_slot(..., reset=False).QA Engineer Review
Unit test updates in
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py:test_prepare_maps_unknown_request_to_scratch_row_not_slot_zerotest_prepare_assigns_slots_to_disagg_generation_requeststest_prepare_keeps_dummy_generation_requests_on_scratch_rowThese validate scratch-row routing, distinct persistent slot assignment for unseeded disaggregated generation requests, dummy-request exclusion from real slot allocation, and updated buffer initialization expectations.
Verdict: sufficient for the DSpark slot-routing logic; needs follow-up for DEP8 end-to-end validation and any required CI/manual-QA test-list registration (not indicated).
Description
Fixes the DSpark speculative-decoding accept-length (AL) collapse under disaggregated serving reported in #16767.
Root cause. DSpark assigns each request a rolling-KV-window slot only inside
_seed_context_windows, which runs only whennum_contexts > 0. On the disaggregated generation server, transferred requests arrive as generation requests (prefill runs on the context server), so_seed_context_windowsnever runs there,DSparkWorker._req_to_slotstays empty, and every real generation request falls through to a single shared fallback rolling-window row. At generation batch size 1 that row is effectively private (AL correct), but at batch size > 1 all concurrent requests collide on it and corrupt each other's draft context, collapsing AL (e.g. ~1.4 vs ~4.5 at draft length 7 in a DEP8 repro). Therequest_id=0noted in the issue is the incidental attention-DP idle-dummy id (ATTENTION_DP_DUMMY_REQUEST_ID), not the mechanism.This PR has two commits:
DSparkWorkerallocates one extra scratch window row (indexmax_batch, never handed out through_free_slots);DSparkSpecMetadata.prepare()routes padded/unknown request IDs (CUDA-graph padding, ADP idle) to it instead of aliasing slot 0.prepare()now assigns each real generation-tail request its own persistent rolling-window slot when seeding never mapped it, skipping the ADP-idle (id 0) and CUDA-graph padding dummies (which stay on the scratch row). Context-prefix entries are still left to_seed_context_windows, so aggregated mode is unchanged.The generation server never prompt-seeds the window in either mode (batch-1 disagg also runs unseeded and reaches AL ~4.5, building the window from generation back-fill), so unique per-request slots fully address the collapse — no KV/window transfer from the context server is required.
Test Coverage
Hardware-agnostic unit tests in
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py(all pass on B200):test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero— unknown/padding IDs route to the scratch row, live slot untouched.test_prepare_assigns_slots_to_disagg_generation_requests— disagg gen requests (never seeded here) get distinct, stable, non-scratch slots.test_prepare_keeps_dummy_generation_requests_on_scratch_row— ADP-idle (id 0) and CUDA-graph padding dummies never consume a real slot.test_worker_lazy_init_window_buffers— scratch row + dummy-id floor.Validation status
PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.