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

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:

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 [None][feat] Add DeepSeek DSpark speculative decoding #15808 (DSpark), PyTorch backend, disaggregated serving, 8-GPU (DEP8) generation server, MEGAMOE_DEEPGEMM MoE.