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/mem_cache/test_nsa_pool_host_unit.py
+145 -1
View File
@@ -147,6 +147,85 @@ class TestLayerPageFirstDirectHostLayout(CustomTestCase):
(host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size),
)
def test_nsa_page_first_direct_indexer_layout_compacts_active_layers(self):
indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype
device_pool = SimpleNamespace(
store_dtype=torch.float16,
size=8,
start_layer=0,
end_layer=4,
device="cpu",
kv_lora_rank=16,
qk_rope_head_dim=4,
kv_cache_dim=24,
layer_num=4,
index_head_dim=16,
quant_block_size=8,
index_active_layer_ids=(0, 2),
index_k_with_scale_buffer=[
torch.empty((5, 96), dtype=indexer_dtype) for _ in range(2)
],
)
host_pool = NSATokenToKVPoolHost(
device_pool=device_pool,
host_to_device_ratio=2.0,
host_size=0,
page_size=4,
layout="page_first_direct",
pin_memory=False,
device="cpu",
host_token_capacity=16,
)
self.assertEqual(host_pool.index_active_layer_ids, (0, 2))
self.assertEqual(host_pool._host_index_layer_slot(0), 0)
self.assertEqual(host_pool._host_index_layer_slot(2), 1)
self.assertEqual(
tuple(host_pool.index_k_with_scale_buffer.shape),
(host_pool.indexer_page_num, 2, 1, host_pool.indexer_page_stride_size),
)
def test_nsa_layer_page_first_indexer_layout_compacts_active_layers(self):
indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype
device_pool = SimpleNamespace(
store_dtype=torch.float16,
size=8,
start_layer=0,
end_layer=4,
device="cpu",
kv_lora_rank=16,
qk_rope_head_dim=4,
kv_cache_dim=24,
layer_num=4,
index_head_dim=16,
quant_block_size=8,
index_active_layer_ids=(0, 2),
index_k_with_scale_buffer=[
torch.empty((5, 96), dtype=indexer_dtype) for _ in range(2)
],
)
host_pool = NSATokenToKVPoolHost(
device_pool=device_pool,
host_to_device_ratio=2.0,
host_size=0,
page_size=4,
layout="layer_page_first",
pin_memory=False,
device="cpu",
host_token_capacity=16,
)
self.assertEqual(host_pool.index_active_layer_ids, (0, 2))
self.assertEqual(host_pool._host_index_layer_slot(0), 0)
self.assertEqual(host_pool._host_index_layer_slot(2), 1)
self.assertEqual(
tuple(host_pool.index_k_with_scale_buffer.shape),
(2, host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size),
)
self.assertEqual(len(host_pool.index_k_data_refs), 2)
def test_mha_layer_page_first_page_buffer_meta_fails_fast_for_storage(self):
device_pool = SimpleNamespace(
store_dtype=torch.float16,
@@ -285,8 +364,9 @@ class TestNSAHiCacheTransfer(CustomTestCase):
device="cuda" if io_backend == "kernel" else "cpu",
dtype=torch.int64,
)
index_device = "cuda" if io_backend == "kernel" else "cpu"
device_indices = self._token_indices_for_pages(
device_pages, page_size, device="cuda"
device_pages, page_size, device=index_device
)
host_indices = self._token_indices_for_pages(
host_pages,
@@ -1719,6 +1799,70 @@ class TestNSAIndexerPageIndices(CustomTestCase):
self.assertTrue(pool.is_index_layer_active(11))
self.assertEqual(pool.get_index_layer_slot(11), 7)
def test_nsa_state_buf_infos_returns_active_index_layers_only(self):
pool = object.__new__(NSATokenToKVPool)
pool.start_layer = 4
pool.end_layer = 12
pool.layer_num = 8
pool.index_k_with_scale_buffer = [
torch.empty((3, 8), dtype=torch.uint8) for _ in range(pool.layer_num)
]
pool._init_index_layer_metadata(
index_active_layer_ids=(4, 8),
compact_index_layers=False,
)
data_ptrs, data_lens, item_lens = pool.get_state_buf_infos()
self.assertEqual(
data_ptrs,
[
pool.index_k_with_scale_buffer[0].data_ptr(),
pool.index_k_with_scale_buffer[4].data_ptr(),
],
)
self.assertEqual(
data_lens,
[
pool.index_k_with_scale_buffer[0].nbytes,
pool.index_k_with_scale_buffer[4].nbytes,
],
)
self.assertEqual(
item_lens,
[
pool.index_k_with_scale_buffer[0][0].nbytes,
pool.index_k_with_scale_buffer[4][0].nbytes,
],
)
self.assertEqual(pool.get_state_layer_ids(), [4, 8])
def test_nsa_device_pool_compact_index_layers_allocates_active_slots_only(self):
if not torch.cuda.is_available():
self.skipTest("CUDA is required for compact NSA pool allocation test.")
pool = NSATokenToKVPool(
size=128,
page_size=64,
kv_lora_rank=128,
dtype=torch.bfloat16,
qk_rope_head_dim=32,
layer_num=8,
device="cuda",
enable_memory_saver=False,
kv_cache_dim=576,
index_head_dim=128,
start_layer=4,
end_layer=12,
index_active_layer_ids=(4, 8),
compact_index_layers=True,
)
self.assertEqual(len(pool.index_k_with_scale_buffer), 2)
self.assertEqual(pool.get_index_layer_slot(4), 0)
self.assertEqual(pool.get_index_layer_slot(8), 1)
self.assertEqual(len(pool.get_state_buf_infos()[0]), 2)
def test_indexer_page_indices_accepts_valid_page_spans(self):
host_pool = self.make_host_pool_stub(page_size=4)