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

Reuse CP shared KV remaps across layer materialization

Browse Source 
CP shared KV materialization repeatedly rebuilt the same logical-page slot remaps and page inverse metadata for each layer. Cache the token and paged remap metadata on the forward batch so MLA KV, index K/scale, and prefetch paths can reuse the layer-independent mapping while still materializing layer-specific data through the existing tai/torch runtime paths.

Constraint: Only mapping metadata is batch-scoped; dense KV/index contents remain layer-specific and are not reused.
Rejected: Cache fully materialized dense KV/index buffers | would add large per-layer memory residency and invalidation complexity.
Confidence: medium
Scope-risk: moderate
Directive: Do not assume this removes materialize or CP all-reduce cost; profile tai fallback logs and Nsight kernels before attributing E2E gains or losses.
Tested: git diff --check
Tested: remote g0034 container PYTHONPATH=python python3 -m pytest test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (52 passed, 5 warnings)
Not-tested: Full GLM-5 disaggregated E2E performance run
This commit is contained in:
laoyao0822
2026-05-12 20:02:22 +08:00
parent 099bcfb41e
commit c5c30a3f50
7 changed files with 476 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9,15 +9,14 @@ import torch
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
_all_reduce_materialized_buffer_async,
_all_reduce_materialized_buffer_range,
build_shared_token_kv_slot_remap,
build_slot_page_inverse_optimized,
build_slot_page_remap,
cp_shared_kv_debug_enabled,
cp_shared_kv_mla_prefetch_enabled,
cp_shared_kv_mla_prefetch_log,
cp_shared_kv_mla_prefetch_should_log_layer,
filter_locs_mappable_to_physical_pool,
filter_pages_mappable_to_physical_pool,
get_or_build_shared_paged_buffer_slot_remap,
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,
@@ -195,9 +194,9 @@ class CpSharedKVMlaPrefetcher:
try:
first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0))
kv_cache = token_to_kv_pool.get_key_buffer(first_layer_id)
remap = build_shared_token_kv_slot_remap(
remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
logical_locs=None,
remap_logical_pages=real_page_table,
layout=layout,
page_size=page_size,
@@ -612,18 +611,11 @@ class CpSharedKVIndexPrefetcher:
page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer(
layer_id=first_layer_id
)
remap_logical_pages = filter_pages_mappable_to_physical_pool(
remap = get_or_build_shared_paged_buffer_slot_remap(
forward_batch,
page_buffer=page_buffer,
logical_pages=real_page_table,
layout=layout,
physical_page_capacity=page_buffer.shape[0],
)
slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (
layout.cp_size
) + 1
page_inverse = build_slot_page_inverse_optimized(
slot_logical_pages,
logical_page_capacity=logical_page_capacity,
)
except Exception as exc:
_index_prefetch_fallback_log(
@@ -639,17 +631,17 @@ class CpSharedKVIndexPrefetcher:
layout.cp_rank,
layout.cp_size,
prefix_pages,
int(slot_logical_pages.numel()),
int(slot_logical_pages.numel()) + 1,
int(remap.slot_logical_pages.numel()),
remap.dense_num_pages,
page_size,
)
return cls(
layout=layout,
prefix_pages=prefix_pages,
slot_logical_pages=slot_logical_pages,
page_inverse=page_inverse,
dense_num_pages=int(slot_logical_pages.numel()) + 1,
slot_logical_pages=remap.slot_logical_pages,
page_inverse=remap.page_inverse,
dense_num_pages=remap.dense_num_pages,
)
def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool:

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

@@ -3,6 +3,7 @@ from __future__ import annotations
import logging
from dataclasses import dataclass
from functools import lru_cache
from typing import Any
import torch
@@ -15,6 +16,7 @@ logger = logging.getLogger(__name__)
_DEBUG_LOG_COUNTS: dict[str, int] = {}
_TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_SLOT_REMAP_CACHE_LOG_COUNTS: dict[str, int] = {}
_MLA_PREFETCH_LOG_PROBE_LAYER = 2
_SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get()
@@ -65,6 +67,102 @@ class SharedTokenKVSlotRemap:
dense_num_pages: int
@dataclass(frozen=True)
class SharedPagedBufferSlotRemap:
slot_logical_pages: torch.Tensor
page_inverse: torch.Tensor
dense_pages: torch.Tensor
logical_page_capacity: int
dense_num_pages: int
def _tensor_identity_key(tensor: torch.Tensor) -> tuple[int, tuple[int, ...], str, str]:
return (
int(tensor.untyped_storage().data_ptr()),
tuple(int(dim) for dim in tensor.shape),
str(tensor.dtype),
str(tensor.device),
)
def _slot_remap_cache_key(
*,
logical_pages: torch.Tensor,
physical_page_capacity: int,
layout: CpSharedKVLayout,
page_size: int,
) -> tuple[object, ...]:
return (
_tensor_identity_key(logical_pages),
int(physical_page_capacity),
int(layout.page_size),
int(layout.cp_size),
int(layout.cp_rank),
int(page_size),
)
def _log_slot_remap_cache_not_reused(
*,
kind: str,
reason: str,
cached_key: tuple[object, ...] | None,
new_key: tuple[object, ...],
forward_batch: Any,
layout: CpSharedKVLayout,
limit: int = 8,
) -> None:
log_key = f"{kind}:{reason}"
count = _SLOT_REMAP_CACHE_LOG_COUNTS.get(log_key, 0)
if count >= limit:
return
_SLOT_REMAP_CACHE_LOG_COUNTS[log_key] = count + 1
logger.warning(
"CP shared KV %s slot remap cache not reused (%s): "
"cp_rank=%s cp_size=%s batch_size=%s forward_mode=%s "
"cached_key=%s new_key=%s",
kind,
reason,
layout.cp_rank,
layout.cp_size,
getattr(forward_batch, "batch_size", None),
getattr(forward_batch, "forward_mode", None),
cached_key,
new_key,
)
def _maybe_log_slot_remap_cache_not_reused(
*,
kind: str,
cached_key: tuple[object, ...] | None,
cached: object | None,
new_key: tuple[object, ...],
forward_batch: Any,
layout: CpSharedKVLayout,
) -> None:
if cached is None and cached_key is None:
# First build for this forward batch. This is the expected cold path.
return
if cached is None:
reason = "missing_cached_value"
elif cached_key is None:
reason = "missing_cached_key"
elif cached_key != new_key:
reason = "key_mismatch"
else:
return
_log_slot_remap_cache_not_reused(
kind=kind,
reason=reason,
cached_key=cached_key,
new_key=new_key,
forward_batch=forward_batch,
layout=layout,
)
@lru_cache(maxsize=1)
def _load_tai_materialize_kernels():
try:
@@ -863,6 +961,126 @@ def build_shared_token_kv_slot_remap(
)
def build_shared_paged_buffer_slot_remap(
page_buffer: torch.Tensor,
logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
) -> SharedPagedBufferSlotRemap:
"""Build the fixed slot-layout remap used by shared paged-buffer materialize."""
_debug_assert_no_negative_tensor_values(
logical_pages,
context="CP shared KV paged materialize",
tensor_name="logical_pages",
)
logical_pages = filter_pages_mappable_to_physical_pool(
logical_pages=logical_pages,
layout=layout,
physical_page_capacity=page_buffer.shape[0],
)
slot_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (
layout.cp_size
) + 1
page_inverse = build_slot_page_inverse_optimized(
slot_logical_pages,
logical_page_capacity=logical_page_capacity,
)
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,
)
def get_or_build_shared_token_kv_slot_remap(
forward_batch: Any,
*,
kv_cache: torch.Tensor,
remap_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
page_size: int,
) -> SharedTokenKVSlotRemap:
"""Cache token-KV materialize slot metadata on a forward batch.
The logical page table and local physical pool capacity are batch-scoped,
not layer-scoped. The actual KV rows are still materialized per layer, but
the page inverse used to map layer-specific top-k locs into the dense view
can be reused across all layers in the same forward pass.
"""
physical_page_capacity = kv_cache.shape[0] // page_size
key = _slot_remap_cache_key(
logical_pages=remap_logical_pages,
physical_page_capacity=physical_page_capacity,
layout=layout,
page_size=page_size,
)
cached_key = getattr(forward_batch, "cp_shared_kv_token_slot_remap_key", None)
cached = getattr(forward_batch, "cp_shared_kv_token_slot_remap", None)
if cached is not None and cached_key == key:
return cached
_maybe_log_slot_remap_cache_not_reused(
kind="token",
cached_key=cached_key,
cached=cached,
new_key=key,
forward_batch=forward_batch,
layout=layout,
)
remap = build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=None,
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
forward_batch.cp_shared_kv_token_slot_remap_key = key
forward_batch.cp_shared_kv_token_slot_remap = remap
return remap
def get_or_build_shared_paged_buffer_slot_remap(
forward_batch: Any,
*,
page_buffer: torch.Tensor,
logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
) -> SharedPagedBufferSlotRemap:
"""Cache paged-buffer materialize slot metadata on a forward batch."""
key = _slot_remap_cache_key(
logical_pages=logical_pages,
physical_page_capacity=page_buffer.shape[0],
layout=layout,
page_size=layout.page_size,
)
cached_key = getattr(forward_batch, "cp_shared_kv_paged_slot_remap_key", None)
cached = getattr(forward_batch, "cp_shared_kv_paged_slot_remap", None)
if cached is not None and cached_key == key:
return cached
_maybe_log_slot_remap_cache_not_reused(
kind="paged",
cached_key=cached_key,
cached=cached,
new_key=key,
forward_batch=forward_batch,
layout=layout,
)
remap = build_shared_paged_buffer_slot_remap(
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
)
forward_batch.cp_shared_kv_paged_slot_remap_key = key
forward_batch.cp_shared_kv_paged_slot_remap = remap
return remap
def remap_logical_locs_to_dense_locs(
logical_locs: torch.Tensor,
unique_logical_pages: torch.Tensor,
@@ -1483,6 +1701,7 @@ def materialize_shared_token_kv_buffer(
page_size: int,
remap_logical_locs: torch.Tensor | None = None,
remap_logical_pages: torch.Tensor | None = None,
slot_remap: SharedTokenKVSlotRemap | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
_debug_assert_no_tensor_values_below(
logical_locs,
@@ -1511,7 +1730,15 @@ def materialize_shared_token_kv_buffer(
)
dense_kv_cache = None
if remap_logical_pages is 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(
logical_locs,
page_inverse=slot_remap.page_inverse,
page_size=page_size,
)
use_slot_materialize = True
elif 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(
@@ -1624,6 +1851,7 @@ def materialize_shared_paged_buffer(
page_buffer: torch.Tensor,
logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
slot_remap: SharedPagedBufferSlotRemap | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
_debug_assert_no_negative_tensor_values(
logical_pages,
@@ -1640,16 +1868,24 @@ def materialize_shared_paged_buffer(
logical_pages=logical_pages,
layout=layout,
)
if tai_result is None:
materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
if tai_result is not None:
dense_page_buffer, dense_pages = tai_result
materialized_logical_pages = logical_pages.reshape(-1)
elif slot_remap is not None:
materialized_logical_pages = slot_remap.slot_logical_pages
dense_pages = slot_remap.dense_pages
dense_page_buffer = materialize_local_paged_buffer_page_slots(
page_buffer=page_buffer,
slot_logical_pages=materialized_logical_pages,
layout=layout,
)
else:
dense_page_buffer, dense_pages = tai_result
materialized_logical_pages = logical_pages.reshape(-1)
materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
dense_page_buffer = materialize_local_paged_buffer_page_slots(
page_buffer=page_buffer,
slot_logical_pages=materialized_logical_pages,
layout=layout,
)
if cp_shared_kv_debug_enabled():
owned_pages = materialized_logical_pages[

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

@@ -19,6 +19,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_debug_log,
cp_shared_kv_current_reuse_enabled,
filter_owned_logical_locs,
get_or_build_shared_paged_buffer_slot_remap,
is_current_only_extend_batch,
materialize_shared_paged_buffer,
tensor_debug_checksum,
@@ -344,10 +345,17 @@ class Indexer(MultiPlatformOp):
layer_id,
tensor_debug_summary(logical_page_table),
)
slot_remap = get_or_build_shared_paged_buffer_slot_remap(
forward_batch,
page_buffer=index_buffer,
logical_pages=logical_page_table,
layout=layout,
)
materialized, dense_pages = materialize_shared_paged_buffer(
page_buffer=index_buffer,
logical_pages=logical_page_table,
layout=layout,
slot_remap=slot_remap,
)
if cp_shared_kv_debug_enabled():
cp_shared_kv_debug_log(

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

@@ -22,6 +22,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_log_enabled,
cp_shared_kv_mla_prefetch_should_log_layer,
filter_owned_logical_locs,
get_or_build_shared_token_kv_slot_remap,
is_current_only_extend_batch,
materialize_shared_token_kv_buffer,
tensor_debug_checksum,
@@ -1730,6 +1731,13 @@ class NativeSparseAttnBackend(
if prefetched_kv is not None:
kv_cache, page_table_1 = prefetched_kv
else:
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
remap_logical_pages=metadata.real_page_table,
layout=forward_batch.cp_shared_kv_layout,
page_size=forward_batch.token_to_kv_pool.page_size,
)
kv_cache, page_table_1 = materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=page_table_1,
@@ -1737,6 +1745,7 @@ class NativeSparseAttnBackend(
remap_logical_pages=metadata.real_page_table,
layout=forward_batch.cp_shared_kv_layout,
page_size=forward_batch.token_to_kv_pool.page_size,
slot_remap=slot_remap,
)
if mla_prefetcher is not None:

4
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -425,6 +425,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
cp_local_out_cache_loc: Optional[torch.Tensor] = None
cp_local_physical_out_cache_loc: Optional[torch.Tensor] = None
cp_shared_mla_direct_write_done: bool = False
cp_shared_kv_token_slot_remap_key: Optional[Any] = None
cp_shared_kv_token_slot_remap: Optional[Any] = None
cp_shared_kv_paged_slot_remap_key: Optional[Any] = None
cp_shared_kv_paged_slot_remap: Optional[Any] = None
cp_shared_kv_mla_prefetcher: Optional[Any] = None
cp_shared_kv_index_prefetcher: Optional[Any] = None