[PD] Share one head-slice helper across mooncake, mori, and nixl

[hnyls2002]

• RFC: PD disaggregation: single protocol layer, per-backend transport #33861
• Staged implementation plan and PR tracking: PD shared-protocol unification #34510
• Preceding PR: Share the prefill->decode failure notification across backends #36612

Background

Step 1 of #34510: the same common semantic is implemented more completely in one backend than another.

Heterogeneous-TP head-index arithmetic exists in four places -- compute_head_slice_params in common/staging_buffer.py (already on main, already used by the staging path), plus hand-rolled copies in mooncake, nixl and mori. mooncake and nixl divide by the GQA replication factor; mori does not, and so maps decode ranks onto the wrong KV heads whenever decode_tp_size > total_kv_heads.

This PR routes the three copies through the shared one. compute_head_slice_params itself is unchanged; only the call sites change. No Transport abstraction, no wire-format change.

Modifications

mooncake -- send_kvcache_slice drops 41 lines of inline arithmetic, including its if attn_tp_size > dst_attn_tp_size direction branch, and calls the shared helper.

nixl -- the gather branch does the same. Its scatter branch already called the helper.

mori -- same call site, and this one changes behavior. Its scatter branch read:

src_head_start = (dst_tp_rank * dst_heads_per_rank) % src_heads_per_rank

Under GQA replication (decode_tp_size > total_kv_heads) consecutive decode ranks share a KV head -- QKVParallelLinear maps by tp_rank // num_kv_head_replicas -- so a plain modulo hands ranks 1..r-1 of each group a head they do not own. With prefill TP=1, decode TP=4 and 2 KV heads the correct map is [0, 0, 1, 1]; mori produced [0, 1, 0, 1]. Wrong head indices do not fail loudly: the transfer delivers the wrong channels and the only symptom is a garbled accuracy score.

Equivalence. The old inline code was extracted verbatim from main and compared against the shared helper over a sweep of src_tp/dst_tp in {1,2,4,8,16,32}, total_kv_heads in {1..8,16,32,64,128}, ranks up to 2x tp to cover modulo wraparound:

call site	combinations	mismatches
mooncake send_kvcache_slice, both directions	46128	0
nixl gather branch, its only reachable domain	62496	0
mori gather branch	14880	0
mori scatter branch	31248	932 -- the fix above

Unit test

test/registered/unit/disaggregation/test_head_slice_params.py, registered CPU-only (est_time=10), so it costs no GPU runner time. Covers gather (4->2), scatter (2->4), GQA replication (1->4) and equal TP. Expected values are derived by hand from the head-distribution rules, never by calling the implementation, so a bug there cannot make both sides agree.

On the Step 1 "backend-parameterized regression test" convention: after this PR there is one implementation rather than three, so the test targets the shared helper directly instead of parameterizing over backends.

Test suite

The silent-failure mode above is why the 8-gpu-h20 suite carried a {plain, staging} x {gather, scatter} matrix. With the indices guarded on CPU, mooncake's send_kvcache_slice has no direction branch left -- gather and scatter differ only in the values the helper returns -- so MooncakeMHADecodeLargerTP and StagingPrefillLargerTP are dropped from test_disaggregation_different_tp.py (8 classes -> 6).

Every transfer path keeps an end-to-end test:

path	remaining coverage
send_kvcache_slice gather	MooncakeMHAPrefillLargerTP (MHA 4->2)
send_kvcache_slice scatter	GDNHybridHeteroTP (GDN 1->4)
_do_staging_transfer, both directions	StagingDecodeLargerTP, StagingRadixPrefillLargerTP
send_kvcache whole-block	MooncakePrefillLargerTP, MooncakeDecodeLargerTP (MLA)

CI time

test_disaggregation_different_tp.py is the longest file in the 8-gpu-h20 suite -- 934s of
a 1728s suite on the 2026-09-25 scheduled run -- so it sets that stage's wall-clock floor.

	samples	mean
main, 8 classes	934 / 995 / 1173 / 1179s	1070s
this PR, 6 classes	650 / 779s	715s

About -355s (-33%) on the means. Run-to-run spread on this runner is wide (129s between
two runs of the same code here, 245s across the four main samples), so pairing the worst
case for this PR against the best for main still gives -155s (-17%); that is the floor,
not the expectation.

main samples are the h20 partitions of scheduled runs 36127993962 / 36071364371 /
35991484462 / 35932079543; this PR's are /rerun-test on the same runner class. The
registered est_time=1171 is left alone -- scripts/ci/update_est_time.py refreshes it
from stats p90 once this lands.

The new unit test is CPU-only (est_time=10, measured 6s) and adds no GPU runner time.

Not included

Mamba/linear-attention state transfer slices on state dimensions rather than KV heads, and has its own duplication (utils.py:862 and mori/conn.py) with the same plain-modulo shape. Different axis, left alone here.

Accuracy Tests

On the CUDA side the h20 disaggregation suite covers the mooncake and nixl call sites.

The mori change is on the AMD path. TestMoriTransferEngineTPMismatchE2E exercises prefill TP=2 / decode TP=4, but on Llama-3.2-1B-Instruct (8 KV heads) dst_replication = max(1, 4 // 8) = 1, so the old and new maps agree and that test is green either way -- it guards against regression on the non-replicating path, not the fix. Reaching the divergent path needs decode_tp > total_kv_heads; no shape in the AMD matrix has it today. The CPU unit test is therefore the only automated guard on the corrected mapping.

Speed Tests and Profiling

Not applicable -- the arithmetic is per-transfer setup, not a data path.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): 🚫 Run #36199005947
Latest PR Test (Extra): ❌ Run #36199005896
Latest PR Test (AMD ROCm 10): 🚫 Run #36199006038

[hnyls2002]

/tag-and-rerun-ci

[hnyls2002]

/rerun-test test_head_slice_params.py test_disaggregation_different_tp.py test_disaggregation_nixl.py test_mori_transfer_engine_e2e.py test_nixl_backend_basic.py

[github-actions]

Results for /rerun-test test_head_slice_params.py test_disaggregation_different_tp.py test_disaggregation_nixl.py test_mori_transfer_engine_e2e.py test_nixl_backend_basic.py:

🚀 ubuntu-latest (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/disaggregation/test_head_slice_params.py
cd test/ && python3 registered/unit/disaggregation/test_nixl_backend_basic.py

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

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

⛔ test_mori_transfer_engine_e2e.py: test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py is registered for AMD (suite stage-b-test-large-8-gpu-mi35x-disaggregation-amd), not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

[BBuf]

Reviewed the refactor itself carefully and it is faithful: I extracted the three old inline blocks and compute_head_slice_params and they agree on both directions for mooncake and on the reachable domain for nixl's gather branch, with the mori scatter change being the documented GQA-replication fix. No objection to the production diff.

The concern is what is left guarding it.

The stated justification for dropping the two e2e classes no longer holds.

The PR body argues the {plain, staging} x {gather, scatter} matrix can shrink because the arithmetic is now guarded on CPU:

With the indices guarded on CPU, mooncake's send_kvcache_slice has no direction branch left [...] so MooncakeMHADecodeLargerTP and StagingPrefillLargerTP are dropped

but the "Test cleanup" section then says:

The newly added CPU head-slice unit-test file was removed in follow-up CI cleanup.

test/registered/unit/disaggregation/test_head_slice_params.py is not on the branch. So the net effect of the PR on the test suite is +0 / -160: two e2e classes removed, and the CPU guard that was supposed to replace them removed as well. What is left is the manual equivalence sweep in the PR description, which is not executable and will not run again.

This matters more than usual for this particular change, because the PR itself documents the failure mode:

Wrong head indices do not fail loudly: the transfer delivers the wrong channels and the only symptom is a garbled accuracy score.

Could the unit test be restored? It is CPU-only at est_time=10, costs no GPU runner time, and was the whole argument for shrinking the GPU matrix. If it was dropped because of a registry/CI issue rather than on purpose, I am happy to help get it landed.

Second: the mori fix ships untested on every lane.

The mori scatter branch is the one real behaviour change here ([0, 1, 0, 1] -> [0, 0, 1, 1] for prefill TP=1 / decode TP=4 / 2 KV heads). Per the PR that path needs AMD CI, PR Test (AMD ROCm 10) is currently red, and test_mori_transfer_engine_e2e.py could not be dispatched from /rerun-test. With the CPU unit test also gone, this fix has no automated coverage at all on any lane. Was the AMD failure triaged as unrelated?

(Review by Claude Opus 5, run by @BBuf.)

[hnyls2002]

/rerun-test registered/unit/disaggregation/test_head_slice_params.py registered/disaggregation/test_disaggregation_different_tp.py registered/disaggregation/test_disaggregation_nixl.py

[github-actions]

Results for /rerun-test registered/unit/disaggregation/test_head_slice_params.py registered/disaggregation/test_disaggregation_different_tp.py registered/disaggregation/test_disaggregation_nixl.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/disaggregation/test_head_slice_params.py

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

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