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

Keep CP batch slot descriptors on IPC-compatible runtime

Browse Source 
The syh rebase kept callers that expect reusable batch slot spans, while the restored CUDA IPC runtime lacked the helper and indexer still passed a symm-only writer-rank argument. Restore the batch-scoped span cache and remove the stale symm argument so the production compose path stays on CUDA IPC with exact per-request spans.\n\nConstraint: Production main-stream compose should use CUDA IPC, not symm writer-rank routing.\nRejected: Re-enable symm writer-rank parameters | benchmark showed no main-stream win and callers fail against IPC runtime contracts.\nConfidence: high\nScope-risk: narrow\nDirective: Keep slot-span metadata batch-scoped; do not rebuild Python span descriptors per layer without benchmarking.\nTested: Local py_compile for cp_shared_kv_runtime.py, cp_shared_kv_prefetch.py, nsa_indexer.py, nsa_backend.py.\nTested: Remote cjy-glm5-new pytest test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 157 passed, 2 subtests passed.\nNot-tested: Full ETE prefill/decode runtime after restarting services.
This commit is contained in:
laoyao0822
2026-06-13 01:15:33 +08:00
parent e0c388076e
commit 254d667853
3 changed files with 90 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3885,6 +3885,62 @@ def build_batch_current_slot_spans(
return _merge_slot_spans(spans)
def get_or_build_batch_slot_spans(
forward_batch: Any,
*,
logical_pages: torch.Tensor,
prefix_lens_cpu,
extend_lens_cpu=None,
page_size: int,
want_prefix: bool,
) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:
"""Cache exact flattened page-table spans for batched prefix/current pages.
The spans are batch metadata: they depend on the logical page table and the
per-request prefix/extend lengths, not on layer contents. Reusing them
across layers avoids rebuilding Python descriptor lists in the shared-KV
compose path while preserving per-request row boundaries.
"""
prefix_lens_key = tuple(int(x) for x in prefix_lens_cpu)
extend_lens_key = (
None
if extend_lens_cpu is None
else tuple(int(x) for x in extend_lens_cpu)
)
key = (
_tensor_identity_key(logical_pages),
prefix_lens_key,
extend_lens_key,
int(page_size),
bool(want_prefix),
)
cached_key = getattr(forward_batch, "cp_shared_kv_batch_slot_spans_key", None)
cached = getattr(forward_batch, "cp_shared_kv_batch_slot_spans", None)
if cached is not None and cached_key == key:
return cached
prefix_slot_spans = (
build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens_cpu,
page_size=page_size,
)
if want_prefix
else []
)
current_slot_spans = build_batch_current_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
cached = (prefix_slot_spans, current_slot_spans)
forward_batch.cp_shared_kv_batch_slot_spans_key = key
forward_batch.cp_shared_kv_batch_slot_spans = cached
return cached
def current_loc_remap_fast_path_args(
forward_batch,
) -> tuple[int | None, int | None]:

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

@@ -31,7 +31,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
log_cp_draft_shared_kv_debug,
materialize_prefix_and_reuse_current_index_page_slots,
materialize_shared_paged_buffer,
maybe_build_current_page_writer_ranks,
should_reuse_current_extend_kv,
tensor_debug_checksum,
tensor_debug_summary,
@@ -825,13 +824,6 @@ class Indexer(MultiPlatformOp):
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
layer_id=layer_id,
current_page_writer_ranks=maybe_build_current_page_writer_ranks(
forward_batch=forward_batch,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
layout=layout,
),
)
)
if (