Skip to content

[Qwen3.8-Next] Add PD state transfer for Flash Next - #36651

Merged
YAMY1234 merged 14 commits into
sgl-project:mainfrom
YAMY1234:yamy/qwen38-next-pd-disagg-20260827
Sep 12, 2026
Merged

YAMY1234 merged 14 commits into
sgl-project:mainfrom
YAMY1234:yamy/qwen38-next-pd-disagg-20260827

Conversation

@YAMY1234

@YAMY1234 YAMY1234 commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator

Follow-up to the Qwen 3.8 Flash Next base support merged in #37500 (superseding #36497).

This PR is now based on main and 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:

  1. Matching attention TP

    • Transfer PLE short-convolution and n-gram state through the Mamba slot lifecycle.
    • Transfer QSA pending raw-key/RoPE state by req_pool_idx, preserving the per-request compress_ratio ring.
    • Transfer compressed QSA K by KV page ID.
    • Attach global full-attention layer metadata, including a request-wide RoPE sentinel, so compact QSA lists map correctly across PP layouts.
  2. Heterogeneous attention TP follow-up

    • Treat PLE/QSA state as replicated within attention TP.
    • Add Mooncake writer election for prefill attention-TP > decode attention-TP and broadcast for the reverse direction.
    • Validate exact transferred index lengths and layouts.
    • Keep NIXL and MORI fail-fast for heterogeneous Qwen 3.8 attention TP; mismatch support is Mooncake-only.
  3. Review follow-up

    • Read QSA pending-state ownership from the allocated request row (req.kv.req_pool_idx) on both prefill and decode.
    • Enable Qwen 3.8 PD in the model override. MORI QSA remains explicitly limited to PP=1 because its protocol does not exchange global state-layer metadata.
    • Do not register the request-wide QSA RoPE ring on a PP stage with no local QSA layers; keep empty state-component slots wire-aligned and make empty NIXL state plans a no-op.
  4. Final review tightening

    • Remove the QSA-only GPUDirect visibility flush after production no-flush parity; QSA now follows the same transfer-completion contract as existing KV, Mamba, and DSA state.
    • Let each slot sibling declare replicated transfer state through the existing dim == 0 protocol marker; NIXL/MORI no longer import a Qwen-specific sentinel.
    • Reuse the paged-state and PLE request-pool helpers, and move the CPU-only QSA allocation regression into the CPU unit suite.
    • Document why draft-extend falls back to 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 main after #37500 landed.

Implementation validation was rerun across the review heads. The current no-flush head is 58581bfcc2cbf71c4b8ddf787c0757b1c3409b11:

  • Changed-file Ruff lint/format, py_compile, and git diff --check pass.
  • Latest-head GitHub lint passed, and Base CPU CI passed all 9/9 partitions. The moved QSA allocation test was assigned to partition 4 and executed successfully (run 34326673927, job 102385676706).
  • Fixed-container targeted tests: 245 passed, 1 skipped, plus 51 subtests passed, across QSA kernel, disaggregation wire, NIXL, decode queue, Mamba/PLE slot lifecycle, CPU QSA-pool, and model-override tests (job 3637841). The moved CPU test passes both direct pytest and the CI-style python test_file.py -f entry: 1 passed in each mode (job 3638023, exit 0:0).
  • A zero-QSA-stage PP2 topology was attempted end to end, but the current model cannot enter serving with PP>1: the real checkpoint first hits the existing PLE PP weight-layout path, while dummy weights reach the explicit support_pp=False guard. 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.
  • Valid ROCm 7.2 stage-A CI passed 8 of 9 test files. Its sole failure is the untouched test_umbp_store.py mock raising TypeError on MagicMock.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 (runs 34309068072 and 34308935616).

The no-flush final head 58581bfcc2cbf71c4b8ddf787c0757b1c3409b11 was rerun end to end with the same production full profile:

  • Matching TP4 -> TP4 aggregate/1P1D Mooncake parity passed all 12 short/8K, OSL1/64, BS1/4/16 cases (84 requests / 2730 generated tokens) with 0 input-hash, output-token-sequence, or finish-reason mismatches.
  • The maximum aggregate-vs-PD selected-token logprob absolute delta was 0.00531830545514822 (mean 0.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 most 0.000473120715469123, identifying the outlier as aggregate run-to-run numerical variation rather than a no-flush PD change.
  • The aggregate and PD jobs both completed 0:0 (jobs 3639656 / 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.
  • The affected decode-queue, QSA-pool, and Mamba state-transfer tests passed in the fixed CUDA container: 14 passed (job 3639632, exit 0: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:

  • Matching TP4 -> TP4 full matrix: exact output-token parity across all 12 short/8K, OSL1/64, BS1/4/16 cases (84 requests / 2730 generated tokens); maximum selected-token logprob absolute error was 0.000640055630356 (aggregate/PD jobs 3635179/3635180).
  • Both jobs completed 0:0 with no request or state-transfer errors.

The heterogeneous and compatibility results below are from the prior stacked head:

  • Matching attention TP with MTP: exact output-token parity (3280031/3280032).
  • Mooncake heterogeneous attention TP:
    • TP1 -> TP1: exact output-token parity (3280185/3280187).
    • TP4 -> TP1: exact output-token parity (3280185/3280188).
    • TP4 -> TP4: exact output-token parity (3280198/3280199).
    • TP1 -> TP4: exact output-token parity (3280198/3280370).
    • MTP TP1 -> TP4 and PLE-offload TP4 -> TP1 also have exact output-token parity (3280184/3280182 and 3280183/3280190).

All accepted jobs above completed without request errors.

Checklist

  • Format and static checks for the changed files.
  • Focused unit tests integrated into existing test owners.
  • Accuracy validation for matching and heterogeneous attention TP.
  • Production no-flush parity comparison; no acceleration claim.
  • No user-facing API or documentation change.
  • Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ✅ Run #34569098711
Latest PR Test (Extra): ✅ Run #34606029716
Latest PR Test (AMD ROCm 10): ❌ Run #34569098683

@Leslie360

Copy link
Copy Markdown

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:

File ".../qwen4_exp.py", line 252, in _prepare_ple_batch
    history = get_req_to_token_pool().get_ngram_context(state_indices)
File ".../memory_pool.py", line 1488, in get_ngram_context
    return self.ngram_pool.get_context(ngram_indices)
File ".../ple_state_pool.py", line 183, in get_context
    assert self.context is not None
AssertionError

Root cause confirmed: in decode mode, HybridMambaDecodeReqToTokenPool (disaggregation/decode.py) was built without the Qwen4 PLE n-gram kwargs — ngram_context_len fell back to 0, so NGramPool.context was never allocated — while qwen4_exp unconditionally reads the shared n-gram history on decode. Standalone/prefill (which use HybridReqToTokenPool) didn't crash, which is why this is PD-decode-only.

We confirmed this is not attention-backend dependent (persists with FA2 installed); it is a pure pool-initialization gap.

Verification with this PR applied

Applied the pool-fix hunks from #36651 (the 4 params on HybridMambaDecodeReqToTokenPool.__init__ + ple_kwargs in _build_hybrid_mamba_decode_req_pool), re-ran the PD stack on the same 8×A800 + RDMA:

  • decode health 200, PD end-to-end 2+2=4 OK
  • transfer confirmed over RDMA (installTransport type=rdma in mooncake logs)
  • PD + NEXTN (steps=2, topk=1, draft=3, no rejection sampling): spec_accept_length mean 2.816 (min 2.225, max 3.000) over 100 GSM8K prompts at temp=0; GSM8K 97/100 (preliminary accuracy signal, not an accuracy validation)

Ask

This 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.

@YAMY1234

YAMY1234 commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator Author

@Leslie360

Thanks for the detailed verification! The Qwen3.8 Flash NEXT rebase PR is expected to merge into main within the next day or two, so it may be cleaner to wait for that and then rebase this PR onto main.

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?

@YAMY1234
YAMY1234 force-pushed the yamy/qwen38-next-pd-disagg-20260827 branch from 7bcf5ba to 39f01f5 Compare September 9, 2026 00:06
@YAMY1234
YAMY1234 changed the base branch from qwen4-main-squashed to main September 9, 2026 00:07

@YAMY1234 YAMY1234 left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Focused review of the QSA PD state-transfer paths.

Comment thread python/sglang/srt/disaggregation/prefill.py Outdated
Comment thread python/sglang/srt/disaggregation/decode.py Outdated
Comment thread python/sglang/srt/disaggregation/mori/conn.py
@YAMY1234 YAMY1234 changed the title [Qwen3.8] Add PD state transfer for Flash Next [Qwen3.8-Next] Add PD state transfer for Flash Next Sep 9, 2026
@YAMY1234
YAMY1234 marked this pull request as ready for review September 9, 2026 03:01
@Leslie360

Copy link
Copy Markdown

Reviewed at head ef784f466 (diff vs origin/main: +893/-101, 6 commits), with attention to the two
latest commits (ef784f466 matching-attention-TP transfer, 39f01f57 heterogeneous attention TP via
Mooncake). Line numbers below refer to the PR head.

Overall

Direction looks right and the implementation is careful. The pieces that usually go wrong in this area
— TP-sharded vs replicated state disambiguation, writer election for replicated state under
heterogeneous attention TP, layer-id pairing across PP stages, empty-PP-stage registration, and
GPUDirect write visibility before CUDA consumes transferred state — are all addressed with explicit
mechanisms rather than assumed-away. I checked the writer-election math against
CommonKVManager._resolve_rank_mapping (python/sglang/srt/disaggregation/common/conn.py:682-720)
and it lines up exactly for both scatter (prefill TP > decode TP) and aggregation directions. I did
not find a definite blocking bug; the one correctness item below needs the author's confirmation, and
the rest are validation-oriented questions from a TP8 dual-node Mooncake RDMA deployment.

Blocking / correctness concerns

1. (needs confirmation) MORI transfers QSA state but is excluded from the GPUDirect visibility flush

_requires_qsa_gpudirect_flush (python/sglang/srt/disaggregation/decode.py:172-185) only returns
True for TransferBackend.MOONCAKE and TransferBackend.NIXL, yet this PR adds
qsa_pending/qsa_compressed to the MORI send path
(python/sglang/srt/disaggregation/mori/conn.py:1340-1341, guard at :1472). MORI stands up an RDMA
backend (python/sglang/srt/disaggregation/mori/conn.py:342, :365-372, MORI_RDMA_DEVICES).

The flush exists because a host-side RDMA poll (KVPoll.Success) does not make third-party GPU writes
visible to CUDA before the decode forward reads the QSA pending ring / compressed pages. If MORI's
engine delivers writes via GPUDirect RDMA without its own cudaDeviceFlushGPUDirectRDMAWrites
(or equivalent), the decode side can read stale QSA state right after success — the exact hazard this
PR fixes for Mooncake/NIXL at decode.py:2457-2463. The docstring says the scoping is intentional,
but it scopes over state types, not over the backend's memory-ordering contract, and MORI is
receiving new QSA state transfers in the same PR.

Could you confirm one of:

  • MORI's RDMA engine already performs a CUDA-owner visibility operation on the destination before
    completion is signalled (in which case a one-line comment next to _requires_qsa_gpudirect_flush
    would prevent someone "fixing" this later), or
  • QSA state over MORI should get the same flush (one-line addition to the tuple)?

What I verified (so the questions above stay focused)

  • should_send_replicated_state (python/sglang/srt/disaggregation/utils.py:790) elects writers
    r % (src_tp // dst_tp) == 0 for src > dst, which matches the decode-side pull mapping
    target_tp_rank = (engine_rank % attn_tp_size) // (src // dst) in common/conn.py:692-697; for
    src < dst every source sends and each decode rank's group is covered. The same election feeds the
    replicated-Mamba zero-dim path in compute_mamba_state_slice_byte_blocks (utils.py:918-934).
  • Pending-ring indexing is consistent: ring rows are req_pool_idx * qsa_compress_ratio + pos % ratio (qsa_kv_pool.py:125-134), the transfer item is tensor[0].nbytes * ratio
    (qsa_kv_pool.py:227-246), and both sides use their own side-local req_pool_idx
    (prefill.py:1323-1326, decode.py:1493-1496) — no cross-side index equality assumption.
  • QSA_COMPRESSED reuses full-KV page indices; compressed capacity/items line up 1:1 with full pages
    given page_size % qsa_compress_ratio == 0 (qsa_kv_pool.py:120-124, :261-270).
  • Sentinels QSA_ROPE_STATE_LAYER_ID = 2^32-2 / PLE_NGRAM_STATE_LAYER_ID = 2^32-1 fit the uint32
    wire format (pack_int_lists(..., "I") on both pack and unpack sides,
    mooncake/conn.py:2642-2644, :175/:178/:186, nixl/conn.py:3169-3171, :267+).
  • Empty PP stage is a clean no-op on both transports: Mooncake via
    build_transfer_entry_pairs returning [] for n_src == 0 (utils.py:971-972) plus
    _transfer_data short-circuit (mooncake/conn.py:642-644); Nixl via the new
    if not layers_params: return None (nixl/conn.py:1614-1615) with if h is not None at the call
    site (nixl/conn.py:2553-2554).
  • The Mooncake heterogeneous-TP branch validates ptr/item-lens consistency and per-pair item-lens
    equality before sending (mooncake/conn.py:1511-1538).

Questions for the author

  1. A800/SM80 flush support. Our fleet is SM80 (A800), TP8 across two nodes, Mooncake RDMA. Does
    cudaDeviceFlushGPUDirectRDMAWrites(CurrentDevice, ToOwner) succeed on A800? If it returns
    cudaErrorNotSupported, the fallback is a full torch.cuda.synchronize() per pop_transferred
    batch with successes (decode.py:150-166), which is a device-wide stall in the scheduler loop —
    fine for correctness, but have you measured the decode-step overhead in that case?
  2. Heterogeneous attention TP test matrix. Which direction was end-to-end validated —
    prefill TP8 → decode TP4, the reverse, or both — and was any run done with PP > 1 where a middle
    PP stage owns no full-attention layer (the empty-registration path,
    qsa_kv_pool.py:234-241)? That combination exercises the positional-only fallback being disabled
    (allow_positional_fallback=self.pp_size == 1).
  3. Page-size symmetry across PD. QSA_COMPRESSED item lengths derive from
    page_size // qsa_compress_ratio (qsa_kv_pool.py:123), so prefill/decode must run the same
    --page-size for the compressed-page indices to be meaningful. Is that enforced anywhere for
    Qwen4-Exp PD, or assumed? (_requires_exact_state_index_match catches count mismatches but not
    equal-count/stride mismatches.)
  4. Prefix-cache hit at prefill. With MambaRadixCache, when a prefill request hits the radix
    cache, is the QSA pending ring still (re)populated for the request tail before
    _qsa_pending_payload reads it (prefill.py:1323-1326)? A fully-cached prompt that skips
    full-attention KV compute would otherwise ship a stale/zeroed ring.
  5. CUDA-graph capture window. Confirm DecodeTransferQueue.pop_transferred (and thus the flush
    at decode.py:2457-2463) only ever runs on the scheduler thread outside graph capture, including
    with --enable-overlap-schedule. We have been burned by capture-period scheduler-state access on
    a related path, so want this stated explicitly.
  6. MORI — see the blocking item above.

Nits

  • decode.py:139-166: _flush_gpudirect_writes_to_cuda_owner catches only
    (ImportError, AttributeError) before falling back to torch.cuda.synchronize(). A driver-level
    failure raised by cuda.bindings as a different exception type would escape into the scheduler
    loop instead of degrading to the sync fallback. Consider widening the catch to Exception.
  • nixl/conn.py:1576-1578: the layer-id-paired path uses item_lens[i] (source) as the destination
    stride with no src/dst item-len equality check, unlike the Mooncake guard at
    mooncake/conn.py:1528-1538. Safe today (Nixl raises on heterogeneous TP for non-MLA at
    nixl/conn.py:2492-2495, so lens are identical), but worth a comment or assert so a future
    heterogeneous-Nixl change doesn't silently inherit the assumption.

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.
@YAMY1234

YAMY1234 commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator Author

@Leslie360 Thanks for the review. The head has moved to abb5dfe856; the main change since the version you reviewed is that the QSA-specific GPUDirect flush has been removed entirely. Replies follow your numbering.

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 main (KV pages, Mamba, DSA, SWA) performs this flush on any of the three backends, and this PR never observed a stale read; the flush was added preventively. It is now removed, and QSA follows the same completion contract as the established state types: poll returns Success, the request is released. The TP4 -> TP4 production parity matrix was rerun without the flush: 12/12 cases, 84 requests, 2730 generated tokens, zero output-token mismatches; the previous flush head vs. the no-flush PD head differs by at most 0.00047 in selected-token logprob. If transport-level write visibility should be enforced, that belongs in a separate change covering every transferred buffer and every backend.

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 support_pp=False guard, so the empty-QSA-stage registration path is covered by focused unit tests rather than an E2E run. That is stated in the description as well.

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 page_size mismatch raises in common/conn.py. This is not Qwen4-specific logic.

Q4 (prefix-cache hit and the pending ring). The compressed QSA pool requires page_size % compress_ratio == 0 at construction, and radix prefix boundaries are page-aligned, so an extend always starts at a compression-group boundary and the current prefill always writes the trailing group's ring before _qsa_pending_payload reads it. This is the same behavior as aggregate mode, not a PD-specific assumption; the 8K-input parity cases exercise this path.

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 abb5dfe856: maybe_send_extra now passes the peer item lengths into _send_kvcache_generic, and the layer-id paired path raises on any src/dst mismatch, with a unit test in test_nixl_backend_basic.py.

@YAMY1234

YAMY1234 commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator Author

/rerun-group disaggregation

@github-actions

github-actions Bot commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-group disaggregation:

🚀 4-gpu-gb300 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_aarch64.py

🚀 2-gpu-h100 (6 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_basic.py
cd test/ && python3 registered/disaggregation/test_disaggregation_chunked_prefill_abort.py
cd test/ && python3 registered/disaggregation/test_disaggregation_decode_offload.py
cd test/ && python3 registered/disaggregation/test_disaggregation_optimistic_prefill.py
cd test/ && python3 registered/disaggregation/test_disaggregation_rust_server.py
cd test/ && python3 registered/disaggregation/test_disaggregation_unified_memory.py

🚀 8-gpu-h20 (5 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_decode_radix_cache.py
cd test/ && python3 registered/disaggregation/test_disaggregation_different_tp.py
cd test/ && python3 registered/disaggregation/test_disaggregation_dp_attention.py
cd test/ && python3 registered/disaggregation/test_disaggregation_nixl.py
cd test/ && python3 registered/disaggregation/test_disaggregation_pp.py

🚀 8-gpu-h200 (4 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_decode_radix_cache_swa.py
cd test/ && python3 registered/disaggregation/test_disaggregation_dsv4.py
cd test/ && python3 registered/disaggregation/test_disaggregation_hisparse.py
cd test/ && python3 registered/disaggregation/test_disaggregation_hybrid_attention.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_dwdp_gpt_oss.py
cd test/ && python3 registered/disaggregation/test_disaggregation_inkling_mxfp8.py

🚀 4-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_kimi_linear.py
cd test/ && python3 registered/disaggregation/test_epd_disaggregation.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_xpu.py

🚀 8-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_kimi_linear_pd_dcp4.py

Comment thread python/sglang/srt/disaggregation/decode.py Outdated

@ShangmingCai ShangmingCai left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, logic is clean and correct

@YAMY1234

Copy link
Copy Markdown
Collaborator Author

/rerun-test registered/kernel/qsa/test_qsa.py

@github-actions

github-actions Bot commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/kernel/qsa/test_qsa.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernel/qsa/test_qsa.py

# Conflicts:
#	python/sglang/srt/disaggregation/mooncake/conn.py
#	python/sglang/srt/disaggregation/nixl/conn.py
@YAMY1234

Copy link
Copy Markdown
Collaborator Author

/rerun-test registered/mock_model/test_e2e_pd.py

@github-actions

github-actions Bot commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/mock_model/test_e2e_pd.py:

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/mock_model/test_e2e_pd.py

@YAMY1234

Copy link
Copy Markdown
Collaborator Author

/rerun-failed-ci bypass-fastfail 9851fb3

@YAMY1234

Copy link
Copy Markdown
Collaborator Author

All NV CI Passed

@Leslie360

Copy link
Copy Markdown

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:

  • Decode node boots clean: CUDA-graph warmup passes, /health 200, no asserts.
  • PD e2e through the router: 6/6 concurrent short requests, 3/3 long-context (~2.4 k-token) requests, zero asserts on either side; mamba conv/temporal state transfers are silent in the logs (a good sign).
  • Determinism baseline (temp0 + seed42, 8 fixed prompts, no-thinking): standalone self 5/8 identical; PD self 6/8; PD↔standalone 3/8. The engine itself is numerically nondeterministic at temp0 on this hardware, so PD-vs-standalone token differences sit within baseline noise; PD self-consistency meets the single-node baseline.
  • Sibling hybrid-GDN regression: Qwen3.5-27B PD 6/6 short requests.

Three environment gotchas that cost us most of the bring-up time (none of them sglang bugs):

  1. Mooncake wheel source — PyPI frequently times out from CN networks; https://pypi.tuna.tsinghua.edu.cn/simple/mooncake-transfer-engine/ works. For CP312, mooncake_transfer_engine-0.3.9-cp312-cp312-manylinux_2_35_x86_64.whl installs fine as a file. Verify with from mooncake.engine import TransferEngine.
  2. libcudart.so.12 soname — the Mooncake .so links libcudart.so.12; a CUDA-13 torch stack only ships .so.13. Extract from pip download nvidia-cuda-runtime-cu12==12.9.79 --no-deps and prepend to LD_LIBRARY_PATH.
  3. Router flag--pd-disaggregation is store-true; passing a value (--pd-disaggregation true) errors. Bring-up order: prefill → decode (each /health 200) → router → first completion. Watch for Topology discovery complete. Found N HCAs and pin --disaggregation-ib-device mlx5_... to what ibv_devices reports.

Happy to PR these notes into the PD cookbook/docs if wanted.

@Leslie360

Copy link
Copy Markdown

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:
The decode-side PLE n-gram history is transferred from the prefill side (true history), not EOS-initialized. _iter_transfer_state_entries (memory_pool.py:1117) iterates the mamba conv/temporal tensors and for sibling in self._slot_siblings: yield from sibling.iter_transfer_state_entries() (L1129-1130); ShortConvPool.iter_transfer_state_entries (ple_state_pool.py:129) yields ("ple_short_conv", conv_state, ...) and NGramPool.iter_transfer_state_entries (L239) yields ("ple_ngram", self.context, ...). Those pointers flow into kv_args.state_data_ptrs (disaggregation/utils.py: setup_state_kv_args) and are RDMA'd by _send_mamba_state (mooncake/conn.py:1700). So the decode side starts from the correct prefill tail — good.

2. Hardware re-verify (new main @ 03ea13a):

Check Result
Qwen4-Exp PD e2e (prefill+decode+router, TP8 each) ✅ 6/6 concurrent short; 3/3 long-context (~2.4 k-token); zero asserts
temp0 determinism (8 prompts, no-thinking) PD self 6/8, standalone self 4/8 → PD meets/beats single-node baseline (engine is numerically nondeterministic at temp0; no evidence of transfer-induced drift)
MTP regression (NEXTN3) ✅ accept 3.00 / 2.83 vs our b852d-based baseline 3.10 — within noise
Qwen3.5-27B PD regression (sibling hybrid GDN) ✅ 6/6
Qwen4-Exp DFLASH draft (decode-graph capture) new regression on merged main — see below

3. New regression we hit (needs a maintainer look): --speculative-algorithm DFLASH with the Qwen4-Exp draft on new main crashes during decode CUDA-graph capture:
ValueError: too many values to unpack (expected 2) at qwen3_vl.py:1642 (hidden_states, aux_hidden_states = hidden_states when capture_aux_hidden_states), reached from qwen4_exp.py:1784 (super().forward) during decode_cuda_graph_runner capture. Root cause is upstream main missing the Qwen4-Exp capture guard from #38642 (the and not self.capture_aux_hidden_states condition on the HC-stream overwrite): b852d33 ran clean because it is the #38642 branch head carrying that guard. We rebased #38642 onto 03ea13a545 and the crash disappears (verified on 2×8 A800: server up, 3/3 OK, decode graph captured, accept ≈1.07). Details in #39685.

Everything else from our earlier report holds on merged main. Happy to share logs.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants