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

Enable index partial-current reuse without replaying prefix materialize

Browse Source 
The index path now mirrors the target MLA partial-current contract: prefetched or synchronously materialized prefix pages are composed with valid current index K/scale rows in slot-dense page buffers. Current-only batches keep the compact current-index fast path, while partial cache-hit batches share one composed dense index buffer across the in-seq prev/next topk pair.\n\nThe prefetch consume path remaps through the slot page inverse instead of treating the slot-dense buffer as physical-pool capacity, and current-index quantization uses valid extend rows so padded out_cache_loc does not disable reuse.\n\nConstraint: CP shared KV remains page-slot based; padding rows must stay invisible to attention/index semantics\nConstraint: Draft/EAGLE partial-current reuse remains guarded by should_reuse_current_extend_kv\nRejected: Replace prefix all-reduce with all-gather | NCCL all-gather still uses SM and would require an additional compose/scatter step\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce current-only gating for index reuse; partial target cache hits must compose prefix + valid current rows\nTested: Local py_compile for touched Python files\nTested: g0034 sglang-glm5-dev-2 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 77 passed, 5 warnings, 2 subtests passed\nNot-tested: Full ETE traffic with latest commit; CUDA perf impact of index partial-current prefetch under production load
This commit is contained in:
laoyao0822
2026-05-31 02:31:09 +08:00
parent 251a48fb0a
commit 3c14b1f127
5 changed files with 604 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+176
View File
@@ -1021,6 +1021,182 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertLess(valid_tokens, padded_pages * page_size)
self.assertEqual(valid_tokens, 10)
def test_index_partial_current_sync_compose_fills_current_page_slots(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
page_size = 4
index_head_dim = 4
scale_bytes = 4
page_bytes = page_size * index_head_dim + page_size * scale_bytes
layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0)
page_buffer = torch.zeros((6, page_bytes), dtype=torch.uint8)
page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8)
page_buffer[2] = 99
logical_pages = torch.tensor([[1, 2]], dtype=torch.int64)
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
page_buffer,
logical_pages,
layout,
)
current_k = torch.tensor(
[[1, 2, 3, 4], [5, 6, 7, 8]],
dtype=torch.uint8,
)
current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32)
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
):
dense_page_buffer, dense_pages = (
runtime.materialize_prefix_and_reuse_current_index_page_slots(
page_buffer=page_buffer,
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=torch.tensor([8, 9], dtype=torch.int64),
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=1,
layer_id=2,
)
)
scale_offset = page_size * index_head_dim
self.assertEqual(dense_pages.tolist(), [[1, 2]])
self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1]))
self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0]))
self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1]))
self.assertTrue(
torch.equal(
dense_page_buffer[2, 8:scale_offset],
torch.zeros(8, dtype=torch.uint8),
)
)
self.assertTrue(
torch.equal(
dense_page_buffer[2, scale_offset : scale_offset + 4],
current_scale[0].view(torch.uint8),
)
)
self.assertTrue(
torch.equal(
dense_page_buffer[2, scale_offset + 4 : scale_offset + 8],
current_scale[1].view(torch.uint8),
)
)
self.assertTrue(
torch.equal(
dense_page_buffer[2, scale_offset + 8 :],
torch.zeros(8, dtype=torch.uint8),
)
)
def test_index_prefetch_partial_current_compose_fills_current_page_slots(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
class FakeCurrentStream:
def __init__(self):
self.events = []
def wait_event(self, event):
self.events.append(event)
page_size = 4
index_head_dim = 4
scale_bytes = 4
page_bytes = page_size * index_head_dim + page_size * scale_bytes
dense_page_buffer = torch.zeros((3, page_bytes), dtype=torch.uint8)
dense_page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8)
fake_event = object()
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0),
prefix_pages=1,
slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64),
page_inverse=torch.tensor(
[-1, 1] + [-1] * 18 + [2],
dtype=torch.int64,
),
dense_num_pages=3,
stream=object(),
)
handle = prefetch.CpSharedKVIndexPrefetchHandle(
layer_id=1,
dense_page_buffer=dense_page_buffer,
prefix_rows=slice(1, 2),
event=fake_event,
)
prefetcher.handles[1] = handle
prefetcher.pending_attention_handle = handle
current_k = torch.tensor(
[[11, 12, 13, 14], [15, 16, 17, 18]],
dtype=torch.uint8,
)
current_scale = torch.tensor([[3.25], [4.5]], dtype=torch.float32)
with envs.SGLANG_DEBUG_CP_SHARED_KV.override(True), patch.object(
prefetch.torch.cuda,
"current_stream",
return_value=FakeCurrentStream(),
):
mixed_buffer, dense_pages = prefetcher.consume_prefix_with_current(
layer_id=1,
logical_pages=torch.tensor([[1, 20]], dtype=torch.int64),
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=torch.tensor([80, 81], dtype=torch.int64),
page_size=page_size,
index_head_dim=index_head_dim,
)
scale_offset = page_size * index_head_dim
self.assertEqual(dense_pages.tolist(), [[1, 2]])
self.assertTrue(torch.equal(mixed_buffer[2, 0:4], current_k[0]))
self.assertTrue(torch.equal(mixed_buffer[2, 4:8], current_k[1]))
self.assertTrue(
torch.equal(
mixed_buffer[2, scale_offset : scale_offset + 4],
current_scale[0].view(torch.uint8),
)
)
self.assertEqual(prefetcher.handles, {})
self.assertIsNone(prefetcher.pending_attention_handle)
def test_index_current_reuse_gate_uses_partial_current_contract(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
).read_text()
start = source.index(" def _can_reuse_current_index_kv")
end = source.index(" @contextlib.contextmanager", start)
method_source = source[start:end]
self.assertIn("should_reuse_current_extend_kv(forward_batch)", method_source)
self.assertNotIn("is_current_only_extend_batch(forward_batch)", method_source)
def test_index_current_reuse_prepare_accepts_padded_out_cache_loc(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
).read_text()
start = source.index(" if self._can_reuse_current_index_kv")
end = source.index(" if _is_cuda or _is_hip:", start)
prepare_source = source[start:end]
self.assertIn("valid_current_rows", prepare_source)
self.assertIn("key[:valid_current_rows]", prepare_source)
self.assertNotIn(
"key.shape[0] == forward_batch.out_cache_loc.numel()",
prepare_source,
)
def test_materialize_local_token_kv_pages(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_dense_page_remap,