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

Repair CP current-slot composition instead of disabling reuse

Browse Source 
CP shared KV current-only and partial-current paths were previously unsafe because owner-local current rows could be filled into dense page slots without synchronizing the current slots across CP ranks. This preserves the fast path and fixes the missing synchronization at the actual contract boundary: prefix slots remain materialized by the existing IPC/collective path, while freshly produced current slots are composed from current forward tensors and reduced only over the current slot ranges.\n\nThe change also keeps page-tail compute padding explicit, restores current-only index compose, fixes TAI index descriptor dtype consistency, and records the investigation ledger to avoid repeating discarded hypotheses.\n\nConstraint: CP shared KV must preserve current reuse and TAI materialize fast paths; blanket disable is not an acceptable fix.\nConstraint: Prefix slots must not be reduced twice after IPC/prefix materialization.\nRejected: Disable current-only MLA/index reuse | hides the owner-local composition bug and regresses the intended fast path.\nRejected: Disable TAI materialize globally | avoids symptoms without proving the byte/layout contract.\nConfidence: medium\nScope-risk: broad\nDirective: Do not remove current-slot range synchronization unless replacing it with an equivalent owner-aware P2P/IPC gather contract.\nTested: Local py_compile for touched files.\nTested: Remote g0034 py_compile for touched runtime/test files.\nTested: Remote test_cp_shared_kv_runtime.py: 111 passed, 5 warnings, 2 subtests passed.\nTested: Remote targeted current-slot compose tests and direct current-only index compose script.\nNot-tested: Full ETE output quality and decode accept len after restart.\nNot-tested: Full test_nsa_cp_utils.py collection on remote due incomplete installed sgl_kernel import/fake-op environment.
This commit is contained in:
laoyao0822
2026-06-05 06:25:06 +08:00
parent e200091638
commit c8510938b9
8 changed files with 898 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -772,6 +772,21 @@ class CpSharedKVMlaPrefetcher:
page_size=self.page_size,
mask_non_current_in_current_pages=True,
)
if self.layout.cp_size > 1 and self.prefix_pages < self.total_slots:
current_rows = slot_range_to_token_slice(
self.page_size,
self.prefix_pages,
self.total_slots,
)
_all_reduce_materialized_buffer_range(
mixed_kv_cache,
self.layout.cp_size,
current_rows.start,
current_rows.stop,
nvtx_source="mla.prefetch_current",
nvtx_layer_id=layer_id,
nvtx_cp_rank=self.layout.cp_rank,
)
remap_ms = _cpu_timing_ms(remap_cpu)
total_ms = _cpu_timing_ms(consume_cpu)
self._log_layer(
@@ -1551,6 +1566,20 @@ class CpSharedKVIndexPrefetcher:
page_size=page_size,
index_head_dim=index_head_dim,
)
if self.layout.cp_size > 1 and self.prefix_pages < self.total_slots:
current_pages = slot_range_to_page_slice(
self.prefix_pages,
self.total_slots,
)
_all_reduce_materialized_buffer_range(
dense_page_buffer,
self.layout.cp_size,
current_pages.start,
current_pages.stop,
nvtx_source="index.prefetch_current",
nvtx_layer_id=layer_id,
nvtx_cp_rank=self.layout.cp_rank,
)
remap_ms = _cpu_timing_ms(remap_cpu)
total_ms = _cpu_timing_ms(consume_cpu)
self._log_layer(

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

@@ -10,6 +10,7 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import (
get_cp_shared_kv_local_out_cache_loc,
log_cp_draft_shared_kv_debug,
) # noqa: F401
from sglang.srt.layers.dp_attention import get_attention_cp_group
@@ -2203,18 +2204,32 @@ def current_extend_kv_rows_for_reuse(
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
valid_current_rows = sum(int(x) for x in extend_seq_lens_cpu)
if valid_current_rows <= 0:
global_current_rows = sum(int(x) for x in extend_seq_lens_cpu)
if global_current_rows <= 0:
return None
out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
if out_cache_loc is None or int(out_cache_loc.numel()) < valid_current_rows:
if out_cache_loc is None or int(out_cache_loc.numel()) < global_current_rows:
return None
tensor_rows: list[int] = []
for tensor in current_kv_tensors:
if tensor is None or int(tensor.shape[0]) < valid_current_rows:
if tensor is None:
return None
return valid_current_rows
tensor_rows.append(int(tensor.shape[0]))
if all(rows >= global_current_rows for rows in tensor_rows):
return global_current_rows
if getattr(forward_batch, "uses_cp_shared_kv", False):
local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
if local_out_cache_loc is None:
return None
local_current_rows = int(local_out_cache_loc.numel())
if local_current_rows > 0 and all(rows >= local_current_rows for rows in tensor_rows):
return local_current_rows
return None
def build_batch_prefix_slot_span(
@@ -2295,6 +2310,108 @@ def build_batch_prefix_slot_span(
return (start_slot, end_slot)
def _merge_slot_spans(spans: list[tuple[int, int]]) -> list[tuple[int, int]]:
normalized = sorted(
(int(start), int(end)) for start, end in spans if int(end) > int(start)
)
if not normalized:
return []
merged: list[tuple[int, int]] = []
cur_start, cur_end = normalized[0]
for start, end in normalized[1:]:
if start <= cur_end:
cur_end = max(cur_end, end)
else:
merged.append((cur_start, cur_end))
cur_start, cur_end = start, end
merged.append((cur_start, cur_end))
return merged
def build_batch_current_slot_spans(
*,
logical_pages: torch.Tensor,
prefix_lens_cpu,
extend_lens_cpu=None,
page_size: int,
) -> list[tuple[int, int]]:
"""Return flattened page-table slot spans that contain current extend pages."""
if prefix_lens_cpu is None:
raise ValueError("prefix_lens_cpu is required for current slot spans")
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
prefix_lens = [int(x) for x in prefix_lens_cpu]
if not prefix_lens:
return []
if logical_pages.dim() == 1:
if len(prefix_lens) != 1:
raise ValueError(
"1D logical_pages can only describe one request for current slot spans: "
f"batch_size={len(prefix_lens)} logical_pages_shape={tuple(logical_pages.shape)}"
)
pages_per_request = int(logical_pages.numel())
else:
if int(logical_pages.shape[0]) < len(prefix_lens):
raise ValueError(
"logical_pages has fewer rows than prefix_lens_cpu: "
f"rows={int(logical_pages.shape[0])} batch_size={len(prefix_lens)}"
)
pages_per_request = int(
logical_pages.reshape(logical_pages.shape[0], -1).shape[1]
)
if extend_lens_cpu is None:
extend_lens: list[int | None] = [None] * len(prefix_lens)
else:
extend_lens = [int(x) for x in extend_lens_cpu]
if len(extend_lens) != len(prefix_lens):
raise ValueError(
"extend_lens_cpu length must match prefix_lens_cpu: "
f"extend={len(extend_lens)} prefix={len(prefix_lens)}"
)
spans: list[tuple[int, int]] = []
for req_id, prefix_len in enumerate(prefix_lens):
if prefix_len < 0:
raise ValueError(
f"prefix_lens_cpu contains negative prefix length: req_id={req_id} "
f"prefix_len={prefix_len}"
)
if prefix_len % page_size != 0:
raise ValueError(
"CP shared KV current slot spans require page-aligned prefixes: "
f"req_id={req_id} prefix_len={prefix_len} page_size={page_size}"
)
prefix_pages = prefix_len // page_size
if prefix_pages > pages_per_request:
raise ValueError(
"prefix pages exceed per-request logical page-table width: "
f"req_id={req_id} prefix_pages={prefix_pages} "
f"pages_per_request={pages_per_request}"
)
extend_len = extend_lens[req_id]
if extend_len is None:
end_pages = pages_per_request
else:
if extend_len < 0:
raise ValueError(
f"extend_lens_cpu contains negative extend length: req_id={req_id} "
f"extend_len={extend_len}"
)
end_pages = (prefix_len + extend_len + page_size - 1) // page_size
end_pages = min(end_pages, pages_per_request)
if end_pages <= prefix_pages:
continue
req_start = req_id * pages_per_request
spans.append((req_start + prefix_pages, req_start + end_pages))
return _merge_slot_spans(spans)
def current_loc_remap_fast_path_args(
forward_batch,
) -> tuple[int | None, int | None]:
@@ -3285,6 +3402,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
page_size: int,
prefix_pages: int,
prefix_slot_span: tuple[int, int] | None = None,
current_slot_spans: list[tuple[int, int]] | None = None,
layer_id: int | None = None,
nvtx_source: str = "mla.partial_current_sync",
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -3379,6 +3497,35 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
page_size=page_size,
mask_non_current_in_current_pages=True,
)
if layout.cp_size > 1:
if current_slot_spans is None:
if prefix_slot_span is not None:
raise ValueError(
"CP shared KV batched current compose requires explicit "
"current_slot_spans to avoid reducing prefix slots twice."
)
current_slot_spans = (
[(int(prefix_pages), total_slots)]
if int(prefix_pages) < total_slots
else []
)
for current_start_slot, current_end_slot in _merge_slot_spans(
current_slot_spans
):
current_rows = slot_range_to_token_slice(
page_size,
current_start_slot,
current_end_slot,
)
_all_reduce_materialized_buffer_range(
mixed_kv_cache,
layout.cp_size,
current_rows.start,
current_rows.stop,
nvtx_source=f"{nvtx_source}.current",
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
return mixed_kv_cache, mixed_locs
@@ -3394,6 +3541,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
index_head_dim: int,
prefix_pages: int,
prefix_slot_span: tuple[int, int] | None = None,
current_slot_spans: list[tuple[int, int]] | None = None,
layer_id: int | None = None,
nvtx_source: str = "index.partial_current_sync",
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -3462,6 +3610,34 @@ def materialize_prefix_and_reuse_current_index_page_slots(
page_size=page_size,
index_head_dim=index_head_dim,
)
if layout.cp_size > 1:
if current_slot_spans is None:
if prefix_slot_span is not None:
raise ValueError(
"CP shared KV batched index current compose requires explicit "
"current_slot_spans to avoid reducing prefix slots twice."
)
current_slot_spans = (
[(int(prefix_pages), total_slots)]
if int(prefix_pages) < total_slots
else []
)
for current_start_slot, current_end_slot in _merge_slot_spans(
current_slot_spans
):
current_pages = slot_range_to_page_slice(
current_start_slot,
current_end_slot,
)
_all_reduce_materialized_buffer_range(
dense_page_buffer,
layout.cp_size,
current_pages.start,
current_pages.stop,
nvtx_source=f"{nvtx_source}.current",
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
return dense_page_buffer, slot_remap.dense_pages

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

@@ -16,6 +16,7 @@ from sglang.jit_kernel.fused_store_index_cache import (
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import index_buf_accessor
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_batch_current_slot_spans,
build_batch_prefix_slot_span,
cp_shared_kv_debug_enabled,
cp_shared_kv_debug_log,
@@ -524,12 +525,12 @@ class Indexer(MultiPlatformOp):
prefix_lens_valid = all(
int(prefix_len) >= 0 and int(prefix_len) % page_size == 0
for prefix_len in prefix_lens_cpu
) and any(int(prefix_len) > 0 for prefix_len in prefix_lens_cpu)
)
if not prefix_lens_valid:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
"CP shared KV index partial-current compose requires "
"positive page-aligned prefix pages. "
"CP shared KV index current compose requires page-aligned "
"prefix metadata. "
f"cp_rank={layout.cp_rank} layer_id={layer_id} "
f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
f"logical_page_table_shape={tuple(logical_page_table.shape)} "
@@ -591,6 +592,12 @@ class Indexer(MultiPlatformOp):
f"out_cache_loc_shape={tuple(current_locs.shape)}"
)
prefix_slot_span = None
current_slot_spans = build_batch_current_slot_spans(
logical_pages=logical_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
else:
@@ -655,6 +662,7 @@ class Indexer(MultiPlatformOp):
index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
prefix_pages=prefix_pages,
prefix_slot_span=prefix_slot_span,
current_slot_spans=current_slot_spans,
layer_id=layer_id,
)
)
@@ -1456,7 +1464,7 @@ class Indexer(MultiPlatformOp):
index_buffer=index_buffer,
block_tables=block_tables,
batch_indices=torch.tensor(
batch_idx_list, dtype=torch.int64, device=descriptor_device
batch_idx_list, dtype=torch.int32, device=descriptor_device
),
kv_lens=torch.tensor(
kv_lens_list, dtype=torch.int32, device=descriptor_device
@@ -1861,7 +1869,7 @@ class Indexer(MultiPlatformOp):
shared_block_tables = None
current_index_kv_for_topk = current_index_kv
current_only_batch = is_current_only_extend_batch(forward_batch)
if current_index_kv is not None and not current_only_batch:
if current_index_kv is not None:
current_index_kv_for_topk = None
shared_block_tables = metadata.get_page_table_64()
shared_index_buffer, shared_block_tables = (
@@ -1973,7 +1981,7 @@ class Indexer(MultiPlatformOp):
shared_block_tables = None
current_index_kv_for_topk = current_index_kv
current_only_batch = is_current_only_extend_batch(forward_batch)
if current_index_kv is not None and not current_only_batch:
if current_index_kv is not None:
current_index_kv_for_topk = None
shared_block_tables = metadata.get_page_table_64()
shared_index_buffer, shared_block_tables = (

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

@@ -716,9 +716,9 @@ def build_batch_page_aligned_in_seq_split_plan(
flat_segment_request_ids=flat_segment_request_ids,
flat_segment_offsets=flat_segment_offsets,
compute_padding_enabled=any(
compute_tokens != valid_tokens
for compute_tokens, valid_tokens in zip(
request_compute_padded_tokens, request_valid_padded_tokens
list(compute_split) != list(valid_split)
for compute_split, valid_split in zip(
request_compute_split_lists, request_valid_split_lists
)
),
request_valid_split_lists=request_valid_split_lists,

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

@@ -15,6 +15,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
CpSharedKVMlaPrefetcher,
)
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_batch_current_slot_spans,
build_batch_prefix_slot_span,
build_current_loc_remap,
cp_shared_kv_debug_enabled,
@@ -1971,6 +1972,16 @@ class NativeSparseAttnBackend(
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
)
current_slot_spans = build_batch_current_slot_spans(
logical_pages=metadata.real_page_table,
prefix_lens_cpu=getattr(
forward_batch, "extend_prefix_lens_cpu", None
),
extend_lens_cpu=getattr(
forward_batch, "extend_seq_lens_cpu", None
),
page_size=page_size,
)
kv_cache, page_table_1 = (
materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
@@ -1981,6 +1992,7 @@ class NativeSparseAttnBackend(
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=0,
current_slot_spans=current_slot_spans,
layer_id=layer.layer_id,
nvtx_source="mla.current_only_page_slots",
)
@@ -2078,6 +2090,12 @@ class NativeSparseAttnBackend(
f"page_size={page_size}"
)
prefix_slot_span = None
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,
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
else:
@@ -2105,6 +2123,7 @@ class NativeSparseAttnBackend(
page_size=page_size,
prefix_pages=prefix_pages,
prefix_slot_span=prefix_slot_span,
current_slot_spans=current_slot_spans,
layer_id=layer.layer_id,
)
)