Repository navigation
[PD] Support heterogeneous TP for hybrid SWA KV transfer - #40909
Merged
Merged
Conversation
Hybrid sliding-window models rejected PD transfer whenever prefill and decode use different attention TP sizes, although ordinary MHA KV transfer already supports that topology. Send the SWA state by token and KV-head slice, reusing the existing head ownership calculation, on both NIXL and Mooncake. Each backend keeps its existing completion handling, and the separate sliding-window cache allocation is retained. SWA state carries no layer ids, so pipeline parallelism and layouts without separate per-layer K/V buffers remain rejected. Other hybrid-state guards are unchanged. Also let NixlKVManager.set_kv_buffer_tensors accept the slot_layer_ids keyword that shared staging initialization already passes, fixing a startup TypeError when NIXL staging is enabled.
Register a GPU regression for PD transfer between unequal attention TP sizes on a hybrid sliding-window model: gpt-oss-20b with prefill TP2 and decode TP1 over NIXL, reusing the PD fixture and the GSM8K eval and floor of the existing equal-TP sliding-window test. Two prefill ranks write disjoint KV heads into the same decode window pages, which the CPU tests cannot exercise end to end. Document the unequal-TP support and its limits for hybrid sliding-window models in the PD disaggregation guide.
github-actions
Bot
added
the
documentation
Jiminator
marked this pull request as ready for review
September 24, 2026 01:23
Jiminator
requested review from
ByronHsu,
Duyi-Wang,
HaiShaw,
JustinTong0323,
ShangmingCai,
hnyls2002,
sogalin,
wisclmy0611 and
zijiexia
as code owners
September 24, 2026 01:23
YAMY1234
approved these changes
Oct 1, 2026
| f"do not split into {src_heads} and {dst_heads} heads of " | ||
| f"{total_kv_heads} total with page_size={page_size}" | ||
| ) | ||
| # Each page item is [page_size, heads, head_dim]: one block per token. |
Collaborator
There was a problem hiding this comment.
This address calculation assumes that each page is laid out as NHD [page_size, heads, head_dim]. A ROCm vectorized_5d cache has the same total item size but a different physical row order, so the length checks above can succeed while these offsets point to the wrong data. Could we explicitly reject that layout here (and ideally cover it with a negative test) to avoid silent corruption?
|
|
||
| > **Note:** The staging buffer is designed for non-MLA models (e.g. GQA, MHA). MLA models (e.g. DeepSeek-V2/V3) should not enable this flag. | ||
|
|
||
| > **Note:** You can use different attention TP sizes for hybrid sliding-window models (e.g. `openai/gpt-oss-120b`) with Mooncake or NIXL. For unequal TP, both backends transfer the sliding-window KV by token and KV-head slice. The staging buffer carries only full-attention KV. Use `--pp-size 1` and leave `--enable-unified-memory` disabled for this configuration. |
Collaborator
There was a problem hiding this comment.
Suggestion: this note should call out that the unequal-TP SWA path still rejects the MXFP8 BLOCK_SCALE / BLOCK_SCALE_SWA states. Stating the limitation here keeps the documented support matrix aligned with the runtime behavior.
Suggested change
| > **Note:** You can use different attention TP sizes for hybrid sliding-window models (e.g. `openai/gpt-oss-120b`) with Mooncake or NIXL. For unequal TP, both backends transfer the sliding-window KV by token and KV-head slice. The staging buffer carries only full-attention KV. Use `--pp-size 1` and leave `--enable-unified-memory` disabled for this configuration. | |
| > **Note:** You can use different attention TP sizes for hybrid sliding-window models (e.g. `openai/gpt-oss-120b`) with Mooncake or NIXL. For unequal TP, both backends transfer the sliding-window KV by token and KV-head slice. The staging buffer carries only full-attention KV. Use `--pp-size 1` and leave `--enable-unified-memory` disabled for this configuration. MXFP8 KV cache is not supported for unequal-TP SWA transfer. |
The unequal-TP SWA head slice computes per-token addresses that assume NHD pages. Record the SWA sub-pool's resolved KV cache layout and reject any other layout (e.g. ROCm vectorized_5d) before building addresses, with a negative case in the existing transfer test. Document that MXFP8 KV cache is not supported for unequal-TP SWA transfer.
Collaborator
Author
|
/tag-and-rerun-ci |
Jiminator
added
the
bypass-fail-fast
This was referenced Oct 6, 2026
Closed
3 of 5 tasks
15 of 40 tasks
Retain both the SWA head-slice transfer helper and upstream deferred-release failure handling.
…ests and comments
hnyls2002
changed the title
Collaborator
|
/rerun-test test_disaggregation_swa_different_tp.py test_disaggregation_decode_radix_cache_swa.py test_disaggregation_different_tp.py |
Contributor
|
🚀 🚀 🚀 |
Contributor
|
🚀 🚀 🚀 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Hybrid sliding-window models currently reject PD KV transfer when prefill and decode use different attention TP sizes. For example,
openai/gpt-oss-120bwith prefill TP1 and decode TP4 reaches the non-MLA SWA guard even though ordinary MHA KV transfer supports that topology.Modifications
Copy SWA state by token and KV-head slice using the existing head-ownership calculation. Both NIXL and Mooncake retain their existing completion handling and the separate sliding-window cache allocation. Fix the pre-existing NIXL staging startup error by accepting the
slot_layer_idskeyword already supplied by shared initialization.Related work: #40861 includes the same NIXL staging setter compatibility fix. It is retained here because staging otherwise fails during initialization on this PR's base; this PR can drop that hunk if #40861 lands first.
Support is limited to PP1 with separate per-layer K/V buffers. Heterogeneous SWA with pipeline parallelism or an incompatible buffer layout remains rejected. MXFP8 KV cache with unequal attention TP on a non-MLA model is now rejected at the first prefill/decode handshake; it previously failed on every request after transfer had started, on main as well. Other hybrid-state guards are unchanged. The deployment guide documents these limits.
Accuracy Tests
>= 0.45floor. A bounded request using the same fixture hits the original SWA rejection on the unfixed base and succeeds on the candidate.9a563bb(after merging main), the wire, NIXL, PP/hybrid, Mooncake batching and staging-slot unit suites pass (165 tests). The handshake check rejects MXFP8 with unequal TP and accepts equal-TP MXFP8, unequal-TP bf16 SWA and MLA.--disable-hybrid-swa-memorypreserves the same prefill/decode TP split and transfers all KV through the existing path. All five cold comparisons in each direction were bit-identical. Overall exact-ID/0.05-logprob parity was 17/21 for TP1→TP4 and 12/21 for TP4→TP1. Every failed warm/concurrent comparison already had different prefill logits before transfer; full-cache and hybrid-cache runs differed in prefix reuse or batching. Those failures remain recorded, with no tolerance change. Pristine colocated TP1/TP4 controls also differ. These GPT-OSS-120B checks do not establish general cross-TP output equivalence or trained quality.Public GPT-OSS-120B test setup: revision
b5c939de8f754692c1647ca79fbf85e8c1e70f8a, based4dcce12d4fcdb3eb694d54d6afd31628c2eb7ce, two nodes with four NVIDIA GB300 GPUs each, Torch 2.13.0+cu130, FlashInfer 0.6.18, sgl-kernel 0.4.7, NIXL 1.4.1 and Mooncake 0.3.13. TCP transport was used. TP1/TP4 deliberately differs from the cookbook's TP8 benchmark to exercise unequal TP; this is not a reproduction of that benchmark.Key flags:
--attention-backend trtllm_mha --page-size 64 --chunked-prefill-size 1024 --context-length 8192 --max-total-tokens 16384 --max-running-requests 16 --cuda-graph-max-bs-decode 16 --mem-fraction-static 0.65 --random-seed 42, plus the respective disaggregation role, backend and TP. Staged runs setSGLANG_DISAGG_STAGING_BUFFER=1; the decode-cache run adds--disaggregation-decode-enable-radix-cache. Hybrid SWA memory optimization stays enabled.The registered GPT-OSS-20B test uses revision
6cee5e81ee83917806bbde320786a8fb61efebee, three GB300 GPUs, Triton attention and page size 64, with hybrid SWA memory enabled. UCX had TCP, CUDA copy/IPC and shared-memory transports enabled. The CI registration targets four H100 GPUs; that runner has not been tested here.Speed Tests and Profiling
Three warm 2049-input/64-output repetitions of the equal-TP control had median request latency 220.9 ms before and 220.2 ms after. This is a small regression sanity check, not a throughput or speedup claim. No RDMA performance claim.
The existing Mooncake non-staged TCP path hit a queue-full error on a large cached prefix; the same failure reproduced on unmodified upstream using its ordinary full-KV path. The Mooncake functional results above are scoped to staging.
Checklist
pre-commit run --all-filespassed with no formatter edits.CI States
Latest PR Test (Base): ⏳ Run #38344788909
Latest PR Test (Extra): ❌ Run #38344788952
Latest PR Test (AMD ROCm 10): ⏳ Run #38344789254