[Jiminator]

Motivation

Hybrid sliding-window models currently reject PD KV transfer when prefill and decode use different attention TP sizes. For example, openai/gpt-oss-120b with 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_ids keyword 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

• A registered GPT-OSS-20B NIXL TP2→TP1 test reuses the PD fixture and GSM8K mixin: 500 examples scored 0.524 with the shared mixin’s >= 0.45 floor. A bounded request using the same fixture hits the original SWA rejection on the unfixed base and succeeds on the candidate.
• 144 targeted unit tests passed across the wire, NIXL, PP/hybrid, Mooncake batching and staging-slot suites; changed production-line coverage is 93%. The SWA transfer, PP-rejection and staging-setter regressions fail on their corresponding unfixed controls.
• At 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.
• Actual NIXL and Mooncake transfers of synthetic CUDA buffers on one GPU passed exact-byte checks for five logical TP/KV-head topologies, including both directions and replicated KV heads. Destination window bytes match and surrounding sentinel bytes remain untouched.
• GPT-OSS-120B: TP1→TP4 and TP4→TP1 passed on both backends, with 21 requests per configuration spanning cold/warm prefixes, page/window boundaries, chunked prefill and concurrent requests. Mooncake used GPU staging. NIXL TP1→TP4 additionally passed with staging and decode radix caching enabled. Baseline TP1→TP4 hit the original SWA rejection.
• Equal-TP NIXL TP4→TP4 produced identical tokens and logprobs before/after on all 21 requests.
• A separate unmodified-base control with --disable-hybrid-swa-memory preserves 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, base d4dcce12d4fcdb3eb694d54d6afd31628c2eb7ce, 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 set SGLANG_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

• Full pre-commit run --all-files passed with no formatter edits.
• Regressions reuse registered unit and GPU fixtures.
• Supported layouts, validation flags and limitations documented above; no new user option.
• Numerical and latency evidence provided with its limits.
• Existing disaggregation patterns and code style followed.

---

CI States

Latest PR Test (Base): ⏳ Run #38344788909
Latest PR Test (Extra): ❌ Run #38344788952
Latest PR Test (AMD ROCm 10): ⏳ Run #38344789254

[YAMY1234]

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?

[YAMY1234]

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.

[Jiminator]

/tag-and-rerun-ci

[hnyls2002]

/rerun-test test_disaggregation_swa_different_tp.py test_disaggregation_decode_radix_cache_swa.py test_disaggregation_different_tp.py

[github-actions]

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

cd test/ && python3 registered/e2e/disaggregation/test_disaggregation_swa_different_tp.py

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

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

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

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

[github-actions]

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

cd test/ && python3 registered/e2e/disaggregation/test_disaggregation_swa_different_tp.py

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

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

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

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