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

Fix RAGGED CP cache-hit current KV reuse

Browse Source 
FP8 flashmla_sparse uses flattened RAGGED page tables that include both cached prefix and the just-computed current suffix. The old cache-hit path materialized the whole flattened range from persistent KV, which could read current rows through the wrong contract under CP shared-KV and compute padding.\n\nThis change makes the RAGGED path use the page-slot partial-current compose contract: prefix pages are materialized from cache slots while current rows are sourced from fresh k/k_rope and packed for FP8 when needed. A new helper accepts the actual current-row contracts seen by attention code: already-local valid rows, CP-local compute-padded rows, or unsplit global valid rows.\n\nConstraint: CP shared-KV stores and consumes cache at page granularity, while attention current rows may be valid-token tensors rather than cache-write local compute rows.\nRejected: Full materialize prefix+current from persistent KV | it can read current suffix through stale or unordered persistent cache state.\nRejected: Reuse select_cp_local_valid_rows_for_cache_write directly | it only accepts CP-local compute-padded rows and killed prefill when RAGGED supplied global valid rows.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not route RAGGED FP8 cache-hit current suffix through full materialize without proving current write ordering and row-contract compatibility.\nTested: Local py_compile for touched runtime/test files.\nTested: Remote container pytest for RAGGED current compose and compute-padding global-current row selection.\nNot-tested: Full GSM8K warm-cache ETE after restart.
This commit is contained in:
laoyao0822
2026-06-08 03:39:34 +08:00
parent b324407def
commit 2aa0b7313e
4 changed files with 261 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1192,6 +1192,60 @@ def select_cp_local_valid_rows_for_cache_write(
return selected
def select_cp_current_valid_rows_for_reuse(
forward_batch,
current_tensor: torch.Tensor,
) -> torch.Tensor | None:
"""Return CP-local valid current rows for attention current reuse.
Cache-write call sites pass CP-local compute-padded rows and can directly use
``select_cp_local_valid_rows_for_cache_write``. Attention current-reuse
call sites can instead see already-local valid rows or unsplit global valid
current rows. This helper accepts those contracts and returns rows aligned
with ``get_cp_shared_kv_local_out_cache_loc``.
"""
local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
local_valid_rows = (
int(local_out_cache_loc.numel()) if local_out_cache_loc is not None else None
)
tensor_rows = int(current_tensor.shape[0])
if local_valid_rows is not None and tensor_rows == local_valid_rows:
return current_tensor
plan = get_cp_shared_kv_batch_plan(forward_batch)
if plan is None:
return None
compute_padding_enabled = bool(getattr(plan, "compute_padding_enabled", False))
if compute_padding_enabled:
_, expected_compute_rows = _get_cp_local_valid_row_indices_cache(
forward_batch,
plan,
current_tensor.device,
)
if tensor_rows == expected_compute_rows:
return select_cp_local_valid_rows_for_cache_write(
forward_batch,
current_tensor,
)
extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) == 0:
return None
global_valid_rows = sum(int(x) for x in extend_seq_lens_cpu)
if global_valid_rows <= 0 or tensor_rows < global_valid_rows:
return None
return split_tensor_by_cp_batch_plan(
current_tensor[:global_valid_rows].contiguous(),
plan,
mode="1d" if current_tensor.dim() == 1 else "data",
split_kind="valid",
static_padded_tokens=_get_forward_batch_static_padded_tokens(forward_batch),
)
def _pad_cp_request_tensor_for_split(
tensor: torch.Tensor,
*,

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

@@ -63,7 +63,7 @@ from sglang.srt.layers.attention.nsa.utils import (
nsa_cp_round_robin_split_q_seqs,
nsa_use_prefill_cp,
pad_nsa_cache_seqlens,
select_cp_local_valid_rows_for_cache_write,
select_cp_current_valid_rows_for_reuse,
)
from sglang.srt.layers.attention.utils import (
concat_mla_absorb_q_general,
@@ -1853,10 +1853,10 @@ class NativeSparseAttnBackend(
current_k_rope = None
if compute_padding_current:
assert k is not None and k_rope is not None
current_k_nope = select_cp_local_valid_rows_for_cache_write(
current_k_nope = select_cp_current_valid_rows_for_reuse(
forward_batch, k
)
current_k_rope = select_cp_local_valid_rows_for_cache_write(
current_k_rope = select_cp_current_valid_rows_for_reuse(
forward_batch, k_rope
)
current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc(
@@ -1869,6 +1869,8 @@ class NativeSparseAttnBackend(
)
can_reuse_current_kv = (
current_kv_rows_for_reuse is not None
and current_k_nope is not None
and current_k_rope is not None
and int(current_k_nope.shape[0]) == current_kv_rows_for_reuse
and int(current_k_rope.shape[0]) == current_kv_rows_for_reuse
)
@@ -2358,6 +2360,7 @@ class NativeSparseAttnBackend(
)
if index_prefetcher is not None:
index_prefetcher.launch_pending_reduce()
try:
if nsa_impl == "tilelang":
if q_rope is not None:
@@ -2381,14 +2384,165 @@ class NativeSparseAttnBackend(
assert page_table_1_flattened is not None
if forward_batch.uses_cp_shared_kv:
assert forward_batch.cp_shared_kv_layout is not None
if k is None or k_rope is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"CP shared KV RAGGED cache-hit compose requires "
"fresh current MLA KV rows."
)
batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
compute_padding_current = batch_plan is not None and bool(
getattr(batch_plan, "compute_padding_enabled", False)
)
current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc(
forward_batch
)
if current_locs_for_reuse is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"CP shared KV RAGGED cache-hit compose requires "
"local out_cache_loc for current suffix rows."
)
if compute_padding_current:
current_k_nope = select_cp_current_valid_rows_for_reuse(
forward_batch,
k,
)
current_k_rope = select_cp_current_valid_rows_for_reuse(
forward_batch,
k_rope,
)
if current_k_nope is None or current_k_rope is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"CP shared KV RAGGED cache-hit compose could "
"not derive valid current rows under "
"compute padding. "
f"k_rows={int(k.shape[0])} "
f"k_rope_rows={int(k_rope.shape[0])} "
f"local_locs={int(current_locs_for_reuse.numel())}"
)
else:
local_current_rows = int(current_locs_for_reuse.numel())
if int(k.shape[0]) == local_current_rows and int(
k_rope.shape[0]
) == local_current_rows:
current_k_nope = k
current_k_rope = k_rope
else:
valid_current_rows = current_extend_kv_rows_for_reuse(
forward_batch,
k,
k_rope,
)
if valid_current_rows is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"CP shared KV RAGGED cache-hit compose could "
"not derive valid current rows. "
f"k_rows={int(k.shape[0])} "
f"k_rope_rows={int(k_rope.shape[0])} "
f"local_locs={local_current_rows}"
)
current_k_nope = cp_split_and_rebuild_data(
forward_batch,
k[: int(valid_current_rows)].contiguous(),
)
current_k_rope = cp_split_and_rebuild_data(
forward_batch,
k_rope[: int(valid_current_rows)].contiguous(),
)
if (
int(current_k_nope.shape[0])
!= int(current_locs_for_reuse.numel())
or int(current_k_rope.shape[0])
!= int(current_locs_for_reuse.numel())
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"CP shared KV RAGGED current rows do not match "
"local out_cache_loc. "
f"k_rows={int(current_k_nope.shape[0])} "
f"k_rope_rows={int(current_k_rope.shape[0])} "
f"local_locs={int(current_locs_for_reuse.numel())}"
)
if is_packed_fp8_mla_kv_cache(kv_cache):
current_kv_cache = pack_current_mla_kv_for_reuse(
current_k_nope,
current_k_rope,
kv_cache=kv_cache,
)
else:
current_kv_cache = _cat(
[current_k_nope, current_k_rope],
dim=-1,
)
page_size = int(forward_batch.token_to_kv_pool.page_size)
prefix_lens_cpu = getattr(
forward_batch,
"extend_prefix_lens_cpu",
None,
)
extend_lens_cpu = getattr(
forward_batch,
"extend_seq_lens_cpu",
None,
)
prefix_lens_valid = (
prefix_lens_cpu is not None
and len(prefix_lens_cpu) > 0
and all(
int(prefix_len) >= 0
and int(prefix_len) % page_size == 0
for prefix_len in prefix_lens_cpu
)
)
if not prefix_lens_valid:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][mla_ragged_current_sync] "
"CP shared KV RAGGED cache-hit compose requires "
"page-aligned prefix metadata. "
f"prefix_lens={prefix_lens_cpu} "
f"extend_lens={extend_lens_cpu} "
f"page_size={page_size}"
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
prefix_slot_spans = None
else:
prefix_pages = 0
prefix_slot_spans = build_batch_prefix_slot_spans(
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
page_size=page_size,
)
current_slot_spans = build_batch_current_slot_spans(
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
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=page_size,
)
kv_cache, page_table_1_flattened = (
materialize_shared_token_kv_buffer(
materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=page_table_1_flattened,
current_kv_cache=current_kv_cache,
current_locs=current_locs_for_reuse,
slot_remap=slot_remap,
layout=forward_batch.cp_shared_kv_layout,
page_size=forward_batch.token_to_kv_pool.page_size,
nvtx_source="mla.ragged_full_materialize",
nvtx_layer_id=layer.layer_id,
page_size=page_size,
prefix_pages=prefix_pages,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
layer_id=layer.layer_id,
nvtx_source="mla.ragged_partial_current_sync",
)
)
kv_cache = dequantize_k_cache_paged(