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

Keep CP shared-KV cache hits page-aligned without short fallback

Browse Source 
CP shared KV and HiCache need a stable contract where physical cache coverage is page-aligned, while scheduler/radix-visible hit length remains the valid token length. This records the contract, adds page-aligned extent metadata, keeps owner assignment on actual tail pages instead of short-prefix fallback, and updates partial current reuse tests around tail-page masking.

Constraint: CP owner lanes operate on page units while scheduler and radix hit accounting must remain token-valid.

Rejected: Pad short suffixes to cp_size or 2*cp_size pages | wastes KV capacity and can turn a small tail into a much larger physical span.

Rejected: Silent direct-write or prefetch fallback | production fallback must be warning-visible for diagnosis.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce replicated short-radix fallback without checking docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Tested: local py_compile for touched runtime, utility, owner, and unit-test files.

Tested: remote g0034 container three-file suite: 122 passed, 5 warnings.

Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson.

Not-tested: CUDA E2E runtime for this commit.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 05:04:50 +08:00
parent 1e1807f2bc
commit c551623ca8
11 changed files with 1369 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -18,11 +18,11 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_should_log_layer,
filter_locs_mappable_to_physical_pool,
filter_pages_mappable_to_physical_pool,
fill_current_kv_page_slots_and_remap_locs,
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,
merge_materialized_and_current_kv,
remap_logical_pages_to_slot_dense_pages,
remap_logical_locs_to_slot_dense_locs_optimized,
slot_range_to_page_slice,
@@ -39,6 +39,14 @@ def _prefetch_log(message: str, *args) -> None:
cp_shared_kv_mla_prefetch_log(message, *args)
def _mla_prefetch_fallback_log(reason: str, message: str, *args) -> None:
logger.warning(
"[CP_SHARED_KV_FALLBACK][mla_prefetch] reason=%s " + message,
reason,
*args,
)
def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None:
logger.warning(
"[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message,
@@ -384,8 +392,10 @@ class CpSharedKVMlaPrefetcher:
return None
extend_prefix_len = int(extend_prefix_lens_cpu[0])
if extend_prefix_len <= 0 or extend_prefix_len % page_size != 0:
_prefetch_log(
"create_skip reason=prefix_not_page_aligned prefix_len=%s page_size=%s",
_mla_prefetch_fallback_log(
"prefix_not_page_aligned",
"prefix length is zero or not page-aligned. "
"prefix_len=%s page_size=%s",
extend_prefix_len,
page_size,
)
@@ -726,14 +736,15 @@ class CpSharedKVMlaPrefetcher:
page_inverse=self.page_inverse,
page_size=self.page_size,
)
mixed_kv_cache, mixed_locs, _ = merge_materialized_and_current_kv(
materialized_kv_cache=dense_kv_cache,
mixed_kv_cache, mixed_locs, _ = fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_kv_cache,
materialized_dense_locs=dense_locs,
current_kv_cache=current_kv_cache,
logical_locs=logical_locs,
current_locs=current_locs,
page_size=current_remap_page_size,
logical_page_capacity=current_remap_logical_page_capacity,
page_inverse=self.page_inverse,
page_size=self.page_size,
mask_non_current_in_current_pages=True,
)
remap_ms = _cpu_timing_ms(remap_cpu)
total_ms = _cpu_timing_ms(consume_cpu)

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

@@ -642,6 +642,135 @@ def remap_logical_locs_to_slot_dense_locs_optimized(
)
def _try_tai_fill_current_kv_page_slots_and_remap_locs(
*,
dense_kv_cache: torch.Tensor,
materialized_dense_locs: torch.Tensor,
current_kv_cache: torch.Tensor,
logical_locs: torch.Tensor,
current_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
mask_non_current_in_current_pages: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
if not _tai_materialize_runtime_enabled():
return None
kernels = _load_tai_materialize_kernels()
if kernels is None:
return None
fill_kernel = getattr(
kernels,
"fill_current_token_kv_page_slots_and_remap_locs",
None,
)
if fill_kernel is None:
_log_tai_materialize_fallback(
"fill_current_missing",
"CP shared KV tai current-slot fill kernel is unavailable; "
"falling back to torch reference. Upgrade tai-kernel to keep this "
"hot path off PyTorch. page_size=%s current_rows=%s query_locs=%s",
page_size,
int(current_kv_cache.shape[0]),
int(logical_locs.numel()),
limit=1,
)
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),
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",
"CP shared KV tai current-slot fill failed; falling back to torch "
"reference. error=%s page_size=%s current_rows=%s query_locs=%s",
exc,
page_size,
int(current_kv_cache.shape[0]),
int(logical_locs.numel()),
)
return None
def fill_current_kv_page_slots_and_remap_locs(
*,
dense_kv_cache: torch.Tensor,
materialized_dense_locs: torch.Tensor,
current_kv_cache: torch.Tensor,
logical_locs: torch.Tensor,
current_locs: torch.Tensor,
page_inverse: torch.Tensor,
page_size: int,
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.
This is the CP shared-KV prefetch compose path. The dense buffer already has
page-aligned slots for prefix and suffix pages; current rows should be copied
into those suffix slots instead of appended to a second compact suffix.
"""
tai_result = _try_tai_fill_current_kv_page_slots_and_remap_locs(
dense_kv_cache=dense_kv_cache,
materialized_dense_locs=materialized_dense_locs,
current_kv_cache=current_kv_cache,
logical_locs=logical_locs,
current_locs=current_locs,
page_inverse=page_inverse,
page_size=page_size,
mask_non_current_in_current_pages=mask_non_current_in_current_pages,
)
if tai_result is not None:
return tai_result
if dense_kv_cache.is_cuda:
_log_tai_materialize_fallback(
"fill_current_torch_reference_cuda",
"CP shared KV current-slot fill is using the torch reference on CUDA; "
"this is a fallback and should not be the steady-state hot path. "
"page_size=%s current_rows=%s query_locs=%s",
page_size,
int(current_kv_cache.shape[0]),
int(logical_locs.numel()),
limit=1,
)
current_dense_locs = remap_logical_locs_to_slot_dense_locs_optimized(
current_locs.reshape(-1),
page_inverse=page_inverse,
page_size=page_size,
)
valid_current_rows = current_dense_locs >= 0
if torch.any(valid_current_rows):
dense_kv_cache[current_dense_locs[valid_current_rows].to(torch.long)] = (
current_kv_cache[valid_current_rows]
)
current_mask, _ = build_current_loc_remap(logical_locs, current_locs)
mixed_locs = materialized_dense_locs
if mask_non_current_in_current_pages:
current_page_mask = build_current_page_mask(
logical_locs,
current_locs,
page_size=page_size,
)
mixed_locs = torch.where(
current_page_mask & (~current_mask),
torch.full_like(materialized_dense_locs, -1),
materialized_dense_locs,
)
return dense_kv_cache, mixed_locs, current_mask
def _copy_tai_dense_slot_range_body(
*,
tai_dense_kv_cache: torch.Tensor,
@@ -866,6 +995,43 @@ def current_loc_remap_fast_path_args(
return page_size, logical_page_capacity
def build_current_page_mask(
query_locs: torch.Tensor,
current_locs: torch.Tensor,
*,
page_size: int,
) -> torch.Tensor:
"""Return true for query locs that fall inside a current suffix page.
Partial-current reuse can have a tail page where only the first N rows are
real current tokens. The page itself is part of the physical suffix, but
rows after the valid current tokens are padding/slack and must be invisible
to attention.
"""
mask = torch.zeros_like(query_locs, dtype=torch.bool)
if page_size <= 1 or query_locs.numel() == 0 or current_locs.numel() == 0:
return mask
query_flat = query_locs.reshape(-1).to(torch.long)
current_flat = current_locs.reshape(-1).to(torch.long)
current_pages = torch.unique(
torch.div(current_flat, page_size, rounding_mode="floor")
)
if current_pages.numel() == 0:
return mask
sorted_pages, _ = torch.sort(current_pages)
valid_query = query_flat >= 0
safe_query = torch.where(valid_query, query_flat, torch.zeros_like(query_flat))
query_pages = torch.div(safe_query, page_size, rounding_mode="floor")
insert_positions = torch.searchsorted(sorted_pages, query_pages)
safe_positions = torch.clamp(insert_positions, max=sorted_pages.numel() - 1)
in_range = insert_positions < sorted_pages.numel()
matched = valid_query & in_range & (sorted_pages[safe_positions] == query_pages)
return matched.reshape(query_locs.shape)
def merge_materialized_and_current_kv(
*,
materialized_kv_cache: torch.Tensor,
@@ -882,6 +1048,11 @@ def merge_materialized_and_current_kv(
materialization path. Entries corresponding to current extend tokens are
replaced with offsets into the appended ``current_kv_cache``. Non-current
entries remain untouched, including ``-1`` invalid sentinels.
This helper is for compact prefix materialization that appends a compact
current suffix. Prefetch paths with existing page slots must use
:func:`fill_current_kv_page_slots_and_remap_locs` instead, so the physical
page padding stays in the original dense slot layout.
"""
current_mask, current_rows = build_current_loc_remap(

113
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -173,6 +173,32 @@ def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens):
return nsa_cache_seqlens
@dataclass(frozen=True)
class PageAlignedCacheExtent:
valid_tokens: int
padded_pages: int
padded_tokens: int
padding_tokens: int
def build_page_aligned_cache_extent(
*, valid_tokens: int, page_size: int
) -> PageAlignedCacheExtent:
if valid_tokens < 0:
raise ValueError(f"valid_tokens must be non-negative, got {valid_tokens}")
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
padded_pages = ceil_div(valid_tokens, page_size) if valid_tokens > 0 else 0
padded_tokens = padded_pages * page_size
return PageAlignedCacheExtent(
valid_tokens=valid_tokens,
padded_pages=padded_pages,
padded_tokens=padded_tokens,
padding_tokens=padded_tokens - valid_tokens,
)
@dataclass
class PageAlignedInSeqSplitInfo:
page_aligned: bool = False
@@ -180,6 +206,10 @@ class PageAlignedInSeqSplitInfo:
extend_prefix_len: int = 0
segment_page_starts: List[int] = None
segment_page_ends: List[int] = None
extend_valid_tokens: int = 0
extend_padded_pages: int = 0
extend_padded_tokens: int = 0
extend_padding_tokens: int = 0
@dataclass
@@ -208,6 +238,10 @@ class NSAContextParallelMetadata:
extend_prefix_len: int = 0
segment_page_starts: List[int] = None
segment_page_ends: List[int] = None
extend_valid_tokens: int = 0
extend_padded_pages: int = 0
extend_padded_tokens: int = 0
extend_padding_tokens: int = 0
def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]:
@@ -233,6 +267,10 @@ def _fallback_page_aligned_split_info(
extend_prefix_len=extend_prefix_len,
segment_page_starts=[],
segment_page_ends=[],
extend_valid_tokens=0,
extend_padded_pages=0,
extend_padded_tokens=0,
extend_padding_tokens=0,
)
@@ -246,13 +284,10 @@ def build_page_aligned_in_seq_split_list(
) -> Tuple[List[int], PageAlignedInSeqSplitInfo]:
"""Build an in-seq split list whose real-token boundaries do not cut pages.
Phase 4 deliberately uses a conservative gate for cache-miss chunks: at
least `2 * cp_size` page units are required so every zigzag segment has at
least one page unit. For radix-hit suffixes with a page-aligned prefix, the
gate is relaxed to `cp_size` page units so every CP rank still receives at
least one local page while second zigzag segments may be empty. When the
gate does not hold, this helper falls back to the existing token-balanced
split and marks the result as not page-aligned.
The split remains valid-token based, but page metadata covers the physical
page units. Short chunks keep the page-aligned contract by assigning zero
valid tokens to surplus zigzag segments instead of falling back to
token-balanced splits that would poison later shared-KV/HiCache reuse.
"""
if extend_len < 0:
@@ -271,14 +306,14 @@ def build_page_aligned_in_seq_split_list(
if page_size <= 1 or extend_len <= 0 or extend_prefix_len % page_size != 0:
return fallback_split, fallback_info
extent = build_page_aligned_cache_extent(
valid_tokens=extend_len,
page_size=page_size,
)
full_pages = extend_len // page_size
tail_tokens = extend_len % page_size
num_page_units = full_pages + (1 if tail_tokens > 0 else 0)
num_page_units = extent.padded_pages
cp_segment_num = cp_size * 2
if num_page_units < cp_size or (
num_page_units < cp_segment_num and extend_prefix_len == 0
):
return fallback_split, fallback_info
base_units = num_page_units // cp_segment_num
remainder_units = num_page_units % cp_segment_num
@@ -315,6 +350,10 @@ def build_page_aligned_in_seq_split_list(
extend_prefix_len=extend_prefix_len,
segment_page_starts=segment_page_starts,
segment_page_ends=segment_page_ends,
extend_valid_tokens=extent.valid_tokens,
extend_padded_pages=extent.padded_pages,
extend_padded_tokens=extent.padded_tokens,
extend_padding_tokens=extent.padding_tokens,
)
@@ -356,50 +395,12 @@ def should_use_replicated_compute_for_short_radix_hit(
) -> bool:
"""Return whether a short radix-hit suffix should avoid CP splitting.
With CP shared KV, radix-hit suffixes can be page-aligned but shorter than
one page per CP rank. A page-aligned CP split would give some ranks zero
local tokens, which is unsafe for parts of the current CP collective/kernel
path. Instead, keep the original non-CP behavior: every rank computes the
short suffix, while shared-KV write filters persist only pages owned by the
local rank.
Kept as a compatibility hook for older callers. The page-aligned cache
contract no longer uses replicated compute for short suffixes: zero-length
CP segments are preferred over breaking the physical page-owner pattern.
"""
if (
forward_batch is None
or cp_size <= 0
or not getattr(forward_batch, "uses_cp_shared_kv", False)
):
return False
extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
if (
extend_seq_lens_cpu is None
or extend_prefix_lens_cpu is None
or len(extend_seq_lens_cpu) != 1
or len(extend_prefix_lens_cpu) != 1
):
return False
token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
page_size = getattr(token_to_kv_pool, "page_size", None)
if page_size is None:
return False
page_size = int(page_size)
if page_size <= 1:
return False
extend_len = int(extend_seq_lens_cpu[0])
extend_prefix_len = int(extend_prefix_lens_cpu[0])
if (
extend_len <= 0
or extend_prefix_len <= 0
or extend_prefix_len % page_size != 0
):
return False
num_page_units = ceil_div(extend_len, page_size)
return 0 < num_page_units < cp_size
return False
def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
@@ -1177,6 +1178,10 @@ def prepare_input_dp_with_cp_dsa(
extend_prefix_len=page_split_info.extend_prefix_len,
segment_page_starts=page_split_info.segment_page_starts,
segment_page_ends=page_split_info.segment_page_ends,
extend_valid_tokens=page_split_info.extend_valid_tokens,
extend_padded_pages=page_split_info.extend_padded_pages,
extend_padded_tokens=page_split_info.extend_padded_tokens,
extend_padding_tokens=page_split_info.extend_padding_tokens,
)
return nsa_cp_metadata

14
python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py
View File

@@ -24,9 +24,6 @@ def get_in_seq_page_compute_owner_unavailable_reason(
full_pages = extend_len // page_size
tail_tokens = extend_len % page_size
num_page_units = full_pages + (1 if tail_tokens > 0 else 0)
if num_page_units < cp_size * 2 and extend_prefix_len == 0:
return "too_short_for_page_aligned"
return None
@@ -39,12 +36,11 @@ def build_in_seq_page_compute_owners(
) -> Optional[List[int]]:
"""Return compute-owner CP rank for each newly allocated current page.
This mirrors the Phase 4 page-aligned `in-seq-split` segmentation for
normal CP chunks, but it only returns page-unit owners for the real extend
chunk. Short radix-hit suffixes with fewer pages than CP ranks are also
allowed: runtime keeps replicated compute for those chunks and the shared
KV write filters persist only locally owned pages. `None` means the batch
must stay on the legacy allocation/write path.
This mirrors the page-aligned `in-seq-split` segmentation and returns one
owner per physical page unit, including a tail page. Short chunks keep the
page-aligned owner pattern by assigning zero pages to surplus zigzag
segments instead of falling back to legacy allocation. `None` means the
batch must stay on the legacy allocation/write path.
"""
if cp_size <= 0: