Repository navigation
[Bugfix][KV Connector][NIXL] Transfer the whole PLE short-conv page in disaggregated P/D - #59997
Conversation
There was a problem hiding this comment.
this should simply be nixl_agent_meta.ple_block_len so if the two workers are out of sync, we don't introduce a bug which can reset the behaviour that the PR is trying to solve.
There was a problem hiding this comment.
Done in 63a74270ec: the remote side now reads nixl_agent_meta.ple_block_len directly and raises if the remote did not advertise it (different vLLM version), so an out-of-sync pair can never fall back to the truncated region length.
96871a7 to
63a7427
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@claude review |
hmellor
added
the
verified
63a7427 to
eb13815
Compare
|
@claude review |
|
/ci run |
|
✅ Triggered Buildkite CI #93103 for commit |
github-actions
Bot
added
the
ready
CI selector (shadow): 14 test steps (18 jobs) instead of 19 (25 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (14)
Would skip (today's rules run them) (6)
Would add (today's rules do not run them) (1)
AMD mirrors: would skip (16)
AMD mirrors: would add (1)
3 changed files · base |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #93103. |
|
/ci cancel |
|
✅ No cancelable CI build is running for branch |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #93103. |
|
CI failures seem due to infra instability. All the others are green. |
Head branch was pushed to by a user without write access
2528cf2 to
97db14b
Compare
|
@arpera it should be fine now |
|
@NickLucche, please, have a look this PR if you have some bandwidth |
|
/ci run |
|
❌ This PR is 99 commits behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #93641 for commit |
The NIXL worker folds every layer that reports the same base address into one region. In the block-outer HMA layout every cache group places its first page at byte 0 of the shared allocation, so on Qwen3.8-Flash-Next the PLE short-conv state (tp-replicated MambaSpec, page = channels x (conv_state_len + num_spec_tokens)) is discovered in the region of layer-0 GDN, the QSA ring and the MLA indexer cache, whose block_len is the MLA page. The PLE descriptors were built from that block_len and transferred only 204800 of the 245760 bytes of each PLE slot; the untransferred channels kept whatever the block held before (blocks written by an async load are exempt from zeroing). With num_speculative_tokens >= 2 the spec-decode short conv read those stale columns, the layer-1 hidden state went NaN, and the rejection sampler rejected every draft: 0% MTP acceptance in disaggregated P/D serving. Record the PLE page length at registration and use it for the PLE descriptors on both sides of the transfer, exchanging it through the handshake metadata so the P/D geometry check still compares PLE pages. Fixes vllm-project#59642 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Alberto Massidda <amassidda@nvidia.com>
…ck_len The PLE page length is always recorded when the PLE region is registered, and a remote that does not advertise it runs a different vLLM version: raise instead of silently transferring the region's (smaller) block_len. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Alberto Massidda <amassidda@nvidia.com>
…xture Advertising ple_block_len changes NixlAgentMetadata, so promote the connector version and assert the field on the remote side. The regression test now widens the PLE page of _make_csa_linear_ple_worker instead of building a second fixture. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Alberto Massidda <amassidda@nvidia.com>
Signed-off-by: Alberto Massidda <amassidda@nvidia.com>
891705e to
f0bc717
Compare
|
/ci run |
|
❌ This PR is 2 commits behind upstream |
|
/ci run |
Artem agreed that his comments fixed
Reconcile with vllm-project#59997 (PLE short-conv page transfer): a PLE page wider than the main KV page it aliases now registers as its own region under the (base_addr, block_len) key, so _ple_region_index points at a dedicated region whose block_len is the full PLE page. Update test_csa_linear_ple_page_wider_than_the_shared_region_page to the split expectations; its descriptor assertions are unchanged since the PLE descs already cover the whole page. Co-authored-by: ZCode <noreply@zcode.ai>
vllm-project#59625 landed on main squashed and reworked, so its files take main's version and only the NIXL commit is replayed on top. Adaptations: - The connector hook is unregister_kv_caches on main; NixlPullConnector now implements that name (the old release_kv_caches was never called). - main took NIXL_CONNECTOR_VERSION 14 (vllm-project#59997); the registration epoch is version 15. - A request whose peer re-registered handshakes directly from _start_read; other requests go through _read_blocks_for_req (vllm-project#55471). - test_config follows main's refusal-only parametrization. Signed-off-by: Dakai An <dakaian108@gmail.com>
…n disaggregated P/D (vllm-project#59997) Signed-off-by: Alberto Massidda <amassidda@nvidia.com> Co-authored-by: Alberto Massidda <amassidda@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Artem Perevedentsev <aperevedents@nvidia.com>
Purpose
Fixes #59642 (Qwen3.8-Flash-Next: 0% MTP acceptance in disaggregated P/D serving with
NixlConnector).The NIXL worker folds every KV cache layer that reports the same base address into one region (
register_kv_caches,if base_addr in seen_base_addresses). In the block-outer HMA layout every cache group places its first page at byte 0 of the shared allocation, so on Qwen3.8-Flash-Next the PLE short-conv state (Qwen4ExpPLELayer, a tp-replicatedMambaSpec, page =channels × (conv_state_len + num_spec_tokens)= 245760 B) lands in the same region as layer-0 GDN, the layer-3 QSA ring and the layer-3 MLA indexer cache, and that region keeps the MLA page length (204800 B). The PLE descriptors built fromblock_len_per_layer[_ple_region_index]therefore transfer only 204800/245760 bytes of each PLE slot: 8534 of 10240 channels arrive on the decode worker, the rest is whatever the block held before (never zeroed: the block is exempt from zeroing because an async load is expected to overwrite it). Withnum_speculative_tokens >= 2the spec-decode short conv reads those stale columns, the layer-1 hidden state becomes NaN, the target logits are NaN and the rejection sampler rejects every draft (the generated tokens are garbage too, not just slow).This PR records the PLE page length at registration (
_ple_block_len) and uses it for the PLE descriptors on both sides of the transfer (_build_mamba_local/_build_mamba_remote), exchanging it throughNixlAgentMetadata.ple_block_len(optional, backward compatible) so the existing P/D geometry check still compares PLE pages. Region discovery is unchanged: the PLE keeps sharing its region with the pages it aliases, astest_csa_linear_registration_discovers_shared_regionsexpects; only the transfer length of its pages changes.Test Plan
tests/v1/kv_connector/unit/test_nixl_ple_region.py: a CSA-linear-like config (MLA + GDN + ring + PLE sharing byte 0 of one block-outer tensor, PLE page larger than the MLA page); asserts the PLE descriptors cover the whole PLE page and that the page length is advertised in the handshake metadata. Before the fix the PLE descriptors carry the MLA page length.tests/v1/kv_connector/unit/test_nixl_connector_hma.py+test_nixl_desc_geometry.pyrun inside the serving image (vLLM109b5363fc) with the patch hot-patched: 307 passed before and after (one HF-download test deselected).109b5363fc+ this patch hot-patched into the image, Dynamo +NixlConnector, 1P/1D,nvidia/Qwen3.8-Flash-Next-NVFP4, MTP k=3, 200 requests of 64/64 tokens): mean acceptance length goes from 1.00 on 20/21 windows (zero accepted drafts) to 2.7–3.5 on every window, no NaN; the same configuration in aggregate mode gives ~2.1–2.4.Test Result
See above; details and the full investigation (per-request probes of the PLE conv state on prefill and decode, per-column comparison of the transferred slot) are in #59642.
🤖 Generated with Claude Code