Skip to content

[PD] Fix TP-skew decode polling and Mooncake staging deadlocks - #36893

Closed
YAMY1234 wants to merge 4 commits into
sgl-project:mainfrom
YAMY1234:agentx/c1-pd-main-fix-20260828
Closed

YAMY1234 wants to merge 4 commits into
sgl-project:mainfrom
YAMY1234:agentx/c1-pd-main-fix-20260828

Conversation

@YAMY1234

@YAMY1234 YAMY1234 commented Aug 28, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation / 动机

A low-concurrency disaggregated configuration (prefill TP1/DEP4, decode TP8, attention TP gather disabled) showed strong native batch-size-1 CUDA graph performance, but longer workloads could permanently stop making request progress.

低并发 PD 分离配置(prefill TP1/DEP4、decode TP8、关闭 attention TP gather)可以获得明显的原生 batch-size-1 CUDA graph 性能收益,但长时间 workload 会永久停止请求推进。

The hang came from three synchronization/state-machine bugs:

该 hang 由三个同步与状态机问题共同造成:

  • Decode KV polling and request broadcast could enter different collective epochs across TP ranks.
  • A deferred heterogeneous-TP Mooncake fanout could replay destinations that had already completed.
  • An empty staging ring could retain a stale head, while a new watermark subscriber could miss the allocator's current state.

对应地:decode KV polling 与请求广播可能在不同 TP rank 上进入不同 collective epoch;异构 TP 的 Mooncake fanout 在 defer 后可能重放已完成 destination;空 staging ring 可能保留旧 head,而新 watermark subscriber 又可能错过 allocator 当前状态。

Changes / 修改

  • Serialize decode polling and request broadcast through one ordered coordinator, using fixed-shape collective payloads and task-bound results.

  • Track completed fanout destinations across requeue so staging transfers are idempotent.

  • Reset an empty staging ring to a fresh round and bootstrap new subscribers with the current watermark.

  • Add focused regression tests for collective ordering, fixed-shape polling, fanout retry, ring wrap, and subscriber bootstrap.

  • 通过单一有序 coordinator 串行化 decode polling 与请求广播,并使用固定形状 collective payload 和任务绑定结果。

  • 在 requeue 期间记录已完成的 fanout destination,保证 staging transfer 幂等。

  • 空 ring 进入新 round 时重置状态,并用当前 watermark 初始化新 subscriber。

  • 为 collective ordering、固定形状 polling、fanout retry、ring wrap 和 subscriber bootstrap 增加定向回归测试。

Validation / 验证

  • Linux targeted tests: 18/18 passed.

  • Exact long-duration workload: 807/807 requests, 0 errors, 0 cancellations, 3630.00 s.

  • TTFT and inter-token-latency coverage: 100% / 100%.

  • P50/P90 output-token throughput per user: 467.70/447.07 → 640.33/685.20 tok/s/user (+36.9% / +53.3%).

  • No model forward or sampling math is changed.

  • Linux 定向测试:18/18 通过。

  • 精确长时间 workload:807/807 requests、0 errors、0 cancellations、3630.00 秒。

  • TTFT 与 inter-token-latency coverage:100% / 100%。

  • P50/P90 每用户输出 token 吞吐:467.70/447.07 → 640.33/685.20 tok/s/user(+36.9% / +53.3%)。

  • 未修改模型 forward 或 sampling 数学逻辑。


CI States

Latest PR Test (Base): ⏳ Run #33196082438
Latest PR Test (Extra): ❌ Run #33196082037
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33196083112

@YAMY1234

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant