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,