chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Compact skipped NSA index-cache state safely

Browse Source 
Index skip reduces the number of target layers that own NSA index state,
but PD transfer and HiCache still assumed dense full-layer state buffers.
This change carries explicit state layer IDs through prefill/decode
registration, compacts device and host index buffers to active layers,
and maps logical layer IDs to compact slots on transfer paths.

The PD side fails fast when prefill/decode disagree on NSA state layer
identity instead of silently truncating or copying mismatched buffers.
Host direct tests now use the same CPU-index descriptor contract required
by the TAI cudaMemcpyBatchAsync path, and host registered memory is
unregistered on tensor finalization to avoid stale cudaHostRegister state
across CUDA tests.

Constraint: CP shared-KV with index_topk skip must keep target/draft state identity explicit before compacting buffers
Constraint: Direct HiCache TAI transfer rejects CUDA indices to avoid hidden D2H copies on the control path
Rejected: Keep full-layer L1/L2 index buffers | wastes the memory/bandwidth that index skip is meant to save
Rejected: Infer state buffer order by count only | can silently corrupt cache when active layer sets differ
Confidence: high
Scope-risk: moderate
Directive: Do not compact or reorder NSA state buffers without carrying logical layer IDs through PD registration and validating both sides
Tested: Remote container py_compile for touched runtime files
Tested: Remote container pytest: test_nsa_pool_host_unit.py, test_model_runner_kv_cache_mixin.py, test_cp_shared_kv_transfer_mapping.py, test_pd_state_layer_ids.py, test_cp_per_layer_transfer.py, test_cp_shared_kv_runtime.py -> 200 passed, 2 subtests passed
Not-tested: Full ETE GSM8K/replay after compacted P3-P6 changes
Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-10 05:37:06 +08:00
co-authored by OmX
parent d21952b903
commit 1ebde44e59
14 changed files with 611 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py
+22
View File
@@ -134,6 +134,28 @@ class TestCPSharedKVTransferMapping(unittest.TestCase):
self.assertEqual(kv_args.state_data_lens, [12, 23, 24])
self.assertEqual(kv_args.state_item_lens, [13, 25, 26])
def test_append_cp_draft_state_buffers_adds_negative_state_layer_ids(self):
kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[11],
state_data_lens=[12],
state_item_lens=[13],
state_layer_ids=[0],
)
appended = append_cp_draft_state_buffers(
kv_args,
draft_state_type="nsa",
draft_state_data_ptrs=[21, 22],
draft_state_data_lens=[23, 24],
draft_state_item_lens=[25, 26],
role="prefill",
cp_rank=2,
)
self.assertTrue(appended)
self.assertEqual(kv_args.state_layer_ids, [0, -1, -2])
def test_append_cp_draft_state_buffers_rejects_state_mismatch_under_shared_kv(self):
kv_args = SimpleNamespace(
state_type="nsa",
test/registered/unit/disaggregation/test_pd_state_layer_ids.py
+182
View File
@@ -0,0 +1,182 @@
import concurrent.futures
import struct
from types import SimpleNamespace
import numpy as np
import pytest
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
mooncake_conn = pytest.importorskip("sglang.srt.disaggregation.mooncake.conn")
KVArgsRegisterInfo = mooncake_conn.KVArgsRegisterInfo
MooncakeKVManager = mooncake_conn.MooncakeKVManager
def _pack_q(values):
return b"".join(struct.pack("Q", value) for value in values)
def _pack_i(values):
return b"".join(struct.pack("i", value) for value in values)
def _pack_u(values):
return b"".join(struct.pack("I", value) for value in values)
def _register_msg(*, state_layer_ids):
return [
b"room-a",
b"127.0.0.1",
b"1234",
b"session-a",
_pack_q([101, 102]),
_pack_q([201]),
_pack_q([301, 302, 303]),
b"0",
b"8",
b"64",
_pack_u([4096, 4096, 4096]),
_pack_u([128, 128, 128]),
_pack_i(state_layer_ids),
]
def test_mooncake_register_info_roundtrips_state_layer_ids():
info = KVArgsRegisterInfo.from_zmq(
_register_msg(state_layer_ids=[0, 4, 8, -1])
)
assert info.dst_state_layer_ids == [0, 4, 8, -1]
def test_mooncake_state_layer_id_mismatch_fails_fast():
manager = MooncakeKVManager.__new__(MooncakeKVManager)
manager.kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[101, 102],
state_item_lens=[64, 64],
state_layer_ids=[0, 4],
draft_state_type="none",
draft_state_buffer_count=0,
)
manager.attn_cp_rank = 3
manager.attn_tp_size = 8
manager.is_mla_backend = True
manager._send_kvcache_generic = lambda **_: 0
req = SimpleNamespace(
room="room-a",
mooncake_session_id="session-a",
dst_state_indices=np.array([11, 12], dtype=np.int32),
)
target_info = SimpleNamespace(
dst_state_layer_ids=[0, 5],
dst_state_data_ptrs=[201, 202],
dst_attn_tp_size=8,
dst_state_item_lens=[64, 64],
dst_state_dim_per_tensor=[],
dst_tp_rank=0,
)
with pytest.raises(RuntimeError, match=r"\[CP_SHARED_KV_FAIL_FAST\]\[state_layer_ids\].*prefill=\[0, 4\].*decode=\[0, 5\]"):
MooncakeKVManager.maybe_send_extra(
manager,
req,
prefill_state_indices=[1, 2],
dst_state_data_ptrs=[201, 202],
executor=concurrent.futures.ThreadPoolExecutor(max_workers=1),
target_rank_registration_info=target_info,
)
def test_mooncake_layer_aware_state_buffer_count_mismatch_fails_fast():
manager = MooncakeKVManager.__new__(MooncakeKVManager)
manager.kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[101, 102, 103],
state_item_lens=[64, 64, 64],
state_layer_ids=[0, 4, 8],
draft_state_type="none",
draft_state_buffer_count=0,
)
manager.attn_cp_rank = 3
manager.attn_tp_size = 8
manager.is_mla_backend = True
manager._send_kvcache_generic = lambda **_: 0
req = SimpleNamespace(
room="room-b",
mooncake_session_id="session-b",
dst_state_indices=np.array([11, 12, 13], dtype=np.int32),
)
target_info = SimpleNamespace(
dst_state_layer_ids=[0, 4, 8],
dst_state_data_ptrs=[201, 202],
dst_attn_tp_size=8,
dst_state_item_lens=[64, 64],
dst_state_dim_per_tensor=[],
dst_tp_rank=0,
)
with pytest.raises(RuntimeError, match=r"\[CP_SHARED_KV_FAIL_FAST\]\[state_buffer_count\].*src=3.*dst=2"):
MooncakeKVManager.maybe_send_extra(
manager,
req,
prefill_state_indices=[1, 2, 3],
dst_state_data_ptrs=[201, 202],
executor=concurrent.futures.ThreadPoolExecutor(max_workers=1),
target_rank_registration_info=target_info,
)
def test_nixl_register_info_roundtrips_state_layer_ids():
nixl_conn = pytest.importorskip("sglang.srt.disaggregation.nixl.conn")
msg = [
b"room-a",
b"127.0.0.1",
b"1234",
b"agent-a",
b"metadata",
_pack_q([101, 102]),
_pack_q([201]),
_pack_q([301, 302]),
b"0",
b"8",
b"3",
b"64",
_pack_i([0, 4, -1]),
]
info = nixl_conn.KVArgsRegisterInfo.from_zmq(msg)
assert info.dst_state_layer_ids == [0, 4, -1]
def test_nixl_state_layer_id_mismatch_fails_fast():
nixl_conn = pytest.importorskip("sglang.srt.disaggregation.nixl.conn")
manager = nixl_conn.NixlKVManager.__new__(nixl_conn.NixlKVManager)
manager.kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[101, 102],
state_item_lens=[64, 64],
state_layer_ids=[0, 4],
)
manager.attn_cp_rank = 3
manager.attn_tp_size = 8
manager.is_mla_backend = True
manager._send_kvcache_generic = lambda **_: 0
with pytest.raises(RuntimeError, match=r"\[CP_SHARED_KV_FAIL_FAST\]\[state_layer_ids\].*prefill=\[0, 4\].*decode=\[0, 5\]"):
nixl_conn.NixlKVManager.maybe_send_extra(
manager,
peer_name="decode-a",
prefill_state_indices=[1, 2],
dst_state_data_ptrs=[201, 202],
dst_state_indices=[11, 12],
dst_gpu_id=0,
notif="notif-a",
decode_tp_size=8,
dst_state_layer_ids=[0, 5],
)