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

Reduce shared KV materialize synchronization

Browse Source 
The shared-KV materialize path was spending time in Python-observed CUDA tensor predicates and dynamic-shape remap helpers. Keep the runtime changes that move the hot paged path to slot-based device remapping, while dropping the NVTX experiment from this commit so profiling annotations do not become part of the runtime surface yet.\n\nThe MLA read path now passes the real page table as the page remap source, which keeps paged topk indices tied to the same logical page-table domain used to build the dense materialized KV view.\n\nConstraint: CP shared KV still needs a dense per-call view before deeper Phase4/Phase5 layout changes remove the materialize cost.\nRejected: Keep NVTX ranges in this commit | user requested reverting NVTX instrumentation before commit\nRejected: Restore compact unique-page remap everywhere | it reintroduces CUDA sync-prone dynamic-shape ops on the hot paged materialize path\nConfidence: medium\nScope-risk: moderate\nDirective: Benchmark slot-remap buffer size against compact unique-page remap before treating this as the final performance path; Phase4/5 should reduce materialize instead of relying on this aggregation path.\nTested: git diff --check on changed files; python -m py_compile on changed runtime/backend/test files; grep confirmed NVTX symbols removed\nNot-tested: pytest blocked locally by missing pybase64 dependency; multi-node PD runtime not rerun
This commit is contained in:
laoyao0822
2026-04-28 04:09:19 +08:00
parent d31535589b
commit 47bd2fdf1f
3 changed files with 393 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -211,13 +211,10 @@ def build_dense_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, t
"""Build a dense per-call page remap for shared-KV runtime materialization."""
dense_pages = logical_pages.clone()
positive_mask = logical_pages > 0
if not torch.any(positive_mask):
empty = logical_pages.new_empty((0,))
return empty, dense_pages
unique_logical_pages = torch.unique(logical_pages[positive_mask], sorted=True)
positive_pages = logical_pages[positive_mask]
unique_logical_pages = torch.unique(positive_pages, sorted=True)
dense_pages[positive_mask] = remap_logical_pages_to_dense_pages(
logical_pages[positive_mask],
positive_pages,
unique_logical_pages=unique_logical_pages,
)
return unique_logical_pages, dense_pages
@@ -229,15 +226,11 @@ def remap_logical_pages_to_dense_pages(
) -> torch.Tensor:
dense_pages = logical_pages.clone()
positive_mask = logical_pages > 0
if not torch.any(positive_mask):
return dense_pages
if unique_logical_pages.numel() == 0:
raise ValueError("unique_logical_pages is empty but logical_pages contains data")
positive_pages = logical_pages[positive_mask]
insert_positions = torch.searchsorted(unique_logical_pages, positive_pages)
if insert_positions.numel() > 0:
if cp_shared_kv_debug_enabled() and insert_positions.numel() > 0:
if unique_logical_pages.numel() == 0:
raise ValueError("unique_logical_pages is empty but logical_pages contains data")
if torch.any(insert_positions >= unique_logical_pages.numel()):
raise ValueError("logical_pages contains entries outside unique_logical_pages")
if not torch.equal(unique_logical_pages[insert_positions], positive_pages):
@@ -247,6 +240,117 @@ def remap_logical_pages_to_dense_pages(
return dense_pages
def build_slot_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Build a fixed-shape page remap without dynamic CUDA output ops.
The compact remap path uses boolean compaction + unique/searchsorted. Those
ops have data-dependent output shapes and can cause device-to-host
synchronization in CUDA. The slot remap instead gives every input page-table
slot a deterministic dense page id (`flat_slot + 1`) and preserves 0/-1
sentinels. It can materialize duplicate logical pages into duplicate dense
slots, trading extra device work for avoiding CPU synchronization.
"""
slot_logical_pages = logical_pages.reshape(-1)
dense_pages_flat = slot_logical_pages.clone()
if slot_logical_pages.numel() == 0:
return slot_logical_pages, dense_pages_flat.reshape(logical_pages.shape)
slot_ids = torch.arange(
1,
slot_logical_pages.numel() + 1,
device=logical_pages.device,
dtype=logical_pages.dtype,
)
dense_pages_flat = torch.where(
slot_logical_pages > 0,
slot_ids,
dense_pages_flat,
)
return slot_logical_pages, dense_pages_flat.reshape(logical_pages.shape)
def _logical_page_capacity_from_physical_page_capacity(
physical_page_capacity: int,
layout: CpSharedKVLayout,
) -> int:
# Physical page 0 is the shared dummy page. Real physical pages
# 1..N correspond to N * cp_size logical pages across the CP group.
return max(physical_page_capacity - 1, 0) * layout.cp_size + 1
def build_slot_page_inverse(
slot_logical_pages: torch.Tensor,
logical_page_capacity: int,
) -> torch.Tensor:
"""Build logical_page -> slot_dense_page map without unique/searchsorted."""
page_inverse = torch.full(
(logical_page_capacity,),
-1,
device=slot_logical_pages.device,
dtype=torch.long,
)
if logical_page_capacity == 0:
return page_inverse
# Page 0 is the dummy/padding page and always maps to dense page 0.
page_inverse[0] = 0
if slot_logical_pages.numel() == 0:
return page_inverse
flat_pages = slot_logical_pages.reshape(-1).to(torch.long)
slot_ids = torch.arange(
1,
flat_pages.numel() + 1,
device=flat_pages.device,
dtype=torch.long,
)
valid_pages = (flat_pages > 0) & (flat_pages < logical_page_capacity)
safe_pages = torch.where(valid_pages, flat_pages, torch.zeros_like(flat_pages))
safe_slot_ids = torch.where(valid_pages, slot_ids, torch.zeros_like(slot_ids))
page_inverse.scatter_(0, safe_pages, safe_slot_ids)
return page_inverse
def remap_logical_locs_to_slot_dense_locs(
logical_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
) -> torch.Tensor:
"""Map logical token locs through a fixed-shape slot page inverse."""
dense_locs = torch.full_like(logical_locs, -1)
if logical_locs.numel() == 0 or page_inverse.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)
pages_in_range = logical_pages < page_inverse.numel()
safe_pages = torch.clamp(logical_pages, max=page_inverse.numel() - 1)
dense_pages = page_inverse[safe_pages]
mapped = valid_locs & pages_in_range & (dense_pages >= 0)
if cp_shared_kv_debug_enabled() and torch.any(
valid_locs & pages_in_range & (dense_pages < 0)
):
missing_pages = logical_pages[valid_locs & pages_in_range & (dense_pages < 0)]
raise RuntimeError(
"CP shared KV slot remap got logical locs outside remap_logical_pages. "
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)}"
)
dense_values = dense_pages.to(logical_locs.dtype) * page_size + offsets.to(
logical_locs.dtype
)
return torch.where(mapped, dense_values, dense_locs)
def remap_logical_locs_to_dense_locs(
logical_locs: torch.Tensor,
unique_logical_pages: torch.Tensor,
@@ -254,9 +358,6 @@ def remap_logical_locs_to_dense_locs(
) -> torch.Tensor:
dense_locs = logical_locs.clone()
valid_mask = logical_locs >= 0
if not torch.any(valid_mask):
return dense_locs
valid_locs = logical_locs[valid_mask]
logical_pages = torch.div(valid_locs, page_size, rounding_mode="floor")
offsets = torch.remainder(valid_locs, page_size)
@@ -314,12 +415,12 @@ def build_current_loc_remap(
def logical_pages_from_locs(logical_locs: torch.Tensor, page_size: int) -> torch.Tensor:
logical_pages = logical_locs.clone()
valid_mask = logical_locs >= 0
if torch.any(valid_mask):
logical_pages[valid_mask] = torch.div(
logical_locs[valid_mask],
page_size,
rounding_mode="floor",
)
valid_locs = logical_locs[valid_mask]
logical_pages[valid_mask] = torch.div(
valid_locs,
page_size,
rounding_mode="floor",
)
return logical_pages
@@ -419,14 +520,13 @@ def materialize_local_token_kv_pages(
page_size: int,
) -> torch.Tensor:
dense_num_pages = int(unique_logical_pages.numel()) + 1
dense_kv_cache = kv_cache.new_zeros((dense_num_pages * page_size, *kv_cache.shape[1:]))
dense_kv_cache = kv_cache.new_zeros(
(dense_num_pages * page_size, *kv_cache.shape[1:])
)
if unique_logical_pages.numel() == 0:
return dense_kv_cache
owned_mask = layout.owned_pages_mask(unique_logical_pages)
if not torch.any(owned_mask):
return dense_kv_cache
owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64)
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
torch.long
@@ -439,6 +539,44 @@ def materialize_local_token_kv_pages(
return dense_kv_cache
def materialize_local_token_kv_page_slots(
kv_cache: torch.Tensor,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
page_size: int,
) -> torch.Tensor:
"""Materialize slot-remapped token KV with fixed-shape device ops only."""
dense_num_pages = int(slot_logical_pages.numel()) + 1
dense_kv_cache = kv_cache.new_zeros(
(dense_num_pages * page_size, *kv_cache.shape[1:])
)
if slot_logical_pages.numel() == 0:
return dense_kv_cache
logical_pages = slot_logical_pages.reshape(-1).to(torch.long)
owned_mask = layout.owned_pages_mask(logical_pages)
physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long)
safe_physical_pages = torch.where(
owned_mask,
physical_pages,
torch.zeros_like(physical_pages),
)
page_offsets = torch.arange(page_size, device=kv_cache.device, dtype=torch.long)
src_tokens = (safe_physical_pages[:, None] * page_size + page_offsets).reshape(-1)
dense_body = dense_kv_cache[page_size:].view(
dense_num_pages - 1,
page_size,
*kv_cache.shape[1:],
)
gathered = kv_cache[src_tokens].view_as(dense_body)
owned_view = owned_mask.view(-1, *([1] * (dense_body.ndim - 1)))
zero = torch.zeros((), dtype=kv_cache.dtype, device=kv_cache.device)
dense_body.copy_(torch.where(owned_view, gathered, zero))
return dense_kv_cache
def token_page_copy_debug_checksum(
kv_cache: torch.Tensor,
dense_kv_cache: torch.Tensor,
@@ -481,9 +619,6 @@ def materialize_local_paged_buffer(
return dense_page_buffer
owned_mask = layout.owned_pages_mask(unique_logical_pages)
if not torch.any(owned_mask):
return dense_page_buffer
owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64)
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
torch.long
@@ -493,6 +628,33 @@ def materialize_local_paged_buffer(
return dense_page_buffer
def materialize_local_paged_buffer_page_slots(
page_buffer: torch.Tensor,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
) -> torch.Tensor:
"""Materialize slot-remapped page buffer without nonzero/boolean compaction."""
dense_num_pages = int(slot_logical_pages.numel()) + 1
dense_page_buffer = page_buffer.new_zeros((dense_num_pages, *page_buffer.shape[1:]))
if slot_logical_pages.numel() == 0:
return dense_page_buffer
logical_pages = slot_logical_pages.reshape(-1).to(torch.long)
owned_mask = layout.owned_pages_mask(logical_pages)
physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long)
safe_physical_pages = torch.where(
owned_mask,
physical_pages,
torch.zeros_like(physical_pages),
)
gathered = page_buffer[safe_physical_pages]
owned_view = owned_mask.view(-1, *([1] * (gathered.ndim - 1)))
zero = torch.zeros((), dtype=page_buffer.dtype, device=page_buffer.device)
dense_page_buffer[1:].copy_(torch.where(owned_view, gathered, zero))
return dense_page_buffer
def paged_copy_debug_checksum(
page_buffer: torch.Tensor,
dense_page_buffer: torch.Tensor,
@@ -547,6 +709,7 @@ def materialize_shared_token_kv_buffer(
layout: CpSharedKVLayout,
page_size: int,
remap_logical_locs: torch.Tensor | None = None,
remap_logical_pages: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
_debug_assert_no_tensor_values_below(
logical_locs,
@@ -574,21 +737,61 @@ def materialize_shared_token_kv_buffer(
physical_token_capacity=kv_cache.shape[0],
)
remap_logical_pages = logical_pages_from_locs(remap_logical_locs, page_size)
unique_logical_pages, _ = build_dense_page_remap(remap_logical_pages)
dense_locs = remap_logical_locs_to_dense_locs(
logical_locs,
unique_logical_pages=unique_logical_pages,
page_size=page_size,
)
dense_kv_cache = materialize_local_token_kv_pages(
kv_cache=kv_cache,
unique_logical_pages=unique_logical_pages,
layout=layout,
page_size=page_size,
)
if remap_logical_pages is None:
remap_pages_from_locs = logical_pages_from_locs(remap_logical_locs, page_size)
materialized_logical_pages, _ = build_dense_page_remap(remap_pages_from_locs)
dense_locs = remap_logical_locs_to_dense_locs(
logical_locs,
unique_logical_pages=materialized_logical_pages,
page_size=page_size,
)
use_slot_materialize = False
else:
_debug_assert_no_negative_tensor_values(
remap_logical_pages,
context="CP shared KV token materialize page remap",
tensor_name="remap_logical_pages",
)
remap_logical_pages = filter_pages_mappable_to_physical_pool(
logical_pages=remap_logical_pages,
layout=layout,
physical_page_capacity=kv_cache.shape[0] // page_size,
)
materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
logical_page_capacity = _logical_page_capacity_from_physical_page_capacity(
kv_cache.shape[0] // page_size,
layout,
)
page_inverse = build_slot_page_inverse(
materialized_logical_pages,
logical_page_capacity=logical_page_capacity,
)
dense_locs = remap_logical_locs_to_slot_dense_locs(
logical_locs,
page_inverse=page_inverse,
page_size=page_size,
)
use_slot_materialize = True
if use_slot_materialize:
dense_kv_cache = materialize_local_token_kv_page_slots(
kv_cache=kv_cache,
slot_logical_pages=materialized_logical_pages,
layout=layout,
page_size=page_size,
)
else:
dense_kv_cache = materialize_local_token_kv_pages(
kv_cache=kv_cache,
unique_logical_pages=materialized_logical_pages,
layout=layout,
page_size=page_size,
)
if cp_shared_kv_debug_enabled():
owned_pages = unique_logical_pages[layout.owned_pages_mask(unique_logical_pages)]
owned_pages = materialized_logical_pages[
layout.owned_pages_mask(materialized_logical_pages)
]
physical_pages = layout.logical_pages_to_physical(owned_pages)
cp_shared_kv_debug_log(
"materialize_token_pre",
@@ -598,7 +801,7 @@ def materialize_shared_token_kv_buffer(
layout.cp_rank,
tensor_debug_summary(logical_locs),
tensor_debug_summary(remap_logical_locs),
tensor_debug_summary(unique_logical_pages),
tensor_debug_summary(materialized_logical_pages),
tensor_debug_summary(owned_pages),
tensor_debug_summary(physical_pages),
tensor_debug_summary(dense_locs),
@@ -607,12 +810,14 @@ def materialize_shared_token_kv_buffer(
token_page_copy_debug_checksum(
kv_cache,
dense_kv_cache,
unique_logical_pages,
materialized_logical_pages,
layout,
page_size,
),
)
dense_kv_cache = _all_reduce_materialized_buffer(dense_kv_cache, layout.cp_size)
if cp_shared_kv_debug_enabled():
cp_shared_kv_debug_log(
"materialize_token_post",
@@ -625,7 +830,6 @@ def materialize_shared_token_kv_buffer(
)
return dense_kv_cache, dense_locs
def materialize_shared_paged_buffer(
page_buffer: torch.Tensor,
logical_pages: torch.Tensor,
@@ -641,14 +845,17 @@ def materialize_shared_paged_buffer(
layout=layout,
physical_page_capacity=page_buffer.shape[0],
)
unique_logical_pages, dense_pages = build_dense_page_remap(logical_pages)
dense_page_buffer = materialize_local_paged_buffer(
materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
dense_page_buffer = materialize_local_paged_buffer_page_slots(
page_buffer=page_buffer,
unique_logical_pages=unique_logical_pages,
slot_logical_pages=materialized_logical_pages,
layout=layout,
)
if cp_shared_kv_debug_enabled():
owned_pages = unique_logical_pages[layout.owned_pages_mask(unique_logical_pages)]
owned_pages = materialized_logical_pages[
layout.owned_pages_mask(materialized_logical_pages)
]
physical_pages = layout.logical_pages_to_physical(owned_pages)
cp_shared_kv_debug_log(
"materialize_paged_pre",
@@ -657,7 +864,7 @@ def materialize_shared_paged_buffer(
"active_local_ck=%s owned_copy_ck=%s",
layout.cp_rank,
tensor_debug_summary(logical_pages),
tensor_debug_summary(unique_logical_pages),
tensor_debug_summary(materialized_logical_pages),
tensor_debug_summary(owned_pages),
tensor_debug_summary(physical_pages),
tensor_debug_summary(dense_pages),
@@ -666,11 +873,15 @@ def materialize_shared_paged_buffer(
paged_copy_debug_checksum(
page_buffer,
dense_page_buffer,
unique_logical_pages,
materialized_logical_pages,
layout,
),
)
dense_page_buffer = _all_reduce_materialized_buffer(dense_page_buffer, layout.cp_size)
dense_page_buffer = _all_reduce_materialized_buffer(
dense_page_buffer, layout.cp_size
)
if cp_shared_kv_debug_enabled():
cp_shared_kv_debug_log(
"materialize_paged_post",

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

@@ -1634,6 +1634,7 @@ class NativeSparseAttnBackend(
kv_cache=kv_cache,
logical_locs=page_table_1,
remap_logical_locs=metadata.page_table_1,
remap_logical_pages=metadata.real_page_table,
layout=forward_batch.cp_shared_kv_layout,
page_size=forward_batch.token_to_kv_pool.page_size,
)

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

@@ -366,6 +366,35 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[8:12]))
self.assertTrue(torch.equal(dense_kv[20:24], kv_cache[20:24]))
def test_materialize_token_kv_fast_path_avoids_python_tensor_predicates(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
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
kv_cache = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1)
# Page 1 is owned by CP rank 0, so this also covers the no-local-page
# branch without using torch.any(owned_mask) in Python control flow.
logical_locs = torch.tensor([4, 5], dtype=torch.int64)
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=False
), patch.object(
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
), patch.object(
runtime.torch, "any", side_effect=AssertionError("torch.any sync")
), patch.object(
runtime.torch, "equal", side_effect=AssertionError("torch.equal sync")
):
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=logical_locs,
layout=layout,
page_size=4,
)
self.assertEqual(dense_locs.tolist(), [4, 5])
self.assertEqual(float(dense_kv.abs().sum().item()), 0.0)
def test_materialize_token_kv_keeps_dense_shape_for_shared_remap_source(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
@@ -414,7 +443,105 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
)
)
def test_materialize_paged_buffer_fast_path_avoids_python_tensor_predicates(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
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
page_buffer = torch.arange(0, 4 * 3, dtype=torch.uint8).view(4, 3)
# Page 1 is owned by CP rank 0, so this also covers the no-local-page
# branch without using torch.any(owned_mask) in Python control flow.
logical_pages = torch.tensor([1], dtype=torch.int32)
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=False
), patch.object(
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
), patch.object(
runtime.torch, "any", side_effect=AssertionError("torch.any sync")
), patch.object(
runtime.torch, "equal", side_effect=AssertionError("torch.equal sync")
):
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
)
self.assertEqual(dense_pages.tolist(), [1])
self.assertEqual(int(dense_page_buffer.sum().item()), 0)
def test_materialize_paged_buffer_fast_path_avoids_dynamic_shape_ops(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
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
page_buffer = torch.arange(0, 6 * 3, dtype=torch.uint8).view(6, 3)
logical_pages = torch.tensor([1, 2, 5, 6, 0], dtype=torch.int32)
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=False
), patch.object(
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
), patch.object(
runtime.torch, "unique", side_effect=AssertionError("torch.unique sync")
), patch.object(
runtime.torch, "nonzero", side_effect=AssertionError("torch.nonzero sync")
), patch.object(
runtime.torch,
"searchsorted",
side_effect=AssertionError("torch.searchsorted sync"),
):
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
)
self.assertEqual(dense_pages.tolist(), [1, 2, 3, 4, 0])
self.assertEqual(list(dense_page_buffer.shape), [6, 3])
self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[1]))
self.assertTrue(torch.equal(dense_page_buffer[4], page_buffer[3]))
self.assertEqual(int(dense_page_buffer[1].sum().item()), 0)
self.assertEqual(int(dense_page_buffer[3].sum().item()), 0)
def test_materialize_token_kv_page_slot_source_avoids_dynamic_shape_ops(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
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
kv_cache = torch.arange(0, 24, dtype=torch.float32).view(24, 1, 1)
logical_locs = torch.tensor([8, 9, 24, 25, -1], dtype=torch.int64)
remap_logical_pages = torch.tensor([1, 2, 5, 6], dtype=torch.int32)
with patch.object(
runtime, "cp_shared_kv_debug_enabled", return_value=False
), patch.object(
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
), patch.object(
runtime.torch, "unique", side_effect=AssertionError("torch.unique sync")
), patch.object(
runtime.torch, "nonzero", side_effect=AssertionError("torch.nonzero sync")
), patch.object(
runtime.torch,
"searchsorted",
side_effect=AssertionError("torch.searchsorted sync"),
):
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=4,
)
self.assertEqual(dense_locs.tolist(), [8, 9, 16, 17, -1])
self.assertEqual(list(dense_kv.shape), [20, 1, 1])
self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[4:8]))
self.assertTrue(torch.equal(dense_kv[16:20], kv_cache[12:16]))
self.assertEqual(float(dense_kv[4:8].abs().sum().item()), 0.0)
class TestCpSharedKVLazyDebugLogging(unittest.TestCase):