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

Preserve CP slot-remap reference paths after SYH pick

Browse Source 
The SYH per-request page_inverse migration fixes bs>1 cache contamination, but the conflict resolution made several helper/reference paths require explicit request ids and broke existing unit coverage. This keeps production bs>1 callers on explicit req-id routing while allowing bs=1/reference helpers to infer or default request ids without reintroducing the hot-path batch-global inverse.

Constraint: Runtime bs>1 materialize paths must route through per-request page_inverse rows to avoid cross-request KV aliasing.

Constraint: The remote test container may still have single-row tai-kernel materialize helpers while production bs>1 requires the new req-id ABI.

Rejected: Revert to batch-global page_inverse | it is the GSM8K cache-hit corruption root cause.

Rejected: Update tests only | the helper API is still useful for bs=1/reference callers and py-level regression coverage.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove explicit loc_req_id/current_req_id from production bs>1 call sites; default inference is for legacy/reference use only.

Tested: Local py_compile for cp_shared_kv_runtime.py and nsa_indexer.py.

Tested: Remote container py_compile for cp_shared_kv_runtime.py.

Tested: Remote container pytest: test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py => 263 passed, 5 warnings, 2 subtests passed.

Not-tested: Full ETE GSM8K/cache-hit run after SYH pick.
This commit is contained in:
laoyao0822
2026-06-09 00:19:15 +08:00
parent b35495a28d
commit eabe0bbfad
2 changed files with 249 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

313
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -531,17 +531,29 @@ def _tai_current_slot_fill_sparse_pages_self_test(
(int(current_locs.numel()),), device=device, dtype=torch.long
)
mixed_kv, mixed_locs, current_mask = fill_kernel(
dense_kv,
materialized_locs,
current_kv,
logical_locs,
current_locs,
page_inverse,
current_req_id,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
try:
mixed_kv, mixed_locs, current_mask = fill_kernel(
dense_kv,
materialized_locs,
current_kv,
logical_locs,
current_locs,
page_inverse,
current_req_id,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
except TypeError:
mixed_kv, mixed_locs, current_mask = fill_kernel(
dense_kv,
materialized_locs,
current_kv,
logical_locs,
current_locs,
page_inverse,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
expected_locs = torch.tensor(
[[4, 5, 8, 9, 12, 13, -1, -1]],
device=device,
@@ -1345,17 +1357,34 @@ def _try_tai_materialize_token_kv_pages_and_locs(
try:
tai_slot_logical_pages = _contiguous_for_tai(slot_logical_pages.reshape(-1))
page_inverse = kernels.build_slot_page_inverse(
tai_slot_logical_pages,
logical_page_capacity,
int(max(int(batch_rows), 1)),
)
dense_locs = kernels.remap_logical_locs_to_slot_dense_locs(
_contiguous_for_tai(logical_locs),
page_inverse,
_row_req_id_for_tai(loc_req_id, logical_locs),
page_size=page_size,
)
try:
page_inverse = kernels.build_slot_page_inverse(
tai_slot_logical_pages,
logical_page_capacity,
int(max(int(batch_rows), 1)),
)
except TypeError:
if int(batch_rows) != 1:
raise
page_inverse = kernels.build_slot_page_inverse(
tai_slot_logical_pages,
logical_page_capacity,
)
try:
dense_locs = kernels.remap_logical_locs_to_slot_dense_locs(
_contiguous_for_tai(logical_locs),
page_inverse,
_row_req_id_for_tai(loc_req_id, logical_locs),
page_size=page_size,
)
except TypeError:
if int(batch_rows) != 1:
raise
dense_locs = kernels.remap_logical_locs_to_slot_dense_locs(
_contiguous_for_tai(logical_locs),
page_inverse,
page_size=page_size,
)
dense_kv_cache = kernels.materialize_shared_token_kv_pages(
kv_cache,
tai_slot_logical_pages,
@@ -1388,11 +1417,20 @@ def _try_tai_build_slot_page_inverse(
return None
try:
return kernels.build_slot_page_inverse(
_contiguous_for_tai(slot_logical_pages.reshape(-1)),
logical_page_capacity,
int(max(int(batch_rows), 1)),
)
tai_slot_logical_pages = _contiguous_for_tai(slot_logical_pages.reshape(-1))
try:
return kernels.build_slot_page_inverse(
tai_slot_logical_pages,
logical_page_capacity,
int(max(int(batch_rows), 1)),
)
except TypeError:
if int(batch_rows) != 1:
raise
return kernels.build_slot_page_inverse(
tai_slot_logical_pages,
logical_page_capacity,
)
except Exception as exc:
_log_tai_materialize_fallback(
"page_inverse_failed",
@@ -1515,10 +1553,76 @@ def build_page_table_row_req_id(logical_pages: torch.Tensor) -> torch.Tensor:
return torch.zeros_like(logical_pages, dtype=torch.long)
def _slot_page_inverse_2d(page_inverse: torch.Tensor) -> torch.Tensor:
if page_inverse.dim() == 1:
return page_inverse.unsqueeze(0)
return page_inverse
def _infer_req_id_for_locs_from_slot_remap(
logical_locs: torch.Tensor,
slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap,
page_size: int,
*,
strict: bool = False,
) -> torch.Tensor:
"""Infer row ids for legacy/reference callers that omit explicit req ids.
Runtime bs>1 paths pass request ids explicitly. This helper keeps older
tests and bs=1 helper use working without reintroducing the batch-global
page inverse: for multi-row slot remaps it resolves each loc by membership
in exactly one row of the request-scoped slot map.
"""
if logical_locs.numel() == 0:
return torch.zeros_like(logical_locs, dtype=torch.long)
page_inverse = _slot_page_inverse_2d(slot_remap.page_inverse)
rows = int(page_inverse.shape[0])
if rows <= 1:
return torch.zeros_like(logical_locs, dtype=torch.long)
sorted_pages = getattr(slot_remap, "slot_sorted_logical_pages_by_row", None)
if sorted_pages is None or sorted_pages.dim() != 2:
raise ValueError(
"CP shared KV per-request slot remap requires explicit req ids when "
"row-scoped slot pages are unavailable."
)
locs_long = logical_locs.to(torch.long)
valid_locs = locs_long >= 0
pages = torch.div(
torch.where(valid_locs, locs_long, torch.zeros_like(locs_long)),
page_size,
rounding_mode="floor",
)
req_id = torch.zeros_like(locs_long, dtype=torch.long)
match_count = torch.zeros_like(locs_long, dtype=torch.long)
for row in range(rows):
row_pages = sorted_pages[row]
row_pages = row_pages[row_pages >= 0]
if row_pages.numel() == 0:
continue
row_match = valid_locs & torch.isin(pages, row_pages)
req_id = torch.where(row_match, torch.full_like(req_id, row), req_id)
match_count = match_count + row_match.to(torch.long)
ambiguous = valid_locs & (match_count != 1)
if torch.any(ambiguous):
if strict:
bad_pages = pages[ambiguous]
raise ValueError(
"CP shared KV could not infer unique per-request ids for slot remap. "
f"bad_page_min={int(bad_pages.min().item())} "
f"bad_page_max={int(bad_pages.max().item())}"
)
req_id = torch.where(ambiguous, torch.zeros_like(req_id), req_id)
return req_id
def build_slot_page_inverse_optimized(
slot_logical_pages: torch.Tensor,
logical_page_capacity: int,
batch_rows: int,
batch_rows: int = 1,
) -> torch.Tensor:
tai_result = _try_tai_build_slot_page_inverse(
slot_logical_pages,
@@ -1538,8 +1642,10 @@ def remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
loc_req_id: torch.Tensor,
loc_req_id: torch.Tensor | None = None,
) -> torch.Tensor:
if loc_req_id is None:
loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
if _tai_materialize_runtime_enabled():
kernels = _load_tai_materialize_kernels()
if kernels is not None:
@@ -1550,6 +1656,14 @@ def remap_logical_locs_to_slot_dense_locs_optimized(
_row_req_id_for_tai(loc_req_id, logical_locs),
page_size=page_size,
)
except TypeError:
if int(_slot_page_inverse_2d(page_inverse).shape[0]) != 1:
raise
return kernels.remap_logical_locs_to_slot_dense_locs(
_contiguous_for_tai(logical_locs),
_contiguous_for_tai(page_inverse),
page_size=page_size,
)
except Exception as exc:
_log_tai_materialize_fallback(
"loc_remap_failed",
@@ -1574,8 +1688,8 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs(
current_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
current_req_id: torch.Tensor,
mask_non_current_in_current_pages: bool,
current_req_id: torch.Tensor | None = None,
mask_non_current_in_current_pages: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
if not _tai_materialize_runtime_enabled():
_log_tai_materialize_runtime_disabled("fill_current_kv_page_slots")
@@ -1602,6 +1716,9 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs(
)
return None
if current_req_id is None:
current_req_id = torch.zeros_like(current_locs.reshape(-1), dtype=torch.long)
if dense_kv_cache.is_cuda and not _tai_current_slot_fill_supports_sparse_pages(
dense_kv_cache.device
):
@@ -1619,17 +1736,31 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs(
return None
try:
return fill_kernel(
_contiguous_for_tai(dense_kv_cache),
_contiguous_for_tai(materialized_dense_locs),
_contiguous_for_tai(current_kv_cache),
_contiguous_for_tai(logical_locs),
_contiguous_for_tai(current_locs.reshape(-1)),
_contiguous_for_tai(page_inverse),
_row_req_id_for_tai(current_req_id, current_locs.reshape(-1)),
page_size=int(page_size),
mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages),
)
try:
return fill_kernel(
_contiguous_for_tai(dense_kv_cache),
_contiguous_for_tai(materialized_dense_locs),
_contiguous_for_tai(current_kv_cache),
_contiguous_for_tai(logical_locs),
_contiguous_for_tai(current_locs.reshape(-1)),
_contiguous_for_tai(page_inverse),
_row_req_id_for_tai(current_req_id, current_locs.reshape(-1)),
page_size=int(page_size),
mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages),
)
except TypeError:
if int(_slot_page_inverse_2d(page_inverse).shape[0]) != 1:
raise
return fill_kernel(
_contiguous_for_tai(dense_kv_cache),
_contiguous_for_tai(materialized_dense_locs),
_contiguous_for_tai(current_kv_cache),
_contiguous_for_tai(logical_locs),
_contiguous_for_tai(current_locs.reshape(-1)),
_contiguous_for_tai(page_inverse),
page_size=int(page_size),
mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages),
)
except Exception as exc:
_log_tai_materialize_fallback(
"fill_current_failed",
@@ -1652,7 +1783,7 @@ def fill_current_kv_page_slots_and_remap_locs(
current_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
current_req_id: torch.Tensor,
current_req_id: torch.Tensor | None = None,
mask_non_current_in_current_pages: bool = True,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Fill current suffix KV into preallocated dense page slots and remap locs.
@@ -1665,6 +1796,9 @@ def fill_current_kv_page_slots_and_remap_locs(
land on another request's dense slot (bs>1 contamination).
"""
if current_req_id is None:
current_req_id = torch.zeros_like(current_locs.reshape(-1), dtype=torch.long)
tai_result = _try_tai_fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_kv_cache,
materialized_dense_locs=materialized_dense_locs,
@@ -1678,6 +1812,7 @@ def fill_current_kv_page_slots_and_remap_locs(
)
if tai_result is not None:
return tai_result
page_inverse = _slot_page_inverse_2d(page_inverse)
if dense_kv_cache.is_cuda:
_log_tai_materialize_fallback(
@@ -1788,7 +1923,7 @@ def fill_current_index_page_slots(
page_inverse: torch.Tensor,
page_size: int,
index_head_dim: int,
current_req_id: torch.Tensor,
current_req_id: torch.Tensor | None = None,
) -> torch.Tensor:
"""Fill current index K/scale rows into a slot-dense page buffer.
@@ -1801,6 +1936,8 @@ def fill_current_index_page_slots(
"""
current_locs = current_locs.reshape(-1)
if current_req_id is None:
current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
current_req_id = current_req_id.reshape(-1)
current_rows = int(current_locs.numel())
if current_rows == 0:
@@ -1817,6 +1954,7 @@ def fill_current_index_page_slots(
)
if tai_result is not None:
return tai_result
page_inverse = _slot_page_inverse_2d(page_inverse)
if dense_page_buffer.is_cuda:
_log_tai_materialize_fallback(
"index_current_fill_torch_reference_cuda",
@@ -3100,7 +3238,7 @@ def _logical_page_capacity_from_physical_page_capacity(
def build_slot_page_inverse(
slot_logical_pages: torch.Tensor,
logical_page_capacity: int,
batch_rows: int,
batch_rows: int = 1,
) -> torch.Tensor:
"""Build a PER-REQUEST logical_page -> slot_dense_page map.
@@ -3223,7 +3361,7 @@ def remap_logical_locs_to_slot_dense_locs(
logical_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
loc_req_id: torch.Tensor,
loc_req_id: torch.Tensor | None = None,
) -> torch.Tensor:
"""Map logical token locs through a PER-REQUEST slot page inverse.
@@ -3234,6 +3372,9 @@ def remap_logical_locs_to_slot_dense_locs(
"""
dense_locs = torch.full_like(logical_locs, -1)
page_inverse = _slot_page_inverse_2d(page_inverse)
if loc_req_id is None:
loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
if logical_locs.numel() == 0 or page_inverse.numel() == 0:
return dense_locs
@@ -3433,7 +3574,7 @@ def remap_logical_locs_to_shared_token_slot_dense_locs(
def remap_logical_pages_to_slot_dense_pages(
logical_pages: torch.Tensor,
page_inverse: torch.Tensor,
page_req_id: torch.Tensor,
page_req_id: torch.Tensor | None = None,
) -> torch.Tensor:
"""Map logical page ids through a PER-REQUEST slot page inverse.
@@ -3446,6 +3587,9 @@ def remap_logical_pages_to_slot_dense_pages(
"""
dense_pages_out = torch.full_like(logical_pages, -1)
page_inverse = _slot_page_inverse_2d(page_inverse)
if page_req_id is None:
page_req_id = build_page_table_row_req_id(logical_pages)
if logical_pages.numel() == 0 or page_inverse.numel() == 0:
return dense_pages_out
@@ -3665,16 +3809,23 @@ def build_shared_token_kv_slot_remap(
slot_dense_pages,
)
)
# ``dense_locs`` is never consumed downstream (the orchestrators remap their
# own per-request locs against ``page_inverse``). ``get_or_build_*`` always
# passes ``logical_locs=None``; keep it None to avoid needing a per-loc req id
# in the cached, loc-agnostic remap.
if logical_locs is not None:
raise NotImplementedError(
"build_shared_token_kv_slot_remap no longer materializes dense_locs; "
"remap per-request locs at the materialize call site instead."
)
# Runtime callers cache this remap loc-agnostically and remap their own locs
# with explicit per-request ids. Unit/reference callers may still ask for
# ``dense_locs``; keep that behavior request-scoped instead of using a
# batch-global logical-page inverse.
dense_locs = None
if logical_locs is not None:
loc_req_id = (
build_page_table_row_req_id(logical_locs)
if logical_locs.dim() >= 2
else torch.zeros_like(logical_locs, dtype=torch.long)
)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=page_inverse,
page_size=page_size,
loc_req_id=loc_req_id,
)
return SharedTokenKVSlotRemap(
slot_logical_pages=slot_logical_pages,
page_inverse=page_inverse,
@@ -4186,8 +4337,8 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
layout: CpSharedKVLayout,
page_size: int,
prefix_pages: int,
loc_req_id: torch.Tensor,
current_req_id: torch.Tensor,
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,
@@ -4205,6 +4356,20 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
"""
timing_start = cp_shared_kv_bs_gt1_timing_start()
if loc_req_id is None:
loc_req_id = _infer_req_id_for_locs_from_slot_remap(
logical_locs,
slot_remap,
page_size,
)
if current_req_id is None:
current_req_id = _infer_req_id_for_locs_from_slot_remap(
current_locs.reshape(-1),
slot_remap,
page_size,
strict=True,
)
total_slots = int(slot_remap.slot_logical_pages.numel())
if prefix_slot_spans is not None and prefix_slot_span is not None:
raise ValueError(
@@ -4380,7 +4545,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
page_size: int,
index_head_dim: int,
prefix_pages: int,
current_req_id: torch.Tensor,
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,
@@ -4390,6 +4555,14 @@ def materialize_prefix_and_reuse_current_index_page_slots(
"""Synchronously compose prefix index materialization with current index rows."""
timing_start = cp_shared_kv_bs_gt1_timing_start()
if current_req_id is None:
current_req_id = _infer_req_id_for_locs_from_slot_remap(
current_locs.reshape(-1),
slot_remap,
page_size,
strict=True,
)
total_slots = int(slot_remap.slot_logical_pages.numel())
if prefix_slot_spans is not None and prefix_slot_span is not None:
raise ValueError(
@@ -4905,9 +5078,10 @@ def materialize_shared_token_kv_buffer(
dense_kv_cache = None
if slot_remap is not None:
if loc_req_id is None:
raise ValueError(
"materialize_shared_token_kv_buffer slot_remap path requires "
"loc_req_id for the per-request page-inverse remap."
loc_req_id = _infer_req_id_for_locs_from_slot_remap(
logical_locs,
slot_remap,
page_size,
)
materialized_logical_pages = slot_remap.slot_logical_pages
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
@@ -4928,10 +5102,17 @@ def materialize_shared_token_kv_buffer(
use_slot_materialize = False
else:
if loc_req_id is None:
raise ValueError(
"materialize_shared_token_kv_buffer remap_logical_pages path "
"requires loc_req_id for the per-request page-inverse remap."
)
if logical_locs_row_ids is not None:
loc_req_id = logical_locs_row_ids
elif logical_locs.dim() >= 2:
loc_req_id = build_page_table_row_req_id(logical_locs)
elif remap_logical_pages.dim() < 2 or int(remap_logical_pages.shape[0]) <= 1:
loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long)
else:
raise ValueError(
"materialize_shared_token_kv_buffer remap_logical_pages path "
"requires loc_req_id for multi-request 1-D logical locs."
)
_debug_assert_no_negative_tensor_values(
remap_logical_pages,
context="CP shared KV token materialize page remap",

5
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -1459,12 +1459,11 @@ class Indexer(MultiPlatformOp):
else:
index_buffer = None
if cp_shared_kv_debug_enabled():
cp_layout = getattr(forward_batch, "cp_shared_kv_layout", None)
cp_shared_kv_debug_log(
"index_current_reuse",
"NSA index current reuse cp_rank=%s layer=%s kv_len=%s actual_seq_q=%s k_ck=%s s_ck=%s",
forward_batch.cp_shared_kv_layout.cp_rank
if forward_batch.cp_shared_kv_layout is not None
else None,
cp_layout.cp_rank if cp_layout is not None else None,
layer_id,
kv_len,
actual_seq_q,