[Qwen3.8-Next] Add PD state transfer for Flash Next - #36651
Conversation
|
Hi, thanks for this PR. We independently hit the exact crash this fixes and would like to share a verification on a hardware/scenario combination that's missing from the thread. Environment
Reproduced crash (without this fix)Decode server, during decode CUDA-graph capture at startup: Root cause confirmed: in decode mode, We confirmed this is not attention-backend dependent (persists with FA2 installed); it is a pure pool-initialization gap. Verification with this PR appliedApplied the pool-fix hunks from #36651 (the 4 params on
AskThis PR has been open for ~7 days and is exactly what unblocks Qwen3.8 PD on the decode side. Happy to help test anything else on the SM80 + RDMA setup. If useful, we can also attach logs for the pre-fix crash and post-fix metrics. |
|
Thanks for the detailed verification! The Qwen3.8 Flash NEXT rebase PR is expected to merge into cc @Qiaolin-Yu — if there is an urgent need to unblock this earlier, we could also consider merging it into the rebase branch first. What do you think? |
7bcf5ba to
39f01f5
Compare
YAMY1234
left a comment
There was a problem hiding this comment.
Focused review of the QSA PD state-transfer paths.
|
Reviewed at head OverallDirection looks right and the implementation is careful. The pieces that usually go wrong in this area Blocking / correctness concerns1. (needs confirmation) MORI transfers QSA state but is excluded from the GPUDirect visibility flush
The flush exists because a host-side RDMA poll (KVPoll.Success) does not make third-party GPU writes Could you confirm one of:
What I verified (so the questions above stay focused)
Questions for the author
Nits
|
The layer-id paired flat path uses the source item length as the destination stride. Pass the peer item lengths from maybe_send_extra and reject mismatched pairs so a future heterogeneous NIXL change cannot silently inherit the equal-layout assumption.
|
@Leslie360 Thanks for the review. The head has moved to Blocking (MORI transfers QSA state but is excluded from the flush). You are right that the inconsistency was real: the flush is a memory-ordering contract of the transport backend, not a property of a state type, so it should not have been scoped per state type. Tracing this back, no existing state type on Q1 (A800 flush support), Q5 (capture window), Nit 1 (widen the catch). These all concerned the flush path and are moot after its removal. Q2 (heterogeneous attention TP matrix). Both directions were validated end to end with exact output-token parity: TP1 -> TP4 and TP4 -> TP1, plus TP4 -> TP4 and TP1 -> TP1, with MTP and PLE-offload variants. See the Accuracy section of the PR description. PP > 1 is not reachable for this model today: the real checkpoint hits the existing PLE PP weight-layout path and dummy weights hit the explicit Q3 (page-size symmetry). Enforced by the existing generic PD gate: when the decode side fetches the prefill server info from the bootstrap server, a Q4 (prefix-cache hit and the pending ring). The compressed QSA pool requires Nit 2 (NIXL paired item-length check). Agreed it is safe today because NIXL rejects heterogeneous TP for non-MLA state. Added the guard anyway in |
|
/rerun-group disaggregation |
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 🚀 🚀 |
ShangmingCai
left a comment
There was a problem hiding this comment.
LGTM, logic is clean and correct
|
/rerun-test registered/kernel/qsa/test_qsa.py |
|
Results for 🚀 |
# Conflicts: # python/sglang/srt/disaggregation/mooncake/conn.py # python/sglang/srt/disaggregation/nixl/conn.py
|
/rerun-test registered/mock_model/test_e2e_pd.py |
|
Results for 🚀 |
# Conflicts: # test/registered/unit/disaggregation/test_disaggregation_wire.py
|
/rerun-failed-ci bypass-fastfail 9851fb3 |
|
All NV CI Passed |
|
Independent verification of the merged result on 2×8 A800 (SM80) — works end-to-end, plus three SM80/Mooncake bring-up gotchas We brought up PD for Qwen4-Exp (BF16 Qwen3.8-Flash-Next) on 16×A800-SXM4-80GB (SM80), TP8 per side, prefill and decode on separate nodes, Mooncake 0.3.9 over mlx5 IB, against current main (@ 03ea13a) carrying this change. Data points in case they are useful for this model family:
Three environment gotchas that cost us most of the bring-up time (none of them sglang bugs):
Happy to PR these notes into the PD cookbook/docs if wanted. |
|
Re-verified on merged main (03ea13a), SM80 2×8×A800 TP8-per-side PD over mooncake/RDMA. Correction to our comment above: those earlier numbers were measured on the pre-merge base (b852d33, the #38642 branch head carrying an equivalent guard); the figures below are the re-run on merged main and supersede them. Following the merge, we re-ran the Qwen4-Exp PD suite on the current main to confirm the fix landed intact and to close the one open question we had flagged earlier. 1. Source-level answer to our earlier "where does decode-side n-gram history come from?" question: 2. Hardware re-verify (new main @ 03ea13a):
3. New regression we hit (needs a maintainer look): Everything else from our earlier report holds on merged main. Happy to share logs. |
Follow-up to the Qwen 3.8 Flash Next base support merged in #37500 (superseding #36497).
This PR is now based on
mainand keeps the diff scoped to PD disaggregation.Motivation
Qwen 3.8 Flash Next has non-KV recurrent and sparse-attention state that must move from prefill to decode together with the normal KV cache. Without an explicit state protocol, PD disaggregation cannot preserve aggregate-mode outputs.
Modifications
The core implementation is split into two commits, followed by review/refinement commits:
Matching attention TP
req_pool_idx, preserving the per-requestcompress_ratioring.Heterogeneous attention TP follow-up
Review follow-up
req.kv.req_pool_idx) on both prefill and decode.Final review tightening
dim == 0protocol marker; NIXL/MORI no longer import a Qwen-specific sentinel.num_tokens_per_req.The focused regressions are integrated into the existing QSA, disaggregation-wire/NIXL, decode-queue, Mamba state-buffer, model-override, and CPU QSA-pool test owners.
Accuracy Tests
The two PD commits were ported onto current
mainafter #37500 landed.Implementation validation was rerun across the review heads. The current no-flush head is
58581bfcc2cbf71c4b8ddf787c0757b1c3409b11:py_compile, andgit diff --checkpass.34326673927, job102385676706).3637841). The moved CPU test passes both direct pytest and the CI-stylepython test_file.py -fentry: 1 passed in each mode (job3638023, exit0:0).support_pp=Falseguard. The reported race is therefore unreachable today; the exact registration and NIXL no-op paths are covered by focused regressions rather than claimed as an E2E pass.test_umbp_store.pymock raisingTypeErroronMagicMock.numel() > 0; an unrelated concurrent PR reproduced the same test and exception, so this is a current ROCm CI baseline failure rather than a change in this PR (runs34309068072and34308935616).The no-flush final head
58581bfcc2cbf71c4b8ddf787c0757b1c3409b11was rerun end to end with the same production full profile:0.00531830545514822(mean0.0000349505220436939). The same maximum appears between the old and new aggregate repeats at the same token, while the prior flush head vs. no-flush PD differs by at most0.000473120715469123, identifying the outlier as aggregate run-to-run numerical variation rather than a no-flush PD change.0:0(jobs3639656/3639658), the probe records the exact final source SHA, and SHA-256 comparison confirmed that every changed file in the remote execution snapshot matched the local final head.3639632, exit0:0).The PP=1 production matrix was run at the immediate pre-follow-up head
a03da1ea59877363e7e6d8a9b6dd0fbc7296c12e; the subsequent production deltas are the review fixes above plus behavior-preserving cleanup around metadata generation and queue dispatch:0.000640055630356(aggregate/PD jobs3635179/3635180).0:0with no request or state-transfer errors.The heterogeneous and compatibility results below are from the prior stacked head:
3280031/3280032).3280185/3280187).3280185/3280188).3280198/3280199).3280198/3280370).3280184/3280182and3280183/3280190).All accepted jobs above completed without request errors.
Checklist
CI States
Latest PR Test (Base): ✅ Run #34569098711
Latest PR Test (Extra): ✅ Run #34606029716
Latest PR Test (AMD ROCm 10): ❌ Run #34569098683