#17095·TensorRT-LLM

[Bug] DSpark: disaggregated gen-only benchmark deadlocks under attention-DP with asymmetric (gen-only) speculative decoding

Author: lishicheng1996-nvCreated Jul 31, 2026Updated Sep 18, 2026
LabelsSpeculative DecodingDisaggregated serving

Description

In disaggregated serving (1 ctx + 1 gen) with attention-DP on the generation server and asymmetric speculative decodingspeculative_config present only on the GEN worker (first-token draft produced on GEN; the CTX runs no draft) — the gen-only benchmark deadlocks once draft_len >= 3. The GEN admits ~1 request, then the PyExecutor HangDetector fires after 300 s -> MPI_Abort (exit 137). draft_len 1-2 complete normally.

Observed with the DSpark block speculative-decoding backend, but the deadlock is in generic attention-DP + disagg gen-only code paths. The symmetric config (speculative_config on both the ctx and gen workers) works correctly — only the asymmetric (gen-only draft) config deadlocks.

Root cause (from investigation)

Under attention-DP each rank owns its own requests. The serial CTX prefill (ctx max_batch=1) delivers transferred requests to different GEN ranks on different iterations, and in the asymmetric config the GEN runs a per-request re-context/seed forward (GEN iter-log shows num_ctx_requests=1, num_ctx_tokens=128). So per-rank batches diverge — some ranks doing a context/seed forward, some decoding, some idle. Any per-iteration attention-DP collective then receives mismatched shapes and deadlocks:

  1. First it hangs at the gen-only-benchmark fill-gate allgather — _sync_gen_only_benchmark_has_insufficient_kv -> _allgather_model_parallel_status (tp_allgather) in py_executor.py. Unlike its sibling _is_benchmark_disagg_fill_complete (explicitly rank-safe — every rank always enters the collective, per its own comment), this one returns early per-rank on a local flag (_benchmark_fill_phase_active) and can diverge.
  2. With the fill gate disabled (TLLM_BENCHMARK_REQ_QUEUES_SIZE=0), the hang moves to tp_allgather(scheduled_batch.batch_size), which is gated on enable_attention_dp (the else branch has no collective) — confirming the fault is generic attention-DP per-rank divergence, not the fill gate specifically.

Confirmed not caused by: tokens_per_block, dep width (dep4 hangs too), max_num_tokens, the rolling-window slot handling, or the benchmark client. Symmetric spec (speculative_config on both ctx+gen) does not diverge — requests arrive decode-ready and ranks stay in lockstep.

Steps to reproduce

  • Disaggregated 1 ctx + 1 gen, gen-only burst benchmark (all requests submitted up front).
  • Generation server: attention-DP (e.g. tp8, enable_attention_dp=true) + speculative_config on the GEN worker only (asymmetric / first-token-draft offload), draft_len >= 3.
  • Result: GEN admits ~1 request, then 300 s HangDetector -> MPI_Abort.

Suggested fix direction

  • Make the per-rank request mix rank-consistent under attention-DP for asymmetric spec-decode, or make the affected collectives rank-safe (mirror _is_benchmark_disagg_fill_complete: every model-parallel rank enters/skips the collective together).
  • A continuous/rate-based client (instead of a burst gen-only benchmark) reduces the per-rank skew as a workaround.