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

Preserve CP ragged topk coordinates under batch planning

Browse Source 
Batch-planning metadata can now be attached even for single-request CP prefill, which routed FP8 flashmla_sparse through the batch cp_index path. That path used compact MQA-buffer row bases for score lookup but did not override the final ragged topk coordinate base consumed by attention, so topk indices could point at the wrong KV rows and produce low accept length or meaningless output.\n\nThis keeps ordinary long page tails out of compute padding, reserves compute padding for truly tiny suffixes, and makes cp_index RAGGED topk emit request-ragged offsets while preserving the compact buffer descriptors used for score materialization. The debug ledger records the rejected intermediate diagnoses and the confirmed coordinate-space failure.\n\nConstraint: CP shared-KV cache residency is page-granular, but attention/index compute must not consume synthetic long-tail rows.\nConstraint: FP8 CP prefill uses flashmla_sparse/RAGGED, where fused topk output is consumed directly as attention page_table_1.\nRejected: Disable current reuse or batch planning | would hide the regression and lose the intended bs>1 fast path.\nRejected: Treat all page tails as compute padding | regresses bs=1 semantics and can corrupt query/topk row contracts.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not change cp_index RAGGED topk offset handling without verifying score-buffer row bases and final attention KV coordinate space independently.\nTested: python -m py_compile on touched Python/test files; git diff --check; remote targeted ragged cp_index offset regression test; remote test_nsa_cp_utils.py; remote test_cp_shared_kv_runtime.py; user-reported ETE output recovered after restart.\nNot-tested: Agent-driven full ETE traffic run; broad multi-request bs>1 production soak.
This commit is contained in:
laoyao0822
2026-06-05 07:45:29 +08:00
parent c8510938b9
commit 5eecc5f9c5
5 changed files with 316 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1416,6 +1416,12 @@ class Indexer(MultiPlatformOp):
q_lens_list = []
k_bases_list = []
q_bases_list = []
topk_offset_list = []
request_kv_bases: List[int] = []
request_kv_base = 0
for seq_len in forward_batch.seq_lens_cpu.tolist():
request_kv_bases.append(int(request_kv_base))
request_kv_base += int(seq_len)
k_cursor = 0
q_cursor = 0
for raw_batch_idx, start_seq_position, end_seq_position in cp_index:
@@ -1453,9 +1459,22 @@ class Indexer(MultiPlatformOp):
q_lens_list.append(extend_seq_len)
k_bases_list.append(k_cursor)
q_bases_list.append(q_cursor)
if batch_idx < 0 or batch_idx >= len(request_kv_bases):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_cp_index_bad_batch_idx "
f"batch_idx={batch_idx} seq_lens={forward_batch.seq_lens_cpu.tolist()}"
)
topk_offset_list.extend(
[request_kv_bases[batch_idx]] * extend_seq_len
)
k_cursor += kv_len_i
q_cursor += extend_seq_len
topk_indices_offset_override = torch.tensor(
topk_offset_list, dtype=torch.int32, device=q_fp8.device
)
if current_index_kv is None:
assert index_buffer is not None
assert block_tables is not None
@@ -1512,6 +1531,7 @@ class Indexer(MultiPlatformOp):
cu_seqlens_q=actual_seq_q,
ke_offset=ke_offset,
batch_idx_list=batch_idx_list,
topk_indices_offset_override=topk_indices_offset_override,
)
return topk_result
else:
@@ -1576,6 +1596,7 @@ class Indexer(MultiPlatformOp):
cu_seqlens_q=actual_seq_q,
ke_offset=ke_offset,
batch_idx_list=batch_idx_list,
topk_indices_offset_override=topk_indices_offset_override,
)
return topk_result
@@ -1670,6 +1691,7 @@ class Indexer(MultiPlatformOp):
cu_seqlens_q=actual_seq_q,
ke_offset=ke_offset,
batch_idx_list=batch_idx_list,
topk_indices_offset_override=topk_indices_offset_override,
)
else:
seq_len = int(forward_batch.seq_lens_cpu[batch_idx].item())

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

@@ -610,15 +610,28 @@ def build_batch_page_aligned_in_seq_split_plan(
f"extend_len={extend_len} page_size={page_size}"
)
compute_pages = max(split_info.extend_padded_pages, cp_size)
compute_split_list, compute_page_starts, compute_page_ends = (
_build_full_page_unit_split(
page_units=compute_pages,
extend_prefix_len=prefix_len,
page_size=page_size,
cp_size=cp_size,
# Cache residency is page-granular, but attention/index compute should
# not consume synthetic rows for an ordinary final partial page. Only
# tiny suffixes with fewer physical pages than CP ranks need compute
# padding; otherwise the valid split is already sufficient and keeps
# the pre-bs>1 bs=1 semantics.
needs_compute_padding = split_info.extend_padded_pages < cp_size
if needs_compute_padding:
compute_pages = cp_size
compute_split_list, compute_page_starts, compute_page_ends = (
_build_full_page_unit_split(
page_units=compute_pages,
extend_prefix_len=prefix_len,
page_size=page_size,
cp_size=cp_size,
)
)
)
else:
compute_pages = split_info.extend_padded_pages
compute_split_list = split_list
compute_page_starts = split_info.segment_page_starts
compute_page_ends = split_info.segment_page_ends
compute_tokens = sum(int(token_count) for token_count in compute_split_list)
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
split_list=compute_split_list,
cp_size=cp_size,
@@ -671,8 +684,8 @@ def build_batch_page_aligned_in_seq_split_plan(
request_compute_segment_page_starts.append(compute_page_starts)
request_compute_segment_page_ends.append(compute_page_ends)
request_compute_padded_pages.append(compute_pages)
request_compute_padded_tokens.append(compute_pages * page_size)
request_compute_padding_tokens.append(compute_pages * page_size - extend_len)
request_compute_padded_tokens.append(compute_tokens)
request_compute_padding_tokens.append(compute_tokens - extend_len)
request_compute_rank_local_tokens.append(compute_rank_local_tokens)
request_compute_actual_seq_q_prev.append(compute_split_list[cp_rank])
request_compute_actual_seq_q_next.append(compute_split_list[mirror_idx])