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

Fix CP shared-KV bs=1 cache-hit zero-prefix corruption; spans are sole compose input

Browse Source 
A bs=1 cache-hit prefill produced deterministically wrong output (e.g. "0.5,0.5,0.5").
Root cause: regression a24111a5f4 changed the bs=1 cache-hit call sites to pass
prefix_slot_spans=[] (get_or_build_batch_slot_spans want_prefix=False) together with
prefix_pages>0. In materialize_prefix_and_reuse_current_kv_page_slots and its index twin,
the guard `if prefix_slot_spans is not None:` let the empty list shadow the prefix_pages
fallback -> prefix_spans=[] -> the cached prefix dense slots were never materialized
(both IPC and local paths gate on `if prefix_spans:`) -> attention read zero KV over the
entire reused prefix. The indexer compose dropped its prefix the same way. bs=1 only;
bs>1 (non-empty per-request spans) and fresh prefill (no prefix) were unaffected.

Fix: make the canonical per-request slot-span list the sole description of the
prefix/current regions and fail loud on a missing list. Remove all four
span-reconstruction fallbacks: prefix_slot_span (singular, dead), prefix_pages->span
(the bug), current_slot_spans=None->span (dead), and the prefetcher's copy. Both twins
now require prefix_slot_spans + current_slot_spans (raise on None; [] = genuine
no-prefix). All call sites build canonical spans via want_prefix=True (cached per-forward,
so a24111a5f4's per-layer-CPU goal is preserved by the cache, not by skipping the build);
the MLA current-only path also uses want_prefix=True to keep want_prefix uniform and the
span cache thrash-free.

Tests: migrate unit tests off the removed prefix_pages param (drop it where spans were
already passed; supply canonical spans otherwise); fix two source-string/captured-kwarg
assertions; add test_materialize_prefix_requires_explicit_prefix_slot_spans (None->raise,
explicit spans -> non-zero prefix). test_cp_shared_kv_runtime.py: 158 passed on the
g0033 container.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-06-17 18:25:22 +00:00
co-authored by Claude Opus 4.8
parent b34e7cb932
commit e23168e7f5
7 changed files with 156 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+24 -18
View File
@@ -443,15 +443,18 @@ class CpSharedKVMlaPrefetcher:
self.owned_prefix_pages = owned_prefix_pages
self.owned_total_pages = owned_total_pages
self.total_slots = int(slot_logical_pages.numel())
self.prefix_slot_spans = (
[(0, int(prefix_pages))] if prefix_slot_spans is None and prefix_pages > 0
else list(prefix_slot_spans or [])
)
self.current_slot_spans = (
[(int(prefix_pages), self.total_slots)]
if current_slot_spans is None and prefix_pages < self.total_slots
else list(current_slot_spans or [])
)
# Spans are the canonical description, always built by
# build_batch_*_slot_spans in build(). Fail loud on a missing list rather
# than silently reconstructing from prefix_pages (an explicit [] is a valid
# no-prefix compose; only None is a caller bug).
if prefix_slot_spans is None or current_slot_spans is None:
raise ValueError(
"CP shared KV prefetcher requires explicit prefix_slot_spans and "
"current_slot_spans ([] for none); got "
f"prefix={prefix_slot_spans} current={current_slot_spans}."
)
self.prefix_slot_spans = list(prefix_slot_spans)
self.current_slot_spans = list(current_slot_spans)
self.prefix_page_count = _slot_spans_page_count(self.prefix_slot_spans)
self.current_page_count = _slot_spans_page_count(self.current_slot_spans)
self.stream = stream if stream is not None else torch.cuda.Stream()
@@ -1305,15 +1308,18 @@ class CpSharedKVIndexPrefetcher:
self.owned_prefix_pages = owned_prefix_pages
self.owned_total_pages = owned_total_pages
self.total_slots = int(slot_logical_pages.numel())
self.prefix_slot_spans = (
[(0, int(prefix_pages))] if prefix_slot_spans is None and prefix_pages > 0
else list(prefix_slot_spans or [])
)
self.current_slot_spans = (
[(int(prefix_pages), self.total_slots)]
if current_slot_spans is None and prefix_pages < self.total_slots
else list(current_slot_spans or [])
)
# Spans are the canonical description, always built by
# build_batch_*_slot_spans in build(). Fail loud on a missing list rather
# than silently reconstructing from prefix_pages (an explicit [] is a valid
# no-prefix compose; only None is a caller bug).
if prefix_slot_spans is None or current_slot_spans is None:
raise ValueError(
"CP shared KV prefetcher requires explicit prefix_slot_spans and "
"current_slot_spans ([] for none); got "
f"prefix={prefix_slot_spans} current={current_slot_spans}."
)
self.prefix_slot_spans = list(prefix_slot_spans)
self.current_slot_spans = list(current_slot_spans)
self.prefix_page_count = _slot_spans_page_count(self.prefix_slot_spans)
self.current_page_count = _slot_spans_page_count(self.current_slot_spans)
self.stream = stream if stream is not None else torch.cuda.Stream()
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+26 -88
View File
@@ -5466,10 +5466,8 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
slot_remap: SharedTokenKVSlotRemap,
layout: CpSharedKVLayout,
page_size: int,
prefix_pages: int,
loc_req_id: torch.Tensor | None = None,
current_req_id: torch.Tensor | None = None,
prefix_slot_span: tuple[int, int] | None = None,
prefix_slot_spans: list[tuple[int, int]] | None = None,
current_slot_spans: list[tuple[int, int]] | None = None,
logical_locs_row_ids: torch.Tensor | None = None,
@@ -5501,35 +5499,17 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
)
total_slots = int(slot_remap.slot_logical_pages.numel())
if prefix_slot_spans is not None and prefix_slot_span is not None:
# Prefix is described ONLY by the canonical per-request slot-span list, built
# once per forward by get_or_build_batch_slot_spans. [] means a genuine
# no-prefix (current-only) compose; None means the caller failed to supply
# spans -> fail loud rather than silently materialize a zero prefix.
if prefix_slot_spans is None:
raise ValueError(
"Specify either prefix_slot_span or prefix_slot_spans, not both."
)
if prefix_slot_spans is not None:
prefix_spans = _merge_slot_spans(prefix_slot_spans)
for prefix_start_slot, prefix_end_slot in prefix_spans:
if (
prefix_start_slot < 0
or prefix_end_slot < prefix_start_slot
or prefix_end_slot > total_slots
):
raise ValueError(
"Invalid CP shared KV partial-current prefix slot span: "
f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} "
f"total_slots={total_slots}"
)
elif prefix_slot_span is None:
if prefix_pages < 0 or prefix_pages > total_slots:
raise ValueError(
"Invalid CP shared KV partial-current prefix range: "
f"prefix_pages={prefix_pages} total_slots={total_slots}"
)
prefix_spans = [(0, int(prefix_pages))] if int(prefix_pages) > 0 else []
else:
prefix_start_slot, prefix_end_slot = (
int(prefix_slot_span[0]),
int(prefix_slot_span[1]),
"CP shared KV partial-current compose requires explicit "
"prefix_slot_spans ([] for a no-prefix current-only compose); got None."
)
prefix_spans = _merge_slot_spans(prefix_slot_spans)
for prefix_start_slot, prefix_end_slot in prefix_spans:
if (
prefix_start_slot < 0
or prefix_end_slot < prefix_start_slot
@@ -5537,13 +5517,9 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
):
raise ValueError(
"Invalid CP shared KV partial-current prefix slot span: "
f"prefix_slot_span={prefix_slot_span} total_slots={total_slots}"
f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} "
f"total_slots={total_slots}"
)
prefix_spans = (
[(prefix_start_slot, prefix_end_slot)]
if prefix_end_slot > prefix_start_slot
else []
)
dense_kv_cache = kv_cache.new_zeros(
(slot_remap.dense_num_pages * page_size, *kv_cache.shape[1:])
@@ -5619,15 +5595,9 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
if layout.cp_size > 1:
current_materialized_by_ipc = False
if current_slot_spans is None:
if prefix_slot_span is not None or prefix_slot_spans is not None:
raise ValueError(
"CP shared KV batched current compose requires explicit "
"current_slot_spans to avoid reducing prefix slots twice."
)
current_slot_spans = (
[(int(prefix_pages), total_slots)]
if int(prefix_pages) < total_slots
else []
raise ValueError(
"CP shared KV batched current compose requires explicit "
"current_slot_spans."
)
merged_current_spans_for_reduce = _merge_slot_spans(current_slot_spans)
if merged_current_spans_for_reduce:
@@ -5747,9 +5717,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
layout: CpSharedKVLayout,
page_size: int,
index_head_dim: int,
prefix_pages: int,
current_req_id: torch.Tensor | None = None,
prefix_slot_span: tuple[int, int] | None = None,
prefix_slot_spans: list[tuple[int, int]] | None = None,
current_slot_spans: list[tuple[int, int]] | None = None,
layer_id: int | None = None,
@@ -5767,35 +5735,15 @@ def materialize_prefix_and_reuse_current_index_page_slots(
)
total_slots = int(slot_remap.slot_logical_pages.numel())
if prefix_slot_spans is not None and prefix_slot_span is not None:
# See materialize_prefix_and_reuse_current_kv_page_slots: prefix is described
# ONLY by the canonical span list. [] = genuine no-prefix; None = caller bug.
if prefix_slot_spans is None:
raise ValueError(
"Specify either prefix_slot_span or prefix_slot_spans, not both."
)
if prefix_slot_spans is not None:
prefix_spans = _merge_slot_spans(prefix_slot_spans)
for prefix_start_slot, prefix_end_slot in prefix_spans:
if (
prefix_start_slot < 0
or prefix_end_slot < prefix_start_slot
or prefix_end_slot > total_slots
):
raise ValueError(
"Invalid CP shared KV index partial-current prefix slot span: "
f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} "
f"total_slots={total_slots}"
)
elif prefix_slot_span is None:
if prefix_pages < 0 or prefix_pages > total_slots:
raise ValueError(
"Invalid CP shared KV index partial-current prefix range: "
f"prefix_pages={prefix_pages} total_slots={total_slots}"
)
prefix_spans = [(0, int(prefix_pages))] if int(prefix_pages) > 0 else []
else:
prefix_start_slot, prefix_end_slot = (
int(prefix_slot_span[0]),
int(prefix_slot_span[1]),
"CP shared KV index partial-current compose requires explicit "
"prefix_slot_spans ([] for a no-prefix current-only compose); got None."
)
prefix_spans = _merge_slot_spans(prefix_slot_spans)
for prefix_start_slot, prefix_end_slot in prefix_spans:
if (
prefix_start_slot < 0
or prefix_end_slot < prefix_start_slot
@@ -5803,13 +5751,9 @@ def materialize_prefix_and_reuse_current_index_page_slots(
):
raise ValueError(
"Invalid CP shared KV index partial-current prefix slot span: "
f"prefix_slot_span={prefix_slot_span} total_slots={total_slots}"
f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} "
f"total_slots={total_slots}"
)
prefix_spans = (
[(prefix_start_slot, prefix_end_slot)]
if prefix_end_slot > prefix_start_slot
else []
)
dense_page_buffer = page_buffer.new_zeros(
(slot_remap.dense_num_pages, *page_buffer.shape[1:])
@@ -5865,15 +5809,9 @@ def materialize_prefix_and_reuse_current_index_page_slots(
if layout.cp_size > 1:
current_materialized_by_ipc = False
if current_slot_spans is None:
if prefix_slot_span is not None or prefix_slot_spans is not None:
raise ValueError(
"CP shared KV batched index current compose requires explicit "
"current_slot_spans to avoid reducing prefix slots twice."
)
current_slot_spans = (
[(int(prefix_pages), total_slots)]
if int(prefix_pages) < total_slots
else []
raise ValueError(
"CP shared KV batched index current compose requires explicit "
"current_slot_spans."
)
merged_current_spans_for_reduce = _merge_slot_spans(current_slot_spans)
if merged_current_spans_for_reduce:
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+2 -7
View File
@@ -762,17 +762,13 @@ class Indexer(MultiPlatformOp):
current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
else:
current_req_id = current_req_id[: int(current_locs.shape[0])]
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
else:
prefix_pages = 0
prefix_slot_spans, current_slot_spans = get_or_build_batch_slot_spans(
forward_batch,
logical_pages=logical_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=len(prefix_lens_cpu) > 1,
want_prefix=True,
)
if index_prefetcher is not None:
prefetched = index_prefetcher.consume_prefix_with_current(
@@ -819,7 +815,6 @@ class Indexer(MultiPlatformOp):
layout=layout,
page_size=page_size,
index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
prefix_pages=prefix_pages,
current_req_id=current_req_id,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
@@ -838,7 +833,7 @@ class Indexer(MultiPlatformOp):
layer_id,
prefix_lens,
extend_lens,
prefix_pages,
sum(int(e) - int(s) for s, e in prefix_slot_spans),
prefix_slot_spans,
int(current_index_kv[0].shape[0]),
int(materialized.shape[0]),
python/sglang/srt/layers/attention/nsa_backend.py
+18 -34
View File
@@ -2085,7 +2085,7 @@ class NativeSparseAttnBackend(
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
)
_, current_slot_spans = get_or_build_batch_slot_spans(
prefix_slot_spans, current_slot_spans = get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=getattr(
@@ -2095,7 +2095,11 @@ class NativeSparseAttnBackend(
forward_batch, "extend_seq_lens_cpu", None
),
page_size=page_size,
want_prefix=False,
# current-only batch -> no prefix; want_prefix=True still
# yields [] here but keeps want_prefix uniformly True across
# the MLA + indexer call sites so the per-forward span cache
# (keyed by want_prefix) does not thrash layer-to-layer.
want_prefix=True,
)
kv_cache, page_table_1 = (
materialize_prefix_and_reuse_current_kv_page_slots(
@@ -2106,7 +2110,7 @@ class NativeSparseAttnBackend(
slot_remap=slot_remap,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
loc_req_id=cp_loc_req_id,
current_req_id=cp_current_req_id_full[
: int(current_locs_for_reuse.shape[0])
@@ -2212,30 +2216,16 @@ class NativeSparseAttnBackend(
f"current_locs_shape={tuple(current_locs_for_reuse.shape)} "
f"page_size={page_size}"
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=False,
)
)
else:
prefix_pages = 0
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=True,
)
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=True,
)
)
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
@@ -2252,7 +2242,6 @@ class NativeSparseAttnBackend(
slot_remap=slot_remap,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=prefix_pages,
loc_req_id=cp_loc_req_id,
current_req_id=cp_current_req_id_full[
: int(current_locs_for_reuse.shape[0])
@@ -2278,7 +2267,7 @@ class NativeSparseAttnBackend(
reason,
prefix_lens,
extend_lens,
prefix_pages,
sum(int(e) - int(s) for s, e in prefix_slot_spans),
prefix_slot_spans,
int(current_kv_cache.shape[0]),
int(kv_cache.shape[0]),
@@ -2585,10 +2574,6 @@ class NativeSparseAttnBackend(
f"extend_lens={extend_lens_cpu} "
f"page_size={page_size}"
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
else:
prefix_pages = 0
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
@@ -2596,7 +2581,7 @@ class NativeSparseAttnBackend(
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=len(prefix_lens_cpu) > 1,
want_prefix=True,
)
)
logical_locs_row_ids = (
@@ -2681,7 +2666,6 @@ class NativeSparseAttnBackend(
slot_remap=slot_remap,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=prefix_pages,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
loc_req_id=logical_locs_row_ids,
test/manual/test_cp_shared_kv_compose_v2_8rank.py
+2 -4
View File
@@ -393,7 +393,6 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
slot_remap=kv_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
loc_req_id=loc_req_id,
current_req_id=zeros_req,
layer_id=60,
@@ -408,7 +407,6 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
layout=layout,
page_size=page_size,
index_head_dim=128,
prefix_pages=0,
current_req_id=zeros_req,
layer_id=60,
**spans,
@@ -422,6 +420,7 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
layout=layout,
page_size=page_size,
prefix_pages=prefix_pages,
**spans,
slot_logical_pages=kv_remap.slot_logical_pages,
page_inverse=kv_remap.page_inverse,
dense_num_pages=kv_remap.dense_num_pages,
@@ -452,6 +451,7 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
idx = CpSharedKVIndexPrefetcher(
layout=layout,
prefix_pages=prefix_pages,
**spans,
slot_logical_pages=idx_remap.slot_logical_pages,
page_inverse=idx_remap.page_inverse,
dense_num_pages=idx_remap.dense_num_pages,
@@ -534,7 +534,6 @@ def _compose(s, layer_id: int, *, writers: list[int] | None = None):
slot_remap=s["slot_remap"],
layout=s["layout"],
page_size=s["page_size"],
prefix_pages=0,
loc_req_id=s["loc_req_id"],
current_req_id=s["current_req_id"],
prefix_slot_spans=s["prefix_slot_spans"],
@@ -662,7 +661,6 @@ def main() -> None:
layout=s["layout"],
page_size=s["page_size"],
index_head_dim=128,
prefix_pages=0,
current_req_id=si["current_req_id"],
prefix_slot_spans=s["prefix_slot_spans"],
current_slot_spans=s["current_slot_spans"],
test/registered/unit/layers/test_nsa_cp_utils.py
+1 -1
View File
@@ -3604,7 +3604,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
)
self.assertEqual(len(compose_calls), 1)
self.assertEqual(compose_calls[0]["prefix_pages"], 0)
self.assertEqual(compose_calls[0]["prefix_slot_spans"], [])
self.assertEqual(compose_calls[0]["current_slot_spans"], [(0, 2)])
self.assertIs(compose_calls[0]["current_index_k"], current_index_kv[0])
self.assertIs(compose_calls[0]["current_index_scale"], current_index_kv[1])
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+83 -14
View File
@@ -218,6 +218,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
page_size=4,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([0, 1, 2], dtype=torch.int64),
page_inverse=torch.tensor([0, 1, 2], dtype=torch.int64),
dense_num_pages=4,
@@ -594,6 +596,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([0, 1, 2], dtype=torch.int64),
page_inverse=torch.tensor([0, 1, 2], dtype=torch.int64),
dense_num_pages=4,
@@ -1292,7 +1296,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertIn(
"materialize_prefix_and_reuse_current_kv_page_slots", branch_source
)
self.assertIn("prefix_pages=0", branch_source)
# current-only compose now passes the canonical (empty) prefix_slot_spans
# instead of the removed prefix_pages scalar fallback.
self.assertIn("prefix_slot_spans=prefix_slot_spans", branch_source)
self.assertNotIn("prefix_pages=", branch_source)
self.assertNotIn("kv_cache = current_kv_cache", branch_source)
def test_mla_partial_current_sync_uses_batch_prefix_slot_spans(self):
@@ -2192,7 +2199,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
)
)
@@ -2202,6 +2209,56 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]])
def test_materialize_prefix_requires_explicit_prefix_slot_spans(self):
# Regression (a24111a5f4): a bs=1 cache-hit passed prefix_slot_spans=[] +
# prefix_pages>0; the `is not None` guard let [] shadow the prefix_pages
# fallback so the cached prefix was never materialized (zero KV over the
# reused prefix -> garbage output). The compose now takes prefix spans as
# the sole canonical input: None must fail loud; an explicit span list must
# materialize a NON-ZERO prefix.
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
layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0)
kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1)
logical_locs = torch.tensor([[4, 8, 20, 21, 22, 23]], dtype=torch.int64)
current_locs = torch.tensor([20, 21], dtype=torch.int64)
current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=torch.tensor([[1, 2, 5]], dtype=torch.int64),
layout=layout,
page_size=page_size,
)
def _compose(prefix_slot_spans):
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
):
return runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_kv_cache=current_kv,
current_locs=current_locs,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_slot_spans=prefix_slot_spans,
)
# None -> fail loud (never silently materialize a zero prefix).
with self.assertRaises(ValueError):
_compose(None)
# The canonical bs=1 cache-hit spans materialize the prefix as NON-ZERO
# (the rows the regression left empty).
mixed_kv, _ = _compose([(0, 2)])
self.assertGreater(mixed_kv[4:12].abs().sum().item(), 0.0)
self.assertTrue(torch.equal(mixed_kv[4:8], kv_cache[4:8]))
self.assertTrue(torch.equal(mixed_kv[8:12], kv_cache[8:12]))
def test_batch_prefix_slot_span_covers_bounding_prefix_range(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -2956,7 +3013,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
)
@@ -3027,7 +3083,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
)
@@ -3084,7 +3139,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
layer_id=2,
@@ -3249,7 +3303,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
)
)
@@ -3303,7 +3358,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_spans=[],
current_slot_spans=[(0, 1)],
layer_id=0,
)
@@ -3386,6 +3441,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
page_size=4,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64),
page_inverse=page_inverse,
dense_num_pages=4,
@@ -3452,6 +3509,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
page_size=4,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64),
page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64),
dense_num_pages=4,
@@ -3499,6 +3558,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
page_size=4,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64),
page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64),
dense_num_pages=4,
@@ -3538,6 +3599,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
page_size=4,
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64),
page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64),
dense_num_pages=4,
@@ -3578,6 +3641,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
prefix_pages=2,
prefix_slot_spans=[(0, 2)],
current_slot_spans=[(2, 3)],
slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64),
page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64),
dense_num_pages=4,
@@ -3623,6 +3688,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=0),
prefix_pages=1,
prefix_slot_spans=[(0, 1)],
current_slot_spans=[(1, 3)],
slot_logical_pages=torch.tensor([1, 2, 3], dtype=torch.int64),
page_inverse=torch.tensor([-1, 1, 2, 3], dtype=torch.int64),
dense_num_pages=padded_pages,
@@ -3699,7 +3766,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=1,
prefix_slot_spans=[(0, 1)],
layer_id=2,
)
)
@@ -3831,7 +3898,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=1,
prefix_slot_spans=[(0, 1)],
current_slot_spans=[(1, 2)],
layer_id=2,
)
)
@@ -3889,7 +3957,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=0,
prefix_slot_spans=[],
current_slot_spans=[(0, 1)],
layer_id=0,
)
@@ -3968,6 +4036,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0),
prefix_pages=1,
prefix_slot_spans=[(0, 1)],
current_slot_spans=[(1, 2)],
slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64),
page_inverse=torch.tensor(
[-1, 1] + [-1] * 18 + [2],
@@ -4036,6 +4106,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
prefetcher = prefetch.CpSharedKVIndexPrefetcher(
layout=CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0),
prefix_pages=1,
prefix_slot_spans=[(0, 1)],
current_slot_spans=[(1, 2)],
slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64),
page_inverse=torch.tensor(
[-1, 1] + [-1] * 18 + [2],
@@ -5620,7 +5692,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=[],
layer_id=0,
@@ -5997,7 +6068,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
layer_id=0,
@@ -6056,7 +6126,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=1,
prefix_slot_spans=[(0, 1)],
current_slot_spans=[(1, 3)],
layer_id=0,
)
@@ -6180,7 +6250,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=0,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
layer_id=0,