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

Reuse IPC descriptors across CP shared-KV layers

Browse Source 
CP shared-KV slot remaps already have forward-batch lifetime, but the IPC materialize path rebuilt owner/source/dense descriptor tensors on every layer. Cache prefix/current IPC descriptors on the token and paged slot-remap objects, keyed by layout, spans, device, descriptor kind, and prefix capacity, so all model layers can reuse the same request/batch-plan descriptors.

Constraint: Small-extend cache-hit workloads showed descriptor setup could exceed the all-reduce baseline before any IPC kernel work ran.

Rejected: Global descriptor cache | slot-remap lifetime is safer and avoids stale entries across request/batch-plan changes.

Rejected: Cache without physical page capacity | prefix descriptors encode capacity-invalid pages and must miss when capacity changes.

Confidence: high

Scope-risk: moderate

Directive: Do not reuse descriptors across different slot_logical_pages identity, CP layout, spans, device, or prefix capacity; stale descriptors can alias dense slots across requests.

Tested: Local py_compile; local git diff --check; remote g0034 cjy-glm5-new targeted descriptor tests 2 passed; remote full test_cp_shared_kv_runtime.py 146 passed, 21 warnings, 2 subtests passed.

Not-tested: Full ETE throughput/accuracy after descriptor cache; CUDA service benchmark still needed to quantify speedup.
(cherry picked from commit addd1ca1571e41458315d15304a0e841682fe8fa)
This commit is contained in:
laoyao0822
2026-06-12 05:03:53 +08:00
parent d7eb90dff2
commit cd4412a4b8
3 changed files with 493 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2,7 +2,7 @@ from __future__ import annotations
import logging
from contextlib import contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any
@@ -196,6 +196,22 @@ def _log_current_reuse_fallback(
)
@dataclass(frozen=True)
class _IpcPrefixSlotDescriptors:
slot_indices: torch.Tensor
owner_ranks: torch.Tensor
src_page_indices: torch.Tensor
dense_page_indices: torch.Tensor
@dataclass(frozen=True)
class _IpcCurrentSlotDescriptors:
slot_indices: torch.Tensor
owner_ranks: torch.Tensor
compact_src_page_indices: torch.Tensor
dense_page_indices: torch.Tensor
@dataclass(frozen=True)
class SharedTokenKVSlotRemap:
slot_logical_pages: torch.Tensor
@@ -206,6 +222,9 @@ class SharedTokenKVSlotRemap:
slot_dense_pages: torch.Tensor | None = None
slot_sorted_logical_pages_by_row: torch.Tensor | None = None
slot_sorted_dense_pages_by_row: torch.Tensor | None = None
ipc_descriptor_cache: dict[tuple[object, ...], object] = field(
default_factory=dict, compare=False, repr=False
)
@dataclass(frozen=True)
@@ -217,6 +236,9 @@ class SharedPagedBufferSlotRemap:
dense_num_pages: int
slot_sorted_logical_pages_by_row: torch.Tensor | None = None
slot_sorted_dense_pages_by_row: torch.Tensor | None = None
ipc_descriptor_cache: dict[tuple[object, ...], object] = field(
default_factory=dict, compare=False, repr=False
)
def _tensor_identity_key(tensor: torch.Tensor) -> tuple[int, tuple[int, ...], str, str]:
@@ -2577,6 +2599,140 @@ def _build_compact_current_staging_src_page_indices(
).contiguous()
def _normalized_ipc_slot_spans_key(spans: list[tuple[int, int]]) -> tuple[tuple[int, int], ...]:
return tuple((int(start), int(end)) for start, end in _merge_slot_spans(spans))
def _ipc_slot_descriptor_cache_key(
*,
descriptor_kind: str,
cache_kind: str,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
device: torch.device,
physical_page_capacity: int | None = None,
) -> tuple[object, ...]:
return (
descriptor_kind,
cache_kind,
_tensor_identity_key(slot_logical_pages),
int(layout.page_size),
int(layout.cp_size),
int(layout.cp_rank),
_normalized_ipc_slot_spans_key(spans),
str(device),
None if physical_page_capacity is None else int(physical_page_capacity),
)
def _get_or_build_prefix_ipc_slot_descriptors(
*,
slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
device: torch.device,
physical_page_capacity: int | None,
cache_kind: str,
) -> _IpcPrefixSlotDescriptors:
slot_logical_pages = slot_remap.slot_logical_pages
key = _ipc_slot_descriptor_cache_key(
descriptor_kind="prefix",
cache_kind=cache_kind,
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=device,
physical_page_capacity=physical_page_capacity,
)
cached = slot_remap.ipc_descriptor_cache.get(key)
if cached is not None:
return cached # type: ignore[return-value]
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=device)
)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=device,
)
if slot_indices.numel() == 0:
empty = torch.empty((0,), dtype=torch.long, device=device)
descriptors = _IpcPrefixSlotDescriptors(empty, empty, empty, empty)
else:
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages.index_select(0, slot_indices)
)
owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages_range,
layout,
physical_page_capacity=physical_page_capacity,
)
descriptors = _IpcPrefixSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
src_page_indices=src_page_indices,
dense_page_indices=(slot_indices + 1).to(torch.long).contiguous(),
)
slot_remap.ipc_descriptor_cache[key] = descriptors
return descriptors
def _get_or_build_current_ipc_slot_descriptors(
*,
slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
device: torch.device,
cache_kind: str,
) -> _IpcCurrentSlotDescriptors:
slot_logical_pages = slot_remap.slot_logical_pages
key = _ipc_slot_descriptor_cache_key(
descriptor_kind="current",
cache_kind=cache_kind,
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=device,
)
cached = slot_remap.ipc_descriptor_cache.get(key)
if cached is not None:
return cached # type: ignore[return-value]
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=device)
)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=device,
)
if slot_indices.numel() == 0:
empty = torch.empty((0,), dtype=torch.long, device=device)
descriptors = _IpcCurrentSlotDescriptors(empty, empty, empty, empty)
else:
owner_ranks, src_page_indices, dense_page_indices = (
_build_current_staging_ipc_descriptors(
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=device,
slot_indices=slot_indices,
)
)
descriptors = _IpcCurrentSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
compact_src_page_indices=_build_compact_current_staging_src_page_indices(
src_page_indices
),
dense_page_indices=dense_page_indices,
)
slot_remap.ipc_descriptor_cache[key] = descriptors
return descriptors
def _try_tai_ipc_materialize_token_kv_page_slot_spans_into(
*,
kv_cache: torch.Tensor,
@@ -2585,6 +2741,7 @@ def _try_tai_ipc_materialize_token_kv_page_slot_spans_into(
layout: CpSharedKVLayout,
page_size: int,
spans: list[tuple[int, int]],
slot_remap: SharedTokenKVSlotRemap | None = None,
) -> bool:
if not spans:
return True
@@ -2607,31 +2764,49 @@ def _try_tai_ipc_materialize_token_kv_page_slot_spans_into(
return False
kernels, peer_ptrs = ipc_state
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=dense_kv_cache.device)
)
try:
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=dense_kv_cache.device,
)
if slot_indices.numel() == 0:
if slot_remap is not None:
descriptors = _get_or_build_prefix_ipc_slot_descriptors(
slot_remap=slot_remap,
layout=layout,
spans=spans,
device=dense_kv_cache.device,
physical_page_capacity=kv_cache.shape[0] // page_size,
cache_kind="token",
)
else:
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=dense_kv_cache.device)
)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=dense_kv_cache.device,
)
if slot_indices.numel() == 0:
return True
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages.index_select(0, slot_indices)
)
owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages_range,
layout,
physical_page_capacity=kv_cache.shape[0] // page_size,
)
descriptors = _IpcPrefixSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
src_page_indices=src_page_indices,
dense_page_indices=(slot_indices + 1).to(torch.long).contiguous(),
)
if descriptors.slot_indices.numel() == 0:
return True
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages.index_select(0, slot_indices)
)
owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages_range,
layout,
physical_page_capacity=kv_cache.shape[0] // page_size,
)
kernels.materialize_cuda_ipc_peer_pages_slot_indices(
peer_ptrs,
dense_kv_cache,
owner_ranks,
src_page_indices,
(slot_indices + 1).contiguous(),
descriptors.owner_ranks,
descriptors.src_page_indices,
descriptors.dense_page_indices,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
)
return True
@@ -2659,20 +2834,48 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into(
layout: CpSharedKVLayout,
page_size: int,
spans: list[tuple[int, int]],
slot_remap: SharedTokenKVSlotRemap | None = None,
) -> bool:
"""Materialize peer current KV slots through persistent IPC staging."""
if not spans:
return True
page_nbytes = _token_kv_page_nbytes(dense_kv_cache, page_size)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(slot_logical_pages.reshape(-1).numel()),
device=dense_kv_cache.device,
)
if slot_indices.numel() == 0:
if slot_remap is not None:
descriptors = _get_or_build_current_ipc_slot_descriptors(
slot_remap=slot_remap,
layout=layout,
spans=spans,
device=dense_kv_cache.device,
cache_kind="token",
)
else:
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(slot_logical_pages.reshape(-1).numel()),
device=dense_kv_cache.device,
)
if slot_indices.numel() == 0:
return True
owner_ranks, src_page_indices, dense_page_indices = (
_build_current_staging_ipc_descriptors(
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=dense_kv_cache.device,
slot_indices=slot_indices,
)
)
descriptors = _IpcCurrentSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
compact_src_page_indices=_build_compact_current_staging_src_page_indices(
src_page_indices
),
dense_page_indices=dense_page_indices,
)
if descriptors.slot_indices.numel() == 0:
return True
dense_page_indices = (slot_indices + 1).to(torch.long).contiguous()
required_nbytes = int(dense_page_indices.numel()) * int(page_nbytes)
required_nbytes = int(descriptors.dense_page_indices.numel()) * int(page_nbytes)
staging_state = _get_or_create_tai_ipc_current_staging(
kind="token",
dense_tensor=dense_kv_cache,
@@ -2683,24 +2886,12 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into(
return False
kernels, state = staging_state
try:
owner_ranks, src_page_indices, dense_page_indices = (
_build_current_staging_ipc_descriptors(
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=dense_kv_cache.device,
slot_indices=slot_indices,
)
)
compact_src_page_indices = _build_compact_current_staging_src_page_indices(
src_page_indices
)
state.ready_seq += 1
ready_seq = int(state.ready_seq)
kernels.publish_cuda_ipc_slot_pages_compact_and_mark_ready(
dense_kv_cache,
state.staging,
dense_page_indices,
descriptors.dense_page_indices,
state.ready,
ready_seq=ready_seq,
page_nbytes=page_nbytes,
@@ -2709,9 +2900,9 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into(
state.peer_ptrs,
state.ready_peer_ptrs,
dense_kv_cache,
owner_ranks,
compact_src_page_indices,
dense_page_indices,
descriptors.owner_ranks,
descriptors.compact_src_page_indices,
descriptors.dense_page_indices,
ready_seq=ready_seq,
page_nbytes=page_nbytes,
)
@@ -2815,6 +3006,7 @@ def _try_tai_ipc_materialize_paged_buffer_page_slot_spans_into(
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
slot_remap: SharedPagedBufferSlotRemap | None = None,
) -> bool:
if not spans:
return True
@@ -2837,31 +3029,49 @@ def _try_tai_ipc_materialize_paged_buffer_page_slot_spans_into(
return False
kernels, peer_ptrs = ipc_state
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=dense_page_buffer.device)
)
try:
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=dense_page_buffer.device,
)
if slot_indices.numel() == 0:
if slot_remap is not None:
descriptors = _get_or_build_prefix_ipc_slot_descriptors(
slot_remap=slot_remap,
layout=layout,
spans=spans,
device=dense_page_buffer.device,
physical_page_capacity=page_buffer.shape[0],
cache_kind="paged",
)
else:
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=dense_page_buffer.device)
)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=dense_page_buffer.device,
)
if slot_indices.numel() == 0:
return True
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages.index_select(0, slot_indices)
)
owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages_range,
layout,
physical_page_capacity=page_buffer.shape[0],
)
descriptors = _IpcPrefixSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
src_page_indices=src_page_indices,
dense_page_indices=(slot_indices + 1).to(torch.long).contiguous(),
)
if descriptors.slot_indices.numel() == 0:
return True
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages.index_select(0, slot_indices)
)
owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages_range,
layout,
physical_page_capacity=page_buffer.shape[0],
)
kernels.materialize_cuda_ipc_peer_pages_slot_indices(
peer_ptrs,
dense_page_buffer,
owner_ranks,
src_page_indices,
(slot_indices + 1).contiguous(),
descriptors.owner_ranks,
descriptors.src_page_indices,
descriptors.dense_page_indices,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
)
return True
@@ -2887,20 +3097,48 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into(
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
slot_remap: SharedPagedBufferSlotRemap | None = None,
) -> bool:
"""Materialize peer current index/page slots through persistent IPC staging."""
if not spans:
return True
page_nbytes = _page_nbytes_from_page_tensor(dense_page_buffer)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(slot_logical_pages.reshape(-1).numel()),
device=dense_page_buffer.device,
)
if slot_indices.numel() == 0:
if slot_remap is not None:
descriptors = _get_or_build_current_ipc_slot_descriptors(
slot_remap=slot_remap,
layout=layout,
spans=spans,
device=dense_page_buffer.device,
cache_kind="paged",
)
else:
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(slot_logical_pages.reshape(-1).numel()),
device=dense_page_buffer.device,
)
if slot_indices.numel() == 0:
return True
owner_ranks, src_page_indices, dense_page_indices = (
_build_current_staging_ipc_descriptors(
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=dense_page_buffer.device,
slot_indices=slot_indices,
)
)
descriptors = _IpcCurrentSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
compact_src_page_indices=_build_compact_current_staging_src_page_indices(
src_page_indices
),
dense_page_indices=dense_page_indices,
)
if descriptors.slot_indices.numel() == 0:
return True
dense_page_indices = (slot_indices + 1).to(torch.long).contiguous()
required_nbytes = int(dense_page_indices.numel()) * int(page_nbytes)
required_nbytes = int(descriptors.dense_page_indices.numel()) * int(page_nbytes)
staging_state = _get_or_create_tai_ipc_current_staging(
kind="paged",
dense_tensor=dense_page_buffer,
@@ -2911,24 +3149,12 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into(
return False
kernels, state = staging_state
try:
owner_ranks, src_page_indices, dense_page_indices = (
_build_current_staging_ipc_descriptors(
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=dense_page_buffer.device,
slot_indices=slot_indices,
)
)
compact_src_page_indices = _build_compact_current_staging_src_page_indices(
src_page_indices
)
state.ready_seq += 1
ready_seq = int(state.ready_seq)
kernels.publish_cuda_ipc_slot_pages_compact_and_mark_ready(
dense_page_buffer,
state.staging,
dense_page_indices,
descriptors.dense_page_indices,
state.ready,
ready_seq=ready_seq,
page_nbytes=page_nbytes,
@@ -2937,9 +3163,9 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into(
state.peer_ptrs,
state.ready_peer_ptrs,
dense_page_buffer,
owner_ranks,
compact_src_page_indices,
dense_page_indices,
descriptors.owner_ranks,
descriptors.compact_src_page_indices,
descriptors.dense_page_indices,
ready_seq=ready_seq,
page_nbytes=page_nbytes,
)
@@ -4942,6 +5168,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
layout=layout,
page_size=page_size,
spans=prefix_spans,
slot_remap=slot_remap,
)
if not materialized_by_ipc and prefix_spans and _should_fail_fast_tai_ipc_materialize(dense_kv_cache):
_raise_tai_ipc_materialize_required(
@@ -5022,6 +5249,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
layout=layout,
page_size=page_size,
spans=merged_current_spans_for_reduce,
slot_remap=slot_remap,
)
)
if (
@@ -5163,6 +5391,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
spans=prefix_spans,
slot_remap=slot_remap,
)
if not materialized_by_ipc and prefix_spans and _should_fail_fast_tai_ipc_materialize(dense_page_buffer):
_raise_tai_ipc_materialize_required(
@@ -5223,6 +5452,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
spans=merged_current_spans_for_reduce,
slot_remap=slot_remap,
)
)
if (