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

Preserve request-slot identity for CP cache-hit reuse

Browse Source 
Warm-cache bs>1 requests can carry duplicate logical page ids across request rows. A batch-global page inverse collapses those request slots and can alias current/prefix KV across requests.

The runtime now carries compact row-scoped sorted page descriptors and remaps flattened logical locs with request row ids where needed, while retaining the legacy global inverse for no-row-context paths.

Constraint: Avoid the rejected dense [batch, logical_page_capacity] inverse because bs up to 10 and long contexts make that memory cost unacceptable
Rejected: Keep global page_inverse for bs>1 duplicate pages | it is lossy and matches the GSM8K warm-cache corruption shape
Rejected: Allocate page_inverse_by_row | correctness-safe but too much GPU memory for production
Confidence: medium
Scope-risk: moderate
Directive: Any future TAI materialize fast path for bs>1 duplicate pages must consume row-scoped descriptors or an equivalent request-row key
Tested: Local py_compile for touched runtime files; remote py_compile; remote pytest test_cp_shared_kv_runtime.py, test_nsa_cp_utils.py, test_cp_shared_kv_layout.py => 263 passed, 5 warnings, 2 subtests passed
Not-tested: Full GSM8K ETE after this cleanup pass
This commit is contained in:
laoyao0822
2026-06-08 20:55:58 +08:00
parent fb5ccaff26
commit b58513cba4
4 changed files with 912 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
View File

@@ -25,8 +25,8 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
get_or_build_shared_token_kv_slot_remap,
materialize_local_paged_buffer_page_slots_into,
materialize_local_token_kv_page_slots_into,
remap_logical_pages_to_slot_dense_pages,
remap_logical_locs_to_slot_dense_locs_optimized,
remap_logical_pages_to_shared_paged_slot_dense_pages,
remap_logical_locs_to_shared_token_slot_dense_locs,
slot_range_to_page_slice,
slot_range_to_token_slice,
)
@@ -295,6 +295,23 @@ class CpSharedKVIndexPrefetchHandle:
event: Optional[torch.cuda.Event] = None
@dataclass(frozen=True)
class _PrefetchTokenSlotRemapView:
slot_logical_pages: torch.Tensor
page_inverse: torch.Tensor
slot_sorted_logical_pages_by_row: torch.Tensor | None
slot_sorted_dense_pages_by_row: torch.Tensor | None
@dataclass(frozen=True)
class _PrefetchPagedSlotRemapView:
slot_logical_pages: torch.Tensor
page_inverse: torch.Tensor
dense_pages: torch.Tensor | None
slot_sorted_logical_pages_by_row: torch.Tensor | None
slot_sorted_dense_pages_by_row: torch.Tensor | None
class CpSharedKVMlaPrefetcher:
"""One-layer-ahead MLA prefix materialize prefetch for CP shared KV.
@@ -311,6 +328,8 @@ class CpSharedKVMlaPrefetcher:
prefix_pages: int,
slot_logical_pages: torch.Tensor,
page_inverse: torch.Tensor,
slot_sorted_logical_pages_by_row: torch.Tensor | None = None,
slot_sorted_dense_pages_by_row: torch.Tensor | None = None,
dense_num_pages: int,
owned_prefix_pages: int = -1,
owned_total_pages: int = -1,
@@ -321,6 +340,8 @@ class CpSharedKVMlaPrefetcher:
self.prefix_pages = prefix_pages
self.slot_logical_pages = slot_logical_pages
self.page_inverse = page_inverse
self.slot_sorted_logical_pages_by_row = slot_sorted_logical_pages_by_row
self.slot_sorted_dense_pages_by_row = slot_sorted_dense_pages_by_row
self.dense_num_pages = dense_num_pages
self.owned_prefix_pages = owned_prefix_pages
self.owned_total_pages = owned_total_pages
@@ -528,6 +549,8 @@ class CpSharedKVMlaPrefetcher:
prefix_pages=prefix_pages,
slot_logical_pages=remap.slot_logical_pages,
page_inverse=remap.page_inverse,
slot_sorted_logical_pages_by_row=remap.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=remap.slot_sorted_dense_pages_by_row,
dense_num_pages=remap.dense_num_pages,
owned_prefix_pages=owned_prefix_pages,
owned_total_pages=owned_total_pages,
@@ -657,9 +680,14 @@ class CpSharedKVMlaPrefetcher:
layout=self.layout,
physical_token_capacity=kv_cache.shape[0],
)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs,
page_inverse=self.page_inverse,
slot_remap=_PrefetchTokenSlotRemapView(
slot_logical_pages=self.slot_logical_pages,
page_inverse=self.page_inverse,
slot_sorted_logical_pages_by_row=self.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=self.slot_sorted_dense_pages_by_row,
),
page_size=self.page_size,
)
remap_ms = _cpu_timing_ms(remap_cpu)
@@ -757,9 +785,14 @@ class CpSharedKVMlaPrefetcher:
layout=self.layout,
physical_token_capacity=kv_cache.shape[0],
)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs,
page_inverse=self.page_inverse,
slot_remap=_PrefetchTokenSlotRemapView(
slot_logical_pages=self.slot_logical_pages,
page_inverse=self.page_inverse,
slot_sorted_logical_pages_by_row=self.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=self.slot_sorted_dense_pages_by_row,
),
page_size=self.page_size,
)
mixed_kv_cache, mixed_locs, _ = fill_current_kv_page_slots_and_remap_locs(
@@ -1078,6 +1111,8 @@ class CpSharedKVIndexPrefetcher:
prefix_pages: int,
slot_logical_pages: torch.Tensor,
page_inverse: torch.Tensor,
slot_sorted_logical_pages_by_row: torch.Tensor | None = None,
slot_sorted_dense_pages_by_row: torch.Tensor | None = None,
dense_num_pages: int,
owned_prefix_pages: int = -1,
owned_total_pages: int = -1,
@@ -1087,6 +1122,8 @@ class CpSharedKVIndexPrefetcher:
self.prefix_pages = prefix_pages
self.slot_logical_pages = slot_logical_pages
self.page_inverse = page_inverse
self.slot_sorted_logical_pages_by_row = slot_sorted_logical_pages_by_row
self.slot_sorted_dense_pages_by_row = slot_sorted_dense_pages_by_row
self.dense_num_pages = dense_num_pages
self.owned_prefix_pages = owned_prefix_pages
self.owned_total_pages = owned_total_pages
@@ -1332,6 +1369,8 @@ class CpSharedKVIndexPrefetcher:
prefix_pages=prefix_pages,
slot_logical_pages=remap.slot_logical_pages,
page_inverse=remap.page_inverse,
slot_sorted_logical_pages_by_row=remap.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=remap.slot_sorted_dense_pages_by_row,
dense_num_pages=remap.dense_num_pages,
owned_prefix_pages=owned_prefix_pages,
owned_total_pages=owned_total_pages,
@@ -1459,9 +1498,15 @@ class CpSharedKVIndexPrefetcher:
layout=self.layout,
physical_page_capacity=page_buffer.shape[0],
)
dense_pages = remap_logical_pages_to_slot_dense_pages(
dense_pages = remap_logical_pages_to_shared_paged_slot_dense_pages(
logical_pages,
page_inverse=self.page_inverse,
slot_remap=_PrefetchPagedSlotRemapView(
slot_logical_pages=self.slot_logical_pages,
page_inverse=self.page_inverse,
dense_pages=None,
slot_sorted_logical_pages_by_row=self.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=self.slot_sorted_dense_pages_by_row,
),
)
remap_ms = _cpu_timing_ms(remap_cpu)
total_ms = _cpu_timing_ms(consume_cpu)
@@ -1553,9 +1598,15 @@ class CpSharedKVIndexPrefetcher:
dense_page_buffer = handle.dense_page_buffer
remap_cpu = _cpu_timing_start()
dense_pages = remap_logical_pages_to_slot_dense_pages(
dense_pages = remap_logical_pages_to_shared_paged_slot_dense_pages(
logical_pages,
page_inverse=self.page_inverse,
slot_remap=_PrefetchPagedSlotRemapView(
slot_logical_pages=self.slot_logical_pages,
page_inverse=self.page_inverse,
dense_pages=None,
slot_sorted_logical_pages_by_row=self.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=self.slot_sorted_dense_pages_by_row,
),
)
dense_page_buffer = fill_current_index_page_slots(
dense_page_buffer=dense_page_buffer,

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

@@ -191,6 +191,9 @@ class SharedTokenKVSlotRemap:
dense_locs: torch.Tensor | None
logical_page_capacity: int
dense_num_pages: int
slot_dense_pages: torch.Tensor | None = None
slot_sorted_logical_pages_by_row: torch.Tensor | None = None
slot_sorted_dense_pages_by_row: torch.Tensor | None = None
@dataclass(frozen=True)
@@ -200,6 +203,8 @@ class SharedPagedBufferSlotRemap:
dense_pages: torch.Tensor
logical_page_capacity: int
dense_num_pages: int
slot_sorted_logical_pages_by_row: torch.Tensor | None = None
slot_sorted_dense_pages_by_row: torch.Tensor | None = None
def _tensor_identity_key(tensor: torch.Tensor) -> tuple[int, tuple[int, ...], str, str]:
@@ -2974,6 +2979,63 @@ def build_slot_page_inverse(
return page_inverse
def build_slot_page_sorted_index_by_row(
logical_pages: torch.Tensor,
dense_pages: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor] | tuple[None, None]:
"""Build compact row-scoped logical_page -> dense_page descriptors.
The global slot inverse is one-dimensional and cannot represent the same
logical page appearing in multiple request rows. A dense
[batch, logical_page_capacity] inverse fixes correctness but wastes GPU
memory. Keep only the actual batch page table, sorted per request row, so
row-aware remap can use searchsorted with O(actual_batch_pages) storage.
"""
if logical_pages.dim() != 2 or dense_pages.dim() != 2:
return None, None
if logical_pages.shape != dense_pages.shape:
raise ValueError(
"CP shared KV row-scoped page index expects matching logical/dense "
f"page shapes, got logical={tuple(logical_pages.shape)} "
f"dense={tuple(dense_pages.shape)}"
)
pages_long = logical_pages.to(torch.long)
dense_long = dense_pages.to(torch.long)
if pages_long.numel() == 0:
return pages_long, dense_long
valid_pages = pages_long > 0
invalid_sort_key = torch.full_like(
pages_long,
torch.iinfo(torch.long).max,
)
sort_keys = torch.where(valid_pages, pages_long, invalid_sort_key)
order = torch.argsort(sort_keys, dim=1)
sorted_pages = torch.gather(pages_long, 1, order)
sorted_dense_pages = torch.gather(dense_long, 1, order)
sorted_dense_pages = torch.where(
sorted_pages > 0,
sorted_dense_pages,
torch.full_like(sorted_dense_pages, -1),
)
return sorted_pages, sorted_dense_pages
def build_flattened_request_row_ids(
seq_lens_cpu: Any,
*,
device: torch.device,
) -> torch.Tensor:
seq_lens = _to_cpu_int_list(seq_lens_cpu, name="seq_lens_cpu")
if len(seq_lens) == 0:
return torch.empty((0,), device=device, dtype=torch.long)
lengths = torch.tensor(seq_lens, device=device, dtype=torch.long)
rows = torch.arange(len(seq_lens), device=device, dtype=torch.long)
return torch.repeat_interleave(rows, lengths)
def remap_logical_locs_to_slot_dense_locs(
logical_locs: torch.Tensor,
page_inverse: torch.Tensor,
@@ -3012,6 +3074,154 @@ def remap_logical_locs_to_slot_dense_locs(
return torch.where(mapped, dense_values, dense_locs)
def remap_logical_locs_to_slot_dense_locs_by_row(
logical_locs: torch.Tensor,
sorted_logical_pages_by_row: torch.Tensor | None,
sorted_dense_pages_by_row: torch.Tensor | None,
page_size: int,
*,
row_ids: torch.Tensor | None = None,
) -> torch.Tensor | None:
"""Map logical locs through a compact row-scoped slot descriptor."""
if (
sorted_logical_pages_by_row is None
or sorted_dense_pages_by_row is None
or sorted_logical_pages_by_row.dim() != 2
or sorted_dense_pages_by_row.dim() != 2
):
return None
if sorted_logical_pages_by_row.shape != sorted_dense_pages_by_row.shape:
raise ValueError(
"CP shared KV row-scoped loc remap got mismatched sorted row maps: "
f"pages={tuple(sorted_logical_pages_by_row.shape)} "
f"dense={tuple(sorted_dense_pages_by_row.shape)}"
)
if logical_locs.numel() == 0:
return torch.full_like(logical_locs, -1)
rows = int(sorted_logical_pages_by_row.shape[0])
if row_ids is None:
if logical_locs.dim() < 2:
return None
if int(logical_locs.shape[0]) != rows:
return None
dense_locs = torch.full_like(logical_locs, -1)
for row in range(rows):
dense_locs[row] = _remap_logical_locs_to_slot_dense_locs_one_row(
logical_locs[row],
sorted_logical_pages_by_row[row],
sorted_dense_pages_by_row[row],
page_size,
)
return dense_locs
else:
row_ids = row_ids.to(device=logical_locs.device, dtype=torch.long)
if row_ids.shape != logical_locs.shape:
if row_ids.numel() != logical_locs.numel():
raise ValueError(
"CP shared KV row-scoped loc remap got incompatible row_ids: "
f"logical_locs_shape={tuple(logical_locs.shape)} "
f"row_ids_shape={tuple(row_ids.shape)}"
)
row_ids = row_ids.reshape(logical_locs.shape)
dense_locs = torch.full_like(logical_locs, -1)
logical_locs_flat = logical_locs.reshape(-1)
row_ids_flat = row_ids.reshape(-1)
dense_locs_flat = dense_locs.reshape(-1)
for row in range(rows):
row_mask = row_ids_flat == row
dense_locs_flat[row_mask] = _remap_logical_locs_to_slot_dense_locs_one_row(
logical_locs_flat[row_mask],
sorted_logical_pages_by_row[row],
sorted_dense_pages_by_row[row],
page_size,
)
return dense_locs
def _remap_logical_locs_to_slot_dense_locs_one_row(
logical_locs: torch.Tensor,
sorted_logical_pages: torch.Tensor,
sorted_dense_pages: torch.Tensor,
page_size: int,
) -> torch.Tensor:
dense_locs = torch.full_like(logical_locs, -1)
if logical_locs.numel() == 0:
return dense_locs
locs_long = logical_locs.to(torch.long)
valid_locs = locs_long >= 0
safe_locs = torch.where(valid_locs, locs_long, torch.zeros_like(locs_long))
logical_pages = torch.div(safe_locs, page_size, rounding_mode="floor")
offsets = torch.remainder(safe_locs, page_size)
dummy = valid_locs & (logical_pages == 0)
dense_locs = torch.where(dummy, offsets.to(logical_locs.dtype), dense_locs)
positive = valid_locs & (logical_pages > 0)
flat_positive = positive.reshape(-1)
positive_indices = torch.nonzero(flat_positive, as_tuple=False).reshape(-1)
if positive_indices.numel() == 0 or sorted_logical_pages.numel() == 0:
return dense_locs
pages_flat = logical_pages.reshape(-1)
offsets_flat = offsets.reshape(-1)
out_flat = dense_locs.reshape(-1)
pages_to_map = pages_flat[positive_indices]
positions = torch.searchsorted(sorted_logical_pages, pages_to_map)
in_range = positions < int(sorted_logical_pages.numel())
safe_positions = torch.clamp(
positions,
min=0,
max=max(int(sorted_logical_pages.numel()) - 1, 0),
)
found_pages = sorted_logical_pages[safe_positions]
dense_pages = sorted_dense_pages[safe_positions]
mapped = in_range & (found_pages == pages_to_map) & (dense_pages >= 0)
mapped_indices = positive_indices[mapped]
if mapped_indices.numel() > 0:
dense_values = (
dense_pages[mapped].to(logical_locs.dtype) * page_size
+ offsets_flat[mapped_indices].to(logical_locs.dtype)
)
out_flat[mapped_indices] = dense_values
if cp_shared_kv_debug_enabled() and bool(torch.any(~mapped).item()):
missing_pages = pages_to_map[~mapped]
raise RuntimeError(
"CP shared KV row slot remap got logical locs outside row map. "
f"missing_page_min={int(missing_pages.min().item())} "
f"missing_page_max={int(missing_pages.max().item())} "
f"logical_locs={tensor_debug_summary(logical_locs)}"
)
return dense_locs
def remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs: torch.Tensor,
slot_remap: SharedTokenKVSlotRemap,
page_size: int,
*,
row_ids: torch.Tensor | None = None,
) -> torch.Tensor:
row_scoped = remap_logical_locs_to_slot_dense_locs_by_row(
logical_locs,
slot_remap.slot_sorted_logical_pages_by_row,
slot_remap.slot_sorted_dense_pages_by_row,
page_size,
row_ids=row_ids,
)
if row_scoped is not None:
return row_scoped
return remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=slot_remap.page_inverse,
page_size=page_size,
)
def remap_logical_pages_to_slot_dense_pages(
logical_pages: torch.Tensor,
page_inverse: torch.Tensor,
@@ -3049,6 +3259,139 @@ def remap_logical_pages_to_slot_dense_pages(
return torch.where(mapped, dense_pages.to(logical_pages.dtype), dense_pages_out)
def remap_logical_pages_to_slot_dense_pages_by_row(
logical_pages: torch.Tensor,
sorted_logical_pages_by_row: torch.Tensor | None,
sorted_dense_pages_by_row: torch.Tensor | None,
*,
row_ids: torch.Tensor | None = None,
) -> torch.Tensor | None:
if (
sorted_logical_pages_by_row is None
or sorted_dense_pages_by_row is None
or sorted_logical_pages_by_row.dim() != 2
or sorted_dense_pages_by_row.dim() != 2
):
return None
if sorted_logical_pages_by_row.shape != sorted_dense_pages_by_row.shape:
raise ValueError(
"CP shared KV row-scoped page remap got mismatched sorted row maps: "
f"pages={tuple(sorted_logical_pages_by_row.shape)} "
f"dense={tuple(sorted_dense_pages_by_row.shape)}"
)
if logical_pages.numel() == 0:
return torch.full_like(logical_pages, -1)
rows = int(sorted_logical_pages_by_row.shape[0])
if row_ids is None:
if logical_pages.dim() < 2:
return None
if int(logical_pages.shape[0]) != rows:
return None
dense_pages = torch.full_like(logical_pages, -1)
for row in range(rows):
dense_pages[row] = _remap_logical_pages_to_slot_dense_pages_one_row(
logical_pages[row],
sorted_logical_pages_by_row[row],
sorted_dense_pages_by_row[row],
)
return dense_pages
else:
row_ids = row_ids.to(device=logical_pages.device, dtype=torch.long)
if row_ids.shape != logical_pages.shape:
if row_ids.numel() != logical_pages.numel():
raise ValueError(
"CP shared KV row-scoped page remap got incompatible row_ids: "
f"logical_pages_shape={tuple(logical_pages.shape)} "
f"row_ids_shape={tuple(row_ids.shape)}"
)
row_ids = row_ids.reshape(logical_pages.shape)
dense_pages = torch.full_like(logical_pages, -1)
logical_pages_flat = logical_pages.reshape(-1)
row_ids_flat = row_ids.reshape(-1)
dense_pages_flat = dense_pages.reshape(-1)
for row in range(rows):
row_mask = row_ids_flat == row
dense_pages_flat[row_mask] = _remap_logical_pages_to_slot_dense_pages_one_row(
logical_pages_flat[row_mask],
sorted_logical_pages_by_row[row],
sorted_dense_pages_by_row[row],
)
return dense_pages
def _remap_logical_pages_to_slot_dense_pages_one_row(
logical_pages: torch.Tensor,
sorted_logical_pages: torch.Tensor,
sorted_dense_pages: torch.Tensor,
) -> torch.Tensor:
dense_pages_out = torch.full_like(logical_pages, -1)
if logical_pages.numel() == 0:
return dense_pages_out
pages_long = logical_pages.to(torch.long)
valid_pages = pages_long >= 0
dummy = valid_pages & (pages_long == 0)
dense_pages_out = torch.where(
dummy,
torch.zeros_like(dense_pages_out),
dense_pages_out,
)
positive = valid_pages & (pages_long > 0)
positive_indices = torch.nonzero(positive.reshape(-1), as_tuple=False).reshape(-1)
if positive_indices.numel() == 0 or sorted_logical_pages.numel() == 0:
return dense_pages_out
pages_flat = pages_long.reshape(-1)
out_flat = dense_pages_out.reshape(-1)
pages_to_map = pages_flat[positive_indices]
positions = torch.searchsorted(sorted_logical_pages, pages_to_map)
in_range = positions < int(sorted_logical_pages.numel())
safe_positions = torch.clamp(
positions,
min=0,
max=max(int(sorted_logical_pages.numel()) - 1, 0),
)
found_pages = sorted_logical_pages[safe_positions]
dense_pages = sorted_dense_pages[safe_positions]
mapped = in_range & (found_pages == pages_to_map) & (dense_pages >= 0)
mapped_indices = positive_indices[mapped]
if mapped_indices.numel() > 0:
out_flat[mapped_indices] = dense_pages[mapped].to(logical_pages.dtype)
if cp_shared_kv_debug_enabled() and bool(torch.any(~mapped).item()):
missing_pages = pages_to_map[~mapped]
raise RuntimeError(
"CP shared KV row slot remap got logical pages outside row map. "
f"missing_page_min={int(missing_pages.min().item())} "
f"missing_page_max={int(missing_pages.max().item())} "
f"logical_pages={tensor_debug_summary(logical_pages)}"
)
return dense_pages_out
def remap_logical_pages_to_shared_paged_slot_dense_pages(
logical_pages: torch.Tensor,
slot_remap: SharedPagedBufferSlotRemap,
*,
row_ids: torch.Tensor | None = None,
) -> torch.Tensor:
row_scoped = remap_logical_pages_to_slot_dense_pages_by_row(
logical_pages,
slot_remap.slot_sorted_logical_pages_by_row,
slot_remap.slot_sorted_dense_pages_by_row,
row_ids=row_ids,
)
if row_scoped is not None:
return row_scoped
return remap_logical_pages_to_slot_dense_pages(
logical_pages,
page_inverse=slot_remap.page_inverse,
)
def build_shared_token_kv_slot_remap(
kv_cache: torch.Tensor,
logical_locs: torch.Tensor | None,
@@ -3079,26 +3422,42 @@ def build_shared_token_kv_slot_remap(
kv_cache.shape[0] // page_size,
layout,
)
slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
slot_logical_pages, slot_dense_pages = build_slot_page_remap(remap_logical_pages)
page_inverse = build_slot_page_inverse_optimized(
slot_logical_pages,
logical_page_capacity=logical_page_capacity,
)
sorted_logical_pages_by_row, sorted_dense_pages_by_row = (
build_slot_page_sorted_index_by_row(
remap_logical_pages,
slot_dense_pages,
)
)
dense_locs = (
remap_logical_locs_to_slot_dense_locs_optimized(
remap_logical_locs_to_slot_dense_locs_by_row(
logical_locs,
page_inverse=page_inverse,
sorted_logical_pages_by_row,
sorted_dense_pages_by_row,
page_size=page_size,
)
if logical_locs is not None
else None
)
if dense_locs is None and logical_locs is not None:
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=page_inverse,
page_size=page_size,
)
return SharedTokenKVSlotRemap(
slot_logical_pages=slot_logical_pages,
page_inverse=page_inverse,
dense_locs=dense_locs,
logical_page_capacity=logical_page_capacity,
dense_num_pages=int(slot_logical_pages.numel()) + 1,
slot_dense_pages=slot_dense_pages,
slot_sorted_logical_pages_by_row=sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=sorted_dense_pages_by_row,
)
@@ -3125,12 +3484,20 @@ def build_shared_paged_buffer_slot_remap(
slot_logical_pages,
logical_page_capacity=logical_page_capacity,
)
sorted_logical_pages_by_row, sorted_dense_pages_by_row = (
build_slot_page_sorted_index_by_row(
logical_pages,
dense_pages,
)
)
return SharedPagedBufferSlotRemap(
slot_logical_pages=slot_logical_pages,
page_inverse=page_inverse,
dense_pages=dense_pages,
logical_page_capacity=logical_page_capacity,
dense_num_pages=int(slot_logical_pages.numel()) + 1,
slot_sorted_logical_pages_by_row=sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=sorted_dense_pages_by_row,
)
@@ -3594,6 +3961,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
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,
layer_id: int | None = None,
nvtx_source: str = "mla.partial_current_sync",
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -3699,10 +4067,11 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
layout=layout,
physical_token_capacity=kv_cache.shape[0],
)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs,
page_inverse=slot_remap.page_inverse,
page_size=page_size,
slot_remap=slot_remap,
row_ids=logical_locs_row_ids,
)
mixed_kv_cache, mixed_locs, _ = fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_kv_cache,
@@ -4262,6 +4631,7 @@ def materialize_shared_token_kv_buffer(
remap_logical_locs: torch.Tensor | None = None,
remap_logical_pages: torch.Tensor | None = None,
slot_remap: SharedTokenKVSlotRemap | None = None,
logical_locs_row_ids: torch.Tensor | None = None,
nvtx_source: str = "mla.full_materialize",
nvtx_layer_id: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -4294,10 +4664,11 @@ def materialize_shared_token_kv_buffer(
dense_kv_cache = None
if slot_remap is not None:
materialized_logical_pages = slot_remap.slot_logical_pages
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs(
logical_locs,
page_inverse=slot_remap.page_inverse,
page_size=page_size,
slot_remap=slot_remap,
row_ids=logical_locs_row_ids,
)
use_slot_materialize = True
elif remap_logical_pages is None:
@@ -4325,7 +4696,11 @@ def materialize_shared_token_kv_buffer(
layout,
)
tai_result = None
if _tai_materialize_runtime_enabled():
row_scoped_remap_required = (
logical_locs_row_ids is not None
or (remap_logical_pages.dim() == 2 and logical_locs.dim() >= 2)
)
if _tai_materialize_runtime_enabled() and not row_scoped_remap_required:
materialized_logical_pages = remap_logical_pages.reshape(-1)
tai_result = _try_tai_materialize_token_kv_pages_and_locs(
kv_cache=kv_cache,
@@ -4336,16 +4711,32 @@ def materialize_shared_token_kv_buffer(
page_size=page_size,
)
if tai_result is None:
materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
materialized_logical_pages, slot_dense_pages = build_slot_page_remap(
remap_logical_pages
)
page_inverse = build_slot_page_inverse_optimized(
materialized_logical_pages,
logical_page_capacity=logical_page_capacity,
)
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=page_inverse,
page_size=page_size,
sorted_logical_pages_by_row, sorted_dense_pages_by_row = (
build_slot_page_sorted_index_by_row(
remap_logical_pages,
slot_dense_pages,
)
)
dense_locs = remap_logical_locs_to_slot_dense_locs_by_row(
logical_locs,
sorted_logical_pages_by_row,
sorted_dense_pages_by_row,
page_size=page_size,
row_ids=logical_locs_row_ids,
)
if dense_locs is None:
dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
logical_locs,
page_inverse=page_inverse,
page_size=page_size,
)
use_slot_materialize = True
else:
dense_kv_cache, dense_locs = tai_result

16
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -16,6 +16,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
)
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_batch_current_slot_spans,
build_flattened_request_row_ids,
build_batch_prefix_slot_spans,
build_current_loc_remap,
cp_shared_kv_debug_enabled,
@@ -2522,6 +2523,20 @@ class NativeSparseAttnBackend(
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
logical_locs_row_ids = build_flattened_request_row_ids(
metadata.indexer_seq_lens_cpu,
device=page_table_1_flattened.device,
)
if int(logical_locs_row_ids.numel()) != int(
page_table_1_flattened.numel()
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"row-id metadata does not match flattened page table. "
f"row_ids={int(logical_locs_row_ids.numel())} "
f"page_table_1_flattened={int(page_table_1_flattened.numel())} "
f"indexer_seq_lens_cpu={metadata.indexer_seq_lens_cpu}"
)
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
@@ -2541,6 +2556,7 @@ class NativeSparseAttnBackend(
prefix_pages=prefix_pages,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
logical_locs_row_ids=logical_locs_row_ids,
layer_id=layer.layer_id,
nvtx_source="mla.ragged_partial_current_sync",
)

431
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
View File

@@ -1374,6 +1374,246 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
)
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
def test_tai_current_slot_fill_matches_torch_for_parallel20_tiny_fp8_rows(
self,
):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
kernels = runtime._load_tai_materialize_kernels()
if kernels is None or not hasattr(
kernels, "fill_current_token_kv_page_slots_and_remap_locs"
):
self.skipTest("TAI current-slot fill kernel is not installed")
device = torch.device("cuda", torch.cuda.current_device())
page_size = 64
prefix_lens = [704, 704, 640, 704, 704]
extend_lens = [29, 9, 60, 9, 17]
pages_per_request = max(
(prefix + extend + page_size - 1) // page_size
for prefix, extend in zip(prefix_lens, extend_lens)
)
rows = []
logical_locs = []
current_locs = []
next_page = 1
for prefix_len, extend_len in zip(prefix_lens, extend_lens):
prefix_pages = prefix_len // page_size
current_pages = (prefix_len + extend_len + page_size - 1) // page_size
req_pages = []
for _ in range(current_pages):
req_pages.append(next_page)
next_page += 1
req_pages.extend([0] * (pages_per_request - len(req_pages)))
rows.append(req_pages)
for page_idx in range(prefix_pages):
page = req_pages[page_idx]
logical_locs.extend(
page * page_size + offset for offset in range(page_size)
)
cur_page = req_pages[prefix_pages]
cur = [cur_page * page_size + offset for offset in range(extend_len)]
logical_locs.extend(cur)
current_locs.extend(cur)
remap_logical_pages = torch.tensor(rows, device=device, dtype=torch.long)
slot_logical_pages, _ = runtime.build_slot_page_remap(remap_logical_pages)
page_inverse = runtime.build_slot_page_inverse(
slot_logical_pages, logical_page_capacity=next_page * page_size
)
logical_locs_t = torch.tensor(logical_locs, device=device, dtype=torch.long)
current_locs_t = torch.tensor(current_locs, device=device, dtype=torch.long)
materialized_locs = runtime.remap_logical_locs_to_slot_dense_locs(
logical_locs_t, page_inverse=page_inverse, page_size=page_size
)
dense_rows = (int(slot_logical_pages.numel()) + 1) * page_size
# Use production FP8 MLA row width. uint8 makes byte-copy exactness visible
# without depending on fp8 value quantization in the unit test.
row_width = 656
base_current = torch.arange(
int(current_locs_t.numel()) * row_width,
device=device,
dtype=torch.int64,
)
current_kv = (base_current % 251).to(torch.uint8).view(
int(current_locs_t.numel()), 1, row_width
)
dense_ref = torch.zeros((dense_rows, 1, row_width), device=device, dtype=torch.uint8)
dense_tai = torch.zeros_like(dense_ref)
with patch.object(runtime, "_tai_materialize_runtime_enabled", return_value=False):
ref_kv, ref_locs, ref_mask = runtime.fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_ref,
materialized_dense_locs=materialized_locs,
current_kv_cache=current_kv,
logical_locs=logical_locs_t,
current_locs=current_locs_t,
page_inverse=page_inverse,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
runtime._tai_current_slot_fill_sparse_pages_self_test.cache_clear()
with envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.override(True), patch.object(
runtime,
"_tai_current_slot_fill_supports_sparse_pages",
return_value=True,
):
tai_kv, tai_locs, tai_mask = runtime.fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_tai,
materialized_dense_locs=materialized_locs,
current_kv_cache=current_kv,
logical_locs=logical_locs_t,
current_locs=current_locs_t,
page_inverse=page_inverse,
page_size=page_size,
mask_non_current_in_current_pages=True,
)
self.assertTrue(torch.equal(tai_locs, ref_locs))
self.assertTrue(torch.equal(tai_mask, ref_mask))
self.assertTrue(torch.equal(tai_kv, ref_kv))
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
def test_fp8_paged_dequant_matches_reference_for_parallel20_slot_locs(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.layers.attention.nsa.dequant_k_cache import (
dequantize_k_cache_paged,
)
device = torch.device("cuda", torch.cuda.current_device())
page_size = 64
prefix_lens = [704, 704, 640, 704, 704]
extend_lens = [29, 9, 60, 9, 17]
pages_per_request = max(
(prefix + extend + page_size - 1) // page_size
for prefix, extend in zip(prefix_lens, extend_lens)
)
rows = []
logical_locs = []
next_page = 1
for prefix_len, extend_len in zip(prefix_lens, extend_lens):
seq_len = prefix_len + extend_len
req_pages = []
for _ in range((seq_len + page_size - 1) // page_size):
req_pages.append(next_page)
next_page += 1
req_pages.extend([0] * (pages_per_request - len(req_pages)))
rows.append(req_pages)
for token_idx in range(seq_len):
page = req_pages[token_idx // page_size]
logical_locs.append(page * page_size + token_idx % page_size)
remap_logical_pages = torch.tensor(rows, device=device, dtype=torch.long)
slot_logical_pages, _ = runtime.build_slot_page_remap(remap_logical_pages)
page_inverse = runtime.build_slot_page_inverse(
slot_logical_pages, logical_page_capacity=next_page * page_size
)
logical_locs_t = torch.tensor(logical_locs, device=device, dtype=torch.long)
mixed_locs = runtime.remap_logical_locs_to_slot_dense_locs(
logical_locs_t, page_inverse=page_inverse, page_size=page_size
)
self.assertFalse(bool(torch.any(mixed_locs < 0).item()))
dense_rows = (int(slot_logical_pages.numel()) + 1) * page_size
packed_u8 = torch.empty((dense_rows, 656), device=device, dtype=torch.uint8)
packed_fp8 = packed_u8.view(torch.float8_e4m3fn)
nope = packed_fp8[:, :512]
scales = packed_fp8[:, 512:528].view(torch.float32)
rope = packed_fp8[:, 528:].view(torch.bfloat16)
token_ids = torch.arange(dense_rows, device=device, dtype=torch.float32).view(-1, 1)
dim_ids = torch.arange(512, device=device, dtype=torch.float32).view(1, -1)
nope.copy_(((token_ids % 7) + (dim_ids % 13) / 16.0).to(torch.float8_e4m3fn))
scales.copy_(torch.ones_like(scales))
rope_token_ids = torch.arange(dense_rows, device=device, dtype=torch.float32).view(-1, 1)
rope_dim_ids = torch.arange(64, device=device, dtype=torch.float32).view(1, -1)
rope.copy_((rope_token_ids * 0.25 + rope_dim_ids / 32.0).to(torch.bfloat16))
out = dequantize_k_cache_paged(
packed_fp8.view(dense_rows, 1, 656),
mixed_locs.to(torch.int64),
).view(-1, 576)
selected = packed_fp8.view(dense_rows, 656)[mixed_locs.to(torch.long)]
selected_nope = selected[:, :512].to(torch.float32)
selected_scales = selected[:, 512:528].view(torch.float32)
selected_rope = selected[:, 528:].view(torch.bfloat16)
expected = torch.empty((int(mixed_locs.numel()), 576), device=device, dtype=torch.bfloat16)
for tile in range(4):
expected[:, tile * 128 : (tile + 1) * 128] = (
selected_nope[:, tile * 128 : (tile + 1) * 128]
* selected_scales[:, tile : tile + 1]
).to(torch.bfloat16)
expected[:, 512:] = selected_rope
self.assertTrue(torch.equal(out, expected))
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
def test_flashmla_sparse_valid_rows_match_compact_with_compute_padding(self):
try:
from sgl_kernel.flash_mla import flash_mla_sparse_fwd
except Exception as exc:
self.skipTest(f"flash_mla_sparse_fwd is unavailable: {exc}")
device = torch.device("cuda", torch.cuda.current_device())
torch.manual_seed(1234)
num_heads = 64
head_dim = 576
value_dim = 512
topk = 2048
valid_lens = [29, 9, 60, 9, 17]
compute_rows_per_req = 64
output_rows = len(valid_lens) * compute_rows_per_req
kv_lens = [704 + x for x in valid_lens]
kv_total = sum(kv_lens)
q_padded = torch.randn(
(output_rows, num_heads, head_dim),
device=device,
dtype=torch.bfloat16,
) / 10
kv = torch.randn((kv_total, 1, head_dim), device=device, dtype=torch.bfloat16) / 10
indices_padded = torch.full(
(output_rows, 1, topk), -1, device=device, dtype=torch.int32
)
valid_rows = []
kv_base = 0
out_base = 0
for valid_len, kv_len in zip(valid_lens, kv_lens):
for row in range(valid_len):
out_row = out_base + row
valid_rows.append(out_row)
indices_padded[out_row, 0, :kv_len] = torch.arange(
kv_base, kv_base + kv_len, device=device, dtype=torch.int32
)
kv_base += kv_len
out_base += compute_rows_per_req
valid_rows_t = torch.tensor(valid_rows, device=device, dtype=torch.long)
sm_scale = 1.0 / (head_dim**0.5)
padded_out, _, _ = flash_mla_sparse_fwd(
q_padded, kv, indices_padded, sm_scale=sm_scale, d_v=value_dim
)
compact_out, _, _ = flash_mla_sparse_fwd(
q_padded[valid_rows_t].contiguous(),
kv,
indices_padded[valid_rows_t].contiguous(),
sm_scale=sm_scale,
d_v=value_dim,
)
torch.testing.assert_close(
padded_out[valid_rows_t].float(),
compact_out.float(),
atol=8e-4,
rtol=2.01 / 128,
)
def test_materialize_prefix_and_reuse_current_kv_page_slots_without_prefetcher(
self,
):
@@ -2547,6 +2787,66 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertEqual(remap.dense_num_pages, 7)
self.assertEqual(remap.dense_locs.tolist(), [[4, 8, -1], [16, 0, 20]])
def test_token_slot_remap_duplicate_pages_keep_request_slot_identity(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_shared_token_kv_slot_remap,
)
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.zeros((64, 1, 1), dtype=torch.float32)
# Two requests share logical page 1 as prefix. The slot layout gives
# the two occurrences distinct dense slots: req0 slot 0 -> dense page 1,
# req1 slot 0 -> dense page 3. A request-packed dense buffer must keep
# those request slots addressable independently.
remap_logical_pages = torch.tensor(
[
[1, 2],
[1, 3],
],
dtype=torch.int64,
)
logical_locs = torch.tensor(
[
[1 * page_size, 2 * page_size],
[1 * page_size, 3 * page_size],
],
dtype=torch.int64,
)
remap = build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
self.assertEqual(remap.slot_logical_pages.tolist(), [1, 2, 1, 3])
self.assertEqual(remap.dense_num_pages, 5)
self.assertEqual(
tuple(remap.slot_sorted_logical_pages_by_row.shape),
tuple(remap_logical_pages.shape),
)
self.assertEqual(
tuple(remap.slot_sorted_dense_pages_by_row.shape),
tuple(remap_logical_pages.shape),
)
self.assertLess(
remap.slot_sorted_logical_pages_by_row.numel(),
remap.logical_page_capacity * remap_logical_pages.shape[0],
"row-aware slot remap must stay compact and avoid dense "
"[batch, logical_page_capacity] inverse allocation",
)
self.assertEqual(
remap.dense_locs.tolist(),
[
[1 * page_size, 2 * page_size],
[3 * page_size, 4 * page_size],
],
)
def test_forward_batch_token_slot_remap_is_cached_across_layers(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
@@ -4086,6 +4386,137 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
"remote-rank valid current loc must remain visible after current slot all-reduce",
)
def test_cp8_batch_kv_partial_current_keeps_request_packed_layout(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
cp_size = 8
prefix_lens = [8, 8, 4, 8]
extend_lens = [2, 3, 1, 2]
# Current pages are deliberately rank0-owned while prefix pages are
# spread across other owners. This matches the tiny warm-cache shape
# where all valid current rows live on CP0 but the prefix is assembled
# from all CP ranks.
logical_pages = torch.tensor(
[
[1, 2, 9],
[3, 4, 17],
[5, 25, 0],
[6, 7, 33],
],
dtype=torch.int64,
)
logical_locs_by_req = []
expected_rows = []
current_locs = []
current_rows = []
for req_id, (prefix_len, extend_len) in enumerate(
zip(prefix_lens, extend_lens, strict=True)
):
locs = []
values = []
for token_idx in range(prefix_len + extend_len):
page_idx = token_idx // page_size
token_off = token_idx % page_size
logical_page = int(logical_pages[req_id, page_idx].item())
loc = logical_page * page_size + token_off
value = float(req_id * 1000 + token_idx)
locs.append(loc)
values.append([value])
if token_idx >= prefix_len:
current_locs.append(loc)
current_rows.append([value])
logical_locs_by_req.append(locs)
expected_rows.extend(values)
max_seq_len = max(len(x) for x in logical_locs_by_req)
logical_locs = torch.full(
(len(logical_locs_by_req), max_seq_len),
-1,
dtype=torch.int64,
)
for req_id, locs in enumerate(logical_locs_by_req):
logical_locs[req_id, : len(locs)] = torch.tensor(locs, dtype=torch.int64)
current_locs_tensor = torch.tensor(current_locs, dtype=torch.int64)
current_kv = torch.tensor(current_rows, dtype=torch.float32)
expected = torch.tensor(expected_rows, dtype=torch.float32)
prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens,
page_size=page_size,
)
current_slot_spans = runtime.build_batch_current_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens,
extend_lens_cpu=extend_lens,
page_size=page_size,
)
rank_outputs = []
mixed_locs_ref = None
physical_tokens = 64 * page_size
for rank in range(cp_size):
layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)
kv_cache = torch.zeros((physical_tokens, 1), dtype=torch.float32)
for req_id, locs in enumerate(logical_locs_by_req):
for token_idx, loc in enumerate(locs):
if token_idx >= prefix_lens[req_id]:
continue
loc_tensor = torch.tensor([loc], dtype=torch.int64)
if bool(layout.owned_by_this_rank(loc_tensor)[0].item()):
phys = int(layout.logical_locs_to_physical(loc_tensor)[0].item())
kv_cache[phys] = float(req_id * 1000 + token_idx)
rank_current_locs = current_locs_tensor[
layout.owned_by_this_rank(current_locs_tensor)
]
if rank == 0:
# All selected current pages are rank0-owned by construction.
rank_current_kv = current_kv
self.assertEqual(rank_current_locs.tolist(), current_locs)
else:
rank_current_kv = torch.empty((0, 1), dtype=torch.float32)
self.assertEqual(rank_current_locs.numel(), 0)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=logical_pages,
layout=layout,
page_size=page_size,
)
with patch.object(
runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
):
mixed_kv, mixed_locs = runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_kv_cache=rank_current_kv,
current_locs=rank_current_locs,
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,
)
if mixed_locs_ref is None:
mixed_locs_ref = mixed_locs
else:
self.assertTrue(torch.equal(mixed_locs_ref, mixed_locs))
rank_outputs.append(mixed_kv)
assert mixed_locs_ref is not None
merged_kv = torch.stack(rank_outputs, dim=0).sum(dim=0)
packed_locs = mixed_locs_ref[mixed_locs_ref >= 0].to(torch.long)
actual = merged_kv[packed_locs]
self.assertTrue(torch.equal(actual, expected))
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
def test_tai_batched_index_mqa_prepare_matches_getk_gets_reference_gsm8k_bs5(self):
from types import SimpleNamespace