[Bug] DSpark speculative decoding: accept length collapses to ~1 at generation batch size > 1 in disaggregated serving
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Description
DSpark speculative decoding (added in #15808) produces the expected accept length only at generation batch size 1. As soon as the disaggregated generation server batches more than one request, the measured accept length collapses toward ~1.0 (drafts are almost never accepted), so DSpark provides essentially no speedup under real serving load.
Measured on disaggregated gen-only serving (DEP8 = TP8/EP8 with attention-DP, MEGAMOE_DEEPGEMM MoE on the generation server, deepseek-ai/DeepSeek-V4-Pro-DSpark, real ~39k-token prompts). AL = avg_decoded_tokens_per_iter.mean from the serving result:
| draft len | AL @ concurrency 8 (batch 1 / rank) | AL @ concurrency 128 (batch 16 / rank) |
|---|---|---|
| 1 | 1.86 | 1.20 |
| 3 | 3.25 | 1.44 |
| 5 | 4.03 | 1.52 |
| 7 | 4.51 | 1.49 |
Accept length is a per-request property and should be batch-invariant. MTP on the identical setup is ~2.85 at both batch 1 and batch 16. DSpark's collapse is therefore a bug, not expected behavior.
Root cause
tensorrt_llm/_torch/speculative/dspark.py (DSparkWorker) keys its per-request rolling-KV-window "slots" (_kv_windows[max_batch, ...], _req_to_slot, _batch_to_slot) on spec_metadata.request_ids.
On the disaggregated generation server, the context/seed forwards that seed each request's draft window (_seed_context_windows -> _assign_slot) are passed request_id=0 for (nearly) every request instead of the real request id. At batch > 1 this means:
- Every seed overwrites the single
_req_to_slot[0]entry (observed:_req_to_slotnever holds more than ~2 entries, regardless of batch size). - At generation time the batch's real, distinct request ids miss the map, so
_batch_to_slot = _req_to_slot.get(rid, 0)returns slot 0 for all but one request (cf. the comment atdspark.py: "Unknown request IDs ... default to slot 0"). - All batched requests collide on slot 0 with an un-seeded / corrupted draft window -> the target rejects every draft -> AL ~ 1.
Instrumentation at generation batch 16, DSpark draft 7:
slots = [9, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] # 15/16 collide on slot 0
nacc = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] # every draft rejected
with repeated _assign_slot(req_id=0, reset=True, ...) for the seed forwards.
Ruled out: the metric itself (it is the legitimate per-request avg_decoded_tokens_per_iter), the SpeculationGate (inactive — no acceptance_rate_threshold/window configured), max_num_tokens (512 vs 128: no change), and CUDA graphs (reproduces in eager). Assigning fresh distinct slots at generation time removes the collision (slots become distinct) but AL still stays ~1.5, because the windows are still never seeded under the real ids — confirming the root cause is the request_id=0 plumbing, upstream of dspark.py.
Steps to reproduce
- Serve
deepseek-ai/DeepSeek-V4-Pro-DSparkdisaggregated (separate context + generation servers), DEP8,speculative_config.decoding_type=DSpark,max_draft_len=block_size(e.g. 5 or 7),MEGAMOE_DEEPGEMMon the generation server. - Benchmark at concurrency 8 (generation batch 1/rank) and at concurrency 128 (batch 16/rank).
- Compare
avg_decoded_tokens_per_iter.meanfrom the two results.
Expected: roughly equal AL (batch-invariant). Actual: ~4.5 at batch 1, ~1.5 at batch 16.
Expected behavior
DSpark accept length should be batch-invariant (equal to the batch-1 values), retaining its advantage over MTP (~4.5 vs ~2.85 at draft 7) under real serving load rather than collapsing to ~1.
Suggested fix area
Propagate the real request ids into the DSpark spec metadata for the disaggregated generation server's context/seed forwards, so _seed_context_windows / _assign_slot key each request's rolling window on its actual request id instead of request_id=0.
Environment
- TensorRT-LLM containing #15808 (DSpark), PyTorch backend, disaggregated serving, 8-GPU (DEP8) generation server,
MEGAMOE_DEEPGEMMMoE.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in tensorrt_llm/_torch/speculative/dspark.py with DSparkWorker, then trace the disaggregated generation server's context/seed forwards into spec metadata. Inspect _seed_context_windows, _assign_slot, and request_id propagation, comparing seed and generation request IDs. Done means each request uses its real ID and accept length remains consistent at batch sizes above one.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai-infra-agents, distributed-systems, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100