NVIDIA / NVIDIA/TensorRT-LLM

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

Open
#17,095 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Disaggregated serving Speculative Decoding
Dominant language
Python
Stars
14.7k
Forks
2.8k
Avg merge
2d 23h
Merged PRs (30d)
489

Description

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.

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 py_executor.py with _sync_gen_only_benchmark_has_insufficient_kv, _allgather_model_parallel_status, and _is_benchmark_disagg_fill_complete; reproduce the asymmetric speculative-decoding gen-only burst case with attention-DP. Compare the collective entry behavior across ranks and then verify that the affected benchmark completes without a HangDetector timeout or MPI_Abort, including when the fill gate is disabled.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.