diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md index cdcdb0506..29a45329d 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md @@ -352,3 +352,91 @@ staging_nbytes >= num_slots * page_nbytes - SGLang runtime 单测覆盖 token/index current helper：publish 使用 dense destination pages，peer materialize 使用 compact source pages。 - `benchmark_cp_shared_kv_ipc_gather.py` 已同步改成 compact current staging 合同；quick smoke： `bs=2 cached=4096 extend=1024 fp8/uint8` 下 dense all_reduce p50 0.459ms，IPC compose p50 0.310ms。 + +### 2026-06-12 小 extend / bs1-2 / 200k context 性能边界 + +用户指出线上模型 context 上限约 200k token。因此早先用 `cached=307200` 得到的 scaling 结论只能说明 kernel 趋势，不能作为线上策略依据。 + +远端用 `benchmark_cp_shared_kv_ipc_gather.py --cache-hit-only` 按 `cached + extend <= 200k` 重新测了 bs=1/2、小 extend。关键结果： + +```text +bs=1 cached=65536 extend=256..4096: all_reduce p50 0.286..0.297ms, IPC p50 0.359..0.361ms => IPC 更慢 +bs=1 cached=102400 extend=256..4096: all_reduce p50 0.414..0.437ms, IPC p50 0.411..0.417ms => 基本持平/略快 +bs=1 cached=160000 extend=256..4096: all_reduce p50 0.622..0.634ms, IPC p50 0.521..0.524ms => IPC 明显更快 +bs=1 cached=190000 extend=256..4096: all_reduce p50 0.731..0.746ms, IPC p50 0.566..0.572ms => IPC 明显更快 +bs=2 cached=65536 extend=256..4096: all_reduce p50 0.518..0.544ms, IPC p50 0.455..0.477ms => IPC 更快 +bs=2 cached=102400 extend=256..4096: all_reduce p50 0.785..0.812ms, IPC p50 0.590..0.610ms => IPC 明显更快 +``` + +结论： + +1. “小 extend + bs1-2 效果差”在 **bs=1 且 cached 约 64k** 时成立；此时 dense all_reduce 本身只有约 0.29ms，IPC 的固定开销无法摊薄。 +2. **bs=1 cached 约 100k** 是近似 break-even 区间；真实 ETE 若包含 Python descriptor 构造、stream wait、index+MLA 双路径等额外开销，IPC 可能从 kernel 持平变成端到端劣化。 +3. **cached >= 160k 或 bs=2** 时，纯 materialize kernel 仍显示 IPC 优于 all_reduce。 +4. 下一步不能只看 kernel microbenchmark，需要补 runtime-style benchmark，把 descriptor 构造、Python/Torch tensor 准备、index/MLA 两条路径、stream 同步一起计入。 +5. 可能的工程策略不是简单回退，而是： + - 做 fused current fill+publish，去掉 current path 的额外 HBM copy / kernel 固定开销； + - 对 bs=1 + cached 较小的区间建立显式 cost model，必要时保留 all_reduce fast path，但必须是可解释的策略选择，不是 silent fallback。 + +### 2026-06-12 runtime-style benchmark 补充 + +TAI benchmark 已新增 runtime-overhead 模式： + +```bash +PYTHONPATH=python torchrun --standalone --nproc_per_node=8 \ + benchmark/nsa_prefill/benchmark_cp_shared_kv_ipc_gather.py \ + --runtime-overhead-only \ + --cache-hit-cached-tokens 65536 102400 160000 190000 \ + --cache-hit-extend-tokens 256 1024 4096 \ + --cache-hit-batch-requests 1 2 \ + --cache-hit-max-context-tokens 200000 \ + --dtype uint8 --kv-dim 656 --warmup 2 --repeat 5 --no-check +``` + +其中 `--cache-hit-max-context-tokens` 按 per-request context 过滤无效点，避免把超过模型 200k 上限的 benchmark 混进结论。 + +新增输出路径： + +- `cache_hit_dense_all_reduce_full`：dense all_reduce baseline。 +- `cache_hit_runtime_descriptor_setup_only`：每次 forward 重建 owner/src/dst descriptor tensor 的 Python/Torch/H2D 成本。 +- `cache_hit_runtime_ipc_prefix_current_compose`：descriptor 重建 + prefix IPC materialize + current compact publish/wait-gather。 + +关键结果： + +```text +bs=1 cached=65k: all_reduce ~0.34ms, descriptor setup ~0.46ms, runtime IPC ~0.84ms +bs=1 cached=102k: all_reduce ~0.45ms, descriptor setup ~0.61ms, runtime IPC ~1.05ms +bs=1 cached=160k: all_reduce ~0.65ms, descriptor setup ~0.87ms, runtime IPC ~1.40ms +bs=1 cached=190k: all_reduce ~0.75ms, descriptor setup ~1.00ms, runtime IPC ~1.58ms +bs=2 cached=65k: all_reduce ~0.55ms, descriptor setup ~0.74ms, runtime IPC ~1.22ms +bs=2 cached=102k: all_reduce ~0.81ms, descriptor setup ~1.05ms, runtime IPC ~1.66ms +``` + +这解释了线上“小 extend / bs1-2 / cache hit”效果差的主要来源：纯 TAI IPC kernel 在较大 prefix 下可以比 all_reduce 快，但当前 runtime 若每次 forward 都重建 descriptor tensor，固定开销已经超过 all_reduce 本身。该 benchmark 是一个上界模型：它刻意把 descriptor 构造放进 timed region，用于暴露未缓存 descriptor 的最坏 hot-path 成本。 + +后续优化优先级： + +1. descriptor 需要 request/batch-plan 级缓存或复用，不能每层/每次 materialize 从 Python list 重建 GPU tensor。 +2. current path 仍需要 fused current fill+publish，减少额外 HBM copy 和 kernel launch。 +3. 在 descriptor 缓存完成前，不应默认假设 IPC 替换 all_reduce 会改善 bs1/2 小 extend ETE；需要显式 cost model 或 gate。 + +### 2026-06-12 runtime descriptor 跨 layer 复用实现 + +已把 prefix/current IPC slot descriptor 缓存在 `SharedTokenKVSlotRemap` 与 `SharedPagedBufferSlotRemap` 上，生命周期与 forward batch 的 slot remap 一致。这样同一个 request/batch-plan 在多层 forward 中不会每层重复构造 `slot_indices / owner_ranks / src_page_indices / dense_page_indices` GPU tensor。 + +缓存 key 包含：descriptor 类型（prefix/current）、cache 类型（token/paged）、`slot_logical_pages` storage identity、CP layout、merged slot spans、device，以及 prefix 的 physical page capacity。capacity 被纳入 key 是为了避免 host/L1 capacity 变化时错误复用已经标 invalid 的 source page descriptor。 + +运行时路径： + +- MLA KV prefix：`materialize_prefix_and_reuse_current_kv_page_slots` -> `_get_or_build_prefix_ipc_slot_descriptors`。 +- MLA KV current：同一函数的 current staging IPC 路径 -> `_get_or_build_current_ipc_slot_descriptors`。 +- index prefix/current：`materialize_prefix_and_reuse_current_index_page_slots` 走同一套 descriptor cache，但用 `cache_kind="paged"` 与 token 描述符隔离。 + +当前实现只复用 Python/Torch descriptor tensor，不改变 TAI kernel 合同，也不新增 collective。后续仍需要 fused current fill+publish 来降低 current path 的额外 HBM copy / kernel launch。 + +验证： + +```text +远端 cjy-glm5-new：PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +146 passed, 21 warnings, 2 subtests passed +``` diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 0ccaa6514..c1fd9ab74 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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 ( diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 4400375b2..bf8c25829 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -3832,6 +3832,92 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ], ) + def test_ipc_prefix_descriptors_are_cached_on_token_slot_remap(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0) + kv_cache = torch.zeros((64, 1, 1), dtype=torch.float32) + remap = runtime.build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=None, + remap_logical_pages=torch.tensor([[1, 2, 3, 4]], dtype=torch.int64), + layout=layout, + page_size=4, + ) + + first = runtime._get_or_build_prefix_ipc_slot_descriptors( + slot_remap=remap, + layout=layout, + spans=[(0, 4)], + device=torch.device("cpu"), + physical_page_capacity=16, + cache_kind="token", + ) + second = runtime._get_or_build_prefix_ipc_slot_descriptors( + slot_remap=remap, + layout=layout, + spans=[(0, 4)], + device=torch.device("cpu"), + physical_page_capacity=16, + cache_kind="token", + ) + capacity_miss = runtime._get_or_build_prefix_ipc_slot_descriptors( + slot_remap=remap, + layout=layout, + spans=[(0, 4)], + device=torch.device("cpu"), + physical_page_capacity=2, + cache_kind="token", + ) + + self.assertIs(first.slot_indices, second.slot_indices) + self.assertIs(first.owner_ranks, second.owner_ranks) + self.assertIs(first.src_page_indices, second.src_page_indices) + self.assertIs(first.dense_page_indices, second.dense_page_indices) + self.assertIsNot(first.owner_ranks, capacity_miss.owner_ranks) + self.assertEqual(first.owner_ranks.tolist(), [0, 1, 0, 1]) + self.assertEqual(first.src_page_indices.tolist(), [1, 1, 2, 2]) + self.assertEqual(first.dense_page_indices.tolist(), [1, 2, 3, 4]) + + def test_ipc_current_descriptors_are_cached_on_token_slot_remap(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0) + kv_cache = torch.zeros((64, 1, 1), dtype=torch.float32) + remap = runtime.build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=None, + remap_logical_pages=torch.tensor([[1, 2, 3, 4]], dtype=torch.int64), + layout=layout, + page_size=4, + ) + + first = runtime._get_or_build_current_ipc_slot_descriptors( + slot_remap=remap, + layout=layout, + spans=[(2, 4)], + device=torch.device("cpu"), + cache_kind="token", + ) + second = runtime._get_or_build_current_ipc_slot_descriptors( + slot_remap=remap, + layout=layout, + spans=[(2, 4)], + device=torch.device("cpu"), + cache_kind="token", + ) + + self.assertIs(first.slot_indices, second.slot_indices) + self.assertIs(first.owner_ranks, second.owner_ranks) + self.assertIs(first.compact_src_page_indices, second.compact_src_page_indices) + self.assertIs(first.dense_page_indices, second.dense_page_indices) + self.assertEqual(first.slot_indices.tolist(), [2, 3]) + self.assertEqual(first.owner_ranks.tolist(), [0, 1]) + self.assertEqual(first.compact_src_page_indices.tolist(), [0, 1]) + self.assertEqual(first.dense_page_indices.tolist(), [3, 4]) + def test_forward_batch_token_slot_remap_is_cached_across_layers(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout