NVIDIA / NVIDIA/TensorRT-LLM

[Bug] DSpark speculative decoding: accept length collapses to ~1 at generation batch size > 1 in disaggregated serving

Open
#16,767 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Disaggregated serving Pytorch Speculative Decoding
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:

  1. Every seed overwrites the single _req_to_slot[0] entry (observed: _req_to_slot never holds more than ~2 entries, regardless of batch size).
  2. 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 at dspark.py: "Unknown request IDs ... default to slot 0").
  3. 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

  1. Serve deepseek-ai/DeepSeek-V4-Pro-DSpark disaggregated (separate context + generation servers), DEP8, speculative_config.decoding_type=DSpark, max_draft_len=block_size (e.g. 5 or 7), MEGAMOE_DEEPGEMM on the generation server.
  2. Benchmark at concurrency 8 (generation batch 1/rank) and at concurrency 128 (batch 16/rank).
  3. Compare avg_decoded_tokens_per_iter.mean from 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_DEEPGEMM MoE.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.