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

Avoid redundant CP collectives on sync shared-KV materialize

Browse Source 
Synchronous CP shared KV full-hit and partial-current paths can now use the tai-kernel CUDA IPC slot-dense materialize path instead of first copying local owner pages and then running a dense CP all-reduce. This keeps the async prefetch pipeline unchanged while routing the safer synchronous runtime path through descriptor-driven owner-page reads.

The runtime builds descriptors from the existing page-aligned slot contract, caches peer pointer tables for long-lived KV/index buffers, and falls back with explicit warnings if the tai-kernel IPC capability is missing. Unit coverage locks offset handle exchange, descriptor construction, and both MLA/index full and partial-current calls.

Constraint: Async prefetch currently has unresolved scheduling contention with GEMM/MoE and current-layer KV collectives, so this commit intentionally does not wire IPC into prefetch.

Constraint: tai-kernel must provide offset-aware CUDA IPC symbols from af9fb67.

Rejected: Replace prefetch all-reduce in this slice | Nsight shows prefetch timing and communicator contention need a separate scheduling design.

Rejected: Fail-fast on missing tai IPC immediately | remote ETE still needs to validate production deployment capability before removing the warning fallback.

Confidence: medium

Scope-risk: moderate

Directive: Keep prefetch and synchronous materialize decisions separate until prefetch has a low-SM/copy-engine schedule.

Related: tai-kernel af9fb67

Tested: git diff --check; remote runtime/unit evidence recorded in docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Not-tested: Fresh GLM5 ETE after this commit; async prefetch IPC path.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-06-01 00:00:15 +08:00
co-authored by OmX
parent b4f5c3bc0e
commit 4342de0463
3 changed files with 844 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
+202
View File
@@ -4088,3 +4088,205 @@ Conclusion:
- The segment-copy path is now slower than tuned direct peer-page gather for this
benchmark. Do not prefer owner-concat segment copy solely for performance
unless a future consumer can avoid the final rewrite entirely.
### C99 — 2026-05-31 multi-warp/page IPC materialize tuning
Change:
- Added a benchmark-selectable `warps_per_page` mode to tai-kernel
`materialize_cuda_ipc_peer_pages_slot_dense`.
- `warps_per_page=1` preserves the previous one-warp-per-page mapping.
- `warps_per_page>1` assigns a whole warp group to one page and stripes the
page copy across `warps_per_page * 32` threads.
- Python wrapper default is now `warps_per_page=0`, an auto selector based on
the measured H200 regions for the production ~72 KiB BF16 MLA page:
- <= 256 pages: 8 warps/page,
- 512--1024 pages: 1 warp/page,
- >= 1536 pages: 4 warps/page,
- synthetic pages smaller than 64 KiB: 1 warp/page.
Remote correctness:
- `PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_cuda_ipc_gather.py`
on `g0034` passed: `6 passed`.
Remote benchmark logs:
- Variant sweep:
`/mnt/beegfs/cjy/log/cp_shared_kv_ipc_warps_per_page_20260531_114147.log`
- Auto selector sweep:
`/mnt/beegfs/cjy/log/cp_shared_kv_ipc_auto_wpp_20260531_114447.log`
Auto selector p50 GPU ms, BF16 MLA page payload, 8 CP ranks:
| tokens | pages | dense all-reduce | IPC direct slot-dense auto | speedup vs dense |
| ---: | ---: | ---: | ---: | ---: |
| 4096 | 64 | 0.093 | 0.057 | 1.63x |
| 8192 | 128 | 0.115 | 0.060 | 1.92x |
| 16384 | 256 | 0.164 | 0.083 | 1.98x |
| 32768 | 512 | 0.235 | 0.134 | 1.75x |
| 65536 | 1024 | 0.394 | 0.236 | 1.67x |
| 98304 | 1536 | 0.554 | 0.348 | 1.59x |
| 122880 | 1920 | 0.670 | 0.421 | 1.59x |
Bandwidth reading for the largest case:
- `122880` tokens = `1920` pages = `135 MiB` logical payload per GPU.
- `0.421 ms` gives about `336 GB/s` payload rate per GPU.
- The kernel performs roughly one full payload read plus one full payload write,
so effective HBM traffic is about `672 GB/s`, ~14% of H200's ~4.8 TB/s HBM
bandwidth.
- Remote NVLink receive traffic is about `7/8` of the payload, or `294 GB/s`,
~61% of the per-direction 18-link H200 NVLink budget (`18 * 26.562 GB/s ~= 478 GB/s`).
Conclusion:
- This kernel is no longer obviously HBM-limited; NVLink plus per-page peer-load
overhead is closer to the limiting side for large prefixes.
- The remaining large-prefix headroom to the per-direction NVLink roof is about
`478 / 294 ~= 1.6x` under this access pattern.
- For small prefixes, the main gain came from increasing per-page parallelism
rather than more CTAs; launch/fixed overhead still dominates below ~8k tokens.
Additional C99 correctness check:
- Ran the multi-rank benchmark with checks enabled for `4096` and `122880`
tokens and `warps_per_page=0`; all benchmark reference assertions passed.
### C100 — 2026-05-31 vectorized 16B copy for multi-warp/page IPC materialize
Change:
- The multi-warp/page IPC materialize path now uses `ld.global.nc.v2.u64` +
`st.global.cg.v2.u64` when the page payload is 16-byte aligned.
- The one-warp/page legacy path is unchanged; the vectorized helper only affects
the measured multi-warp variants used by the auto selector.
Remote verification:
- `PYTHONPATH=python python -m pytest -q tests/nsa_prefill/test_cuda_ipc_gather.py`
on `g0034` passed: `6 passed`.
- Multi-rank benchmark checks enabled for `4096` and `122880` tokens passed.
Remote benchmark log:
- `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_auto_wpp_vec16_20260531_114903.log`
Auto selector p50 GPU ms after vectorization:
| tokens | pages | dense all-reduce | IPC direct slot-dense auto | previous auto IPC | delta |
| ---: | ---: | ---: | ---: | ---: | ---: |
| 4096 | 64 | 0.137 | 0.042 | 0.057 | -26% |
| 8192 | 128 | 0.112 | 0.052 | 0.060 | -13% |
| 16384 | 256 | 0.162 | 0.082 | 0.083 | -1% |
| 32768 | 512 | 0.236 | 0.136 | 0.134 | +1% |
| 65536 | 1024 | 0.392 | 0.238 | 0.236 | +1% |
| 98304 | 1536 | 0.553 | 0.339 | 0.348 | -3% |
| 122880 | 1920 | 0.668 | 0.416 | 0.421 | -1% |
Bandwidth reading for `122880` tokens:
- Payload rate: `135 MiB / 0.416 ms ~= 340 GB/s` per GPU.
- Effective HBM read+write traffic: `~680 GB/s`, still only ~14% of H200 HBM.
- Remote NVLink receive traffic: `(7/8) * 340 ~= 298 GB/s`, ~62% of the
`~478 GB/s` per-direction H200 NVLink link budget.
Conclusion:
- 16B vectorization materially improves the short-prefix auto path and slightly
improves the large-prefix path; it does not change the bottleneck conclusion.
- Large-prefix headroom remains primarily in reducing peer-load overhead / better
NVLink utilization, not in HBM bandwidth.
### C101 — 2026-05-31 index page payload also benefits from IPC materialize
Clarification:
- The descriptor builder needed by SGLang runtime is the translation from a
consumer slot page table to the IPC kernel descriptors:
`slot_logical_pages -> owner_ranks, src_page_indices`.
- The IPC kernel itself does not understand radix/HiCache/current-reuse state;
it only consumes owner/source descriptors. Full hit, multi-extend hit, and
partial-current prefix reuse are all covered if the runtime builds these
descriptors from the same page-aligned slot contract.
Index payload benchmark:
- NSA index page buffer is `uint8` with per-token bytes
`index_head_dim + index_head_dim / quant_block_size * 4 = 128 + 4 = 132`.
- With `page_size=64`, each index page is `8448` bytes.
- Benchmark command used `--dtype uint8 --kv-dim 132` to match the real index
page payload size and dtype.
Remote log:
- `/mnt/beegfs/cjy/log/cp_shared_kv_ipc_index_uint8_20260531_120416.log`
p50 GPU ms, 8 CP ranks:
| tokens | pages | dense uint8 all-reduce | IPC direct slot-dense | speedup |
| ---: | ---: | ---: | ---: | ---: |
| 4096 | 64 | 0.119 | 0.033 | 3.61x |
| 8192 | 128 | 0.070 | 0.031 | 2.26x |
| 16384 | 256 | 0.070 | 0.033 | 2.12x |
| 32768 | 512 | 0.081 | 0.037 | 2.19x |
| 65536 | 1024 | 0.102 | 0.049 | 2.08x |
| 98304 | 1536 | 0.120 | 0.061 | 1.97x |
| 122880 | 1920 | 0.139 | 0.072 | 1.93x |
Additional check:
- Multi-rank correctness benchmark with checks enabled for `4096` and `122880`
index payloads passed.
Conclusion:
- IPC materialize is beneficial for index pages too. Because index pages are
much smaller than MLA pages, the absolute saving is smaller, but the relative
saving versus uint8 all-reduce is still about `~2x` for large prefixes and
stronger for short prefixes.
- This supports wiring both MLA and index full/partial prefix materialize to the
same descriptor-driven IPC path.
### C102 — 2026-05-31 runtime IPC materialize is wired only for synchronous non-prefetch paths
Implementation note:
- The first SGLang runtime hookup targets only the synchronous materialize paths:
- MLA full materialize with a precomputed slot remap,
- MLA partial-current synchronous prefix materialize,
- index full materialize with a precomputed slot remap,
- index partial-current synchronous prefix materialize.
- `CpSharedKVMlaPrefetcher` and `CpSharedKVIndexPrefetcher` are intentionally
unchanged in this slice. Async prefetch still uses the existing local
materialize + async collective path until stream-ordering and CPU-submit cost
are re-evaluated separately.
- The runtime descriptor contract is:
`slot_logical_pages -> owner_ranks, src_page_indices`, where source page index
is the peer CP-local physical page id including dummy page 0. Invalid slots
use `-1/-1` and are zero-filled by the tai-kernel IPC slot-dense materialize
kernel.
- CUDA IPC handle exchange is cached per long-lived source tensor pointer/shape
and CP group. The cache intentionally avoids per-layer handle collectives;
a tiny handle/offset all-gather is still required the first time a KV/index
source buffer is registered on a rank.
Important correction:
- Runtime IPC must carry CUDA allocation offsets. Production source tensors are
PyTorch caching-allocator tensors, not guaranteed `cudaMalloc` base pointers.
`tai-kernel` now exposes `get_cuda_ipc_mem_handle_with_offset()` and
`open_cuda_ipc_mem_handles_with_offsets()` so peers dereference the exact
tensor data pointer instead of the allocation base.
Risk / not yet covered:
- Peer pointers are cached for long-lived KV/index buffers and are not currently
closed before process exit. This is acceptable for the scheduler lifetime but
should not be reused for short-lived staging tensors.
- The IPC path currently falls back to the old collective materialize path if the
tai-kernel IPC capability is missing or peer opening fails; the fallback logs
`[CP_SHARED_KV_FALLBACK][tai_ipc_materialize]`. Once remote ETE validates the
capability, this can be tightened to fail-fast for the production GLM5 config.
- Remote CUDA verification is still required after syncing both `tai-kernel` and
`sglang-dev`; local CPU import is blocked by missing optional dependencies.
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+421 -50
View File
@@ -19,9 +19,11 @@ logger = logging.getLogger(__name__)
_DEBUG_LOG_COUNTS: dict[str, int] = {}
_TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_TAI_IPC_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_CURRENT_REUSE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
_SLOT_REMAP_CACHE_LOG_COUNTS: dict[str, int] = {}
_TAI_IPC_PEER_PTR_CACHE: dict[tuple[object, ...], torch.Tensor] = {}
_MLA_PREFETCH_LOG_PROBE_LAYER = 2
_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS = max(
int(envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS.get()),
@@ -300,6 +302,38 @@ def _load_tai_materialize_kernels():
return None
@lru_cache(maxsize=1)
def _load_tai_ipc_kernels():
try:
from tai_kernel.nsa_prefill import ipc as tai_ipc
required = (
"get_cuda_ipc_mem_handle_with_offset",
"open_cuda_ipc_mem_handles_with_offsets",
"materialize_cuda_ipc_peer_pages_slot_dense",
)
missing = [name for name in required if not hasattr(tai_ipc, name)]
if missing:
_log_tai_ipc_materialize_fallback(
"missing_symbols",
"CP shared KV tai IPC materialize kernels are incomplete; "
"falling back to collective materialize. missing=%s",
missing,
limit=1,
)
return None
return tai_ipc
except Exception as exc:
_log_tai_ipc_materialize_fallback(
"import_failed",
"CP shared KV tai IPC materialize kernels are unavailable; "
"falling back to collective materialize. error=%s",
exc,
limit=1,
)
return None
@lru_cache(maxsize=16)
def _tai_current_slot_fill_sparse_pages_self_test(
device_type: str,
@@ -480,6 +514,23 @@ def _log_tai_materialize_fallback(
)
def _log_tai_ipc_materialize_fallback(
key: str,
message: str,
*args,
limit: int = 8,
) -> None:
count = _TAI_IPC_MATERIALIZE_FALLBACK_LOG_COUNTS.get(key, 0)
if count >= limit:
return
_TAI_IPC_MATERIALIZE_FALLBACK_LOG_COUNTS[key] = count + 1
logger.warning(
"[CP_SHARED_KV_FALLBACK][tai_ipc_materialize] reason=%s " + message,
key,
*args,
)
def _log_tai_fused_mla_store_fallback(
key: str,
message: str,
@@ -1120,6 +1171,266 @@ def _try_tai_materialize_token_kv_page_slots_into(
return False
def build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
*,
physical_page_capacity: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build owner/source descriptors for slot-dense CUDA IPC materialize.
The returned vectors are flattened in consumer slot order. Source page
indices are physical page ids in the peer's CP-local pool, including dummy
page 0. Invalid/sentinel slots are encoded as ``-1`` owner/source and are
zero-filled by the IPC kernel.
"""
logical_pages = slot_logical_pages.reshape(-1).to(torch.long)
owner_ranks = layout.owner_for_logical_pages(logical_pages).to(torch.long)
src_page_indices = layout.logical_pages_to_physical(logical_pages).to(torch.long)
invalid = logical_pages <= 0
if physical_page_capacity is not None:
invalid = invalid | (src_page_indices >= int(physical_page_capacity))
invalid_value = torch.full_like(owner_ranks, -1)
owner_ranks = torch.where(invalid, invalid_value, owner_ranks)
src_page_indices = torch.where(invalid, invalid_value, src_page_indices)
return owner_ranks.contiguous(), src_page_indices.contiguous()
def _tai_ipc_peer_ptr_cache_key(
tensor: torch.Tensor,
layout: CpSharedKVLayout,
cp_group: Any,
) -> tuple[object, ...]:
return (
int(tensor.untyped_storage().data_ptr()),
int(tensor.data_ptr()),
tuple(int(dim) for dim in tensor.shape),
tuple(int(stride) for stride in tensor.stride()),
str(tensor.dtype),
str(tensor.device),
int(layout.cp_size),
int(layout.cp_rank),
getattr(cp_group, "unique_name", None),
)
def _get_or_open_tai_ipc_peer_ptrs(
tensor: torch.Tensor,
layout: CpSharedKVLayout,
) -> tuple[Any, torch.Tensor] | None:
if layout.cp_size <= 1:
return None
if not _tai_materialize_runtime_enabled():
return None
if not tensor.is_cuda or not tensor.is_contiguous():
_log_tai_ipc_materialize_fallback(
"unsupported_tensor",
"CP shared KV tai IPC materialize requires a contiguous CUDA tensor; "
"falling back to collective materialize. device=%s contiguous=%s",
tensor.device,
tensor.is_contiguous(),
)
return None
kernels = _load_tai_ipc_kernels()
if kernels is None:
return None
cp_group = get_attention_cp_group()
if int(getattr(cp_group, "world_size", layout.cp_size)) != int(layout.cp_size):
_log_tai_ipc_materialize_fallback(
"group_size_mismatch",
"CP shared KV tai IPC materialize cp group size mismatch; "
"falling back to collective materialize. layout_cp_size=%s group_world_size=%s",
layout.cp_size,
getattr(cp_group, "world_size", None),
limit=1,
)
return None
cache_key = _tai_ipc_peer_ptr_cache_key(tensor, layout, cp_group)
cached = _TAI_IPC_PEER_PTR_CACHE.get(cache_key)
if cached is not None:
return kernels, cached
try:
handle, offset = kernels.get_cuda_ipc_mem_handle_with_offset(tensor)
handle_cuda = handle.to(device=tensor.device, non_blocking=False)
offset_cuda = torch.tensor(
[int(offset)], dtype=torch.int64, device=tensor.device
)
gathered_handles = torch.empty(
(layout.cp_size, int(handle.numel())),
dtype=torch.uint8,
device=tensor.device,
)
gathered_offsets = torch.empty(
(layout.cp_size,),
dtype=torch.int64,
device=tensor.device,
)
cp_group.all_gather_into_tensor(gathered_handles, handle_cuda)
cp_group.all_gather_into_tensor(gathered_offsets, offset_cuda)
peer_ptrs = kernels.open_cuda_ipc_mem_handles_with_offsets(
gathered_handles.cpu(),
gathered_offsets.cpu(),
tensor,
self_rank=layout.cp_rank,
)
_TAI_IPC_PEER_PTR_CACHE[cache_key] = peer_ptrs
return kernels, peer_ptrs
except Exception as exc:
_log_tai_ipc_materialize_fallback(
"peer_open_failed",
"CP shared KV tai IPC peer pointer setup failed; falling back to "
"collective materialize. cp_rank=%s cp_size=%s tensor_shape=%s "
"dtype=%s error=%s",
layout.cp_rank,
layout.cp_size,
tuple(tensor.shape),
tensor.dtype,
exc,
limit=4,
)
return None
def _page_nbytes_from_page_tensor(tensor: torch.Tensor) -> int:
if tensor.ndim == 0:
raise ValueError("CP shared KV IPC page tensor must have at least one dim")
return int(tensor.stride(0)) * int(tensor.element_size())
def _token_kv_page_nbytes(kv_cache: torch.Tensor, page_size: int) -> int:
return int(page_size) * _page_nbytes_from_page_tensor(kv_cache)
def _try_tai_ipc_materialize_token_kv_page_slots_into(
*,
kv_cache: torch.Tensor,
dense_kv_cache: torch.Tensor,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
page_size: int,
start_slot: int,
end_slot: int,
) -> bool:
if start_slot == end_slot:
return True
if start_slot != 0:
return False
if not dense_kv_cache.is_cuda or not dense_kv_cache.is_contiguous():
return False
ipc_state = _get_or_open_tai_ipc_peer_ptrs(kv_cache, layout)
if ipc_state is None:
return False
kernels, peer_ptrs = ipc_state
flat_slot_logical_pages = slot_logical_pages.reshape(-1)
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages[start_slot:end_slot]
)
if slot_logical_pages_range.numel() == 0:
return True
try:
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_dense(
peer_ptrs,
dense_kv_cache,
owner_ranks,
src_page_indices,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
)
return True
except Exception as exc:
_log_tai_ipc_materialize_fallback(
"token_kernel_failed",
"CP shared KV tai IPC token materialize failed; falling back to "
"collective materialize. cp_rank=%s cp_size=%s start_slot=%s "
"end_slot=%s slot_count=%s page_size=%s kv_shape=%s dense_shape=%s "
"error=%s",
layout.cp_rank,
layout.cp_size,
start_slot,
end_slot,
int(slot_logical_pages_range.numel()),
page_size,
tuple(kv_cache.shape),
tuple(dense_kv_cache.shape),
exc,
)
return False
def _try_tai_ipc_materialize_paged_buffer_page_slots_into(
*,
page_buffer: torch.Tensor,
dense_page_buffer: torch.Tensor,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
start_slot: int,
end_slot: int,
) -> bool:
if start_slot == end_slot:
return True
if start_slot != 0:
return False
if not dense_page_buffer.is_cuda or not dense_page_buffer.is_contiguous():
return False
ipc_state = _get_or_open_tai_ipc_peer_ptrs(page_buffer, layout)
if ipc_state is None:
return False
kernels, peer_ptrs = ipc_state
flat_slot_logical_pages = slot_logical_pages.reshape(-1)
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages[start_slot:end_slot]
)
if slot_logical_pages_range.numel() == 0:
return True
try:
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_dense(
peer_ptrs,
dense_page_buffer,
owner_ranks,
src_page_indices,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
)
return True
except Exception as exc:
_log_tai_ipc_materialize_fallback(
"paged_kernel_failed",
"CP shared KV tai IPC paged materialize failed; falling back to "
"collective materialize. cp_rank=%s cp_size=%s start_slot=%s "
"end_slot=%s slot_count=%s page_shape=%s dense_shape=%s error=%s",
layout.cp_rank,
layout.cp_size,
start_slot,
end_slot,
int(slot_logical_pages_range.numel()),
tuple(page_buffer.shape),
tuple(dense_page_buffer.shape),
exc,
)
return False
def is_current_only_extend_batch(forward_batch) -> bool:
"""Return whether an extend batch has no cached/history tokens.
@@ -2256,7 +2567,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
dense_kv_cache = kv_cache.new_zeros(
(slot_remap.dense_num_pages * page_size, *kv_cache.shape[1:])
)
materialize_local_token_kv_page_slots_into(
materialized_by_ipc = _try_tai_ipc_materialize_token_kv_page_slots_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=slot_remap.slot_logical_pages,
@@ -2265,17 +2576,27 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
start_slot=0,
end_slot=prefix_pages,
)
if not materialized_by_ipc:
materialize_local_token_kv_page_slots_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
page_size=page_size,
start_slot=0,
end_slot=prefix_pages,
)
prefix_rows = slot_range_to_token_slice(page_size, 0, prefix_pages)
_all_reduce_materialized_buffer_range(
dense_kv_cache,
layout.cp_size,
prefix_rows.start,
prefix_rows.stop,
nvtx_source=nvtx_source,
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
prefix_rows = slot_range_to_token_slice(page_size, 0, prefix_pages)
_all_reduce_materialized_buffer_range(
dense_kv_cache,
layout.cp_size,
prefix_rows.start,
prefix_rows.stop,
nvtx_source=nvtx_source,
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
logical_locs = filter_locs_mappable_to_physical_pool(
logical_locs=logical_locs,
@@ -2326,7 +2647,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
dense_page_buffer = page_buffer.new_zeros(
(slot_remap.dense_num_pages, *page_buffer.shape[1:])
)
materialize_local_paged_buffer_page_slots_into(
materialized_by_ipc = _try_tai_ipc_materialize_paged_buffer_page_slots_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=slot_remap.slot_logical_pages,
@@ -2334,16 +2655,25 @@ def materialize_prefix_and_reuse_current_index_page_slots(
start_slot=0,
end_slot=prefix_pages,
)
prefix_rows = slot_range_to_page_slice(0, prefix_pages)
_all_reduce_materialized_buffer_range(
dense_page_buffer,
layout.cp_size,
prefix_rows.start,
prefix_rows.stop,
nvtx_source=nvtx_source,
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
if not materialized_by_ipc:
materialize_local_paged_buffer_page_slots_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=slot_remap.slot_logical_pages,
layout=layout,
start_slot=0,
end_slot=prefix_pages,
)
prefix_rows = slot_range_to_page_slice(0, prefix_pages)
_all_reduce_materialized_buffer_range(
dense_page_buffer,
layout.cp_size,
prefix_rows.start,
prefix_rows.stop,
nvtx_source=nvtx_source,
nvtx_layer_id=layer_id,
nvtx_cp_rank=layout.cp_rank,
)
dense_page_buffer = fill_current_index_page_slots(
dense_page_buffer=dense_page_buffer,
current_index_k=current_index_k,
@@ -2777,13 +3107,33 @@ def materialize_shared_token_kv_buffer(
dense_kv_cache, dense_locs = tai_result
use_slot_materialize = False
materialized_by_ipc = False
if use_slot_materialize:
dense_kv_cache = materialize_local_token_kv_page_slots(
dense_kv_cache = kv_cache.new_zeros(
(
(int(materialized_logical_pages.numel()) + 1) * page_size,
*kv_cache.shape[1:],
)
)
materialized_by_ipc = _try_tai_ipc_materialize_token_kv_page_slots_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=materialized_logical_pages,
layout=layout,
page_size=page_size,
start_slot=0,
end_slot=int(materialized_logical_pages.numel()),
)
if not materialized_by_ipc:
materialize_local_token_kv_page_slots_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=materialized_logical_pages,
layout=layout,
page_size=page_size,
start_slot=0,
end_slot=int(materialized_logical_pages.numel()),
)
elif dense_kv_cache is None:
dense_kv_cache = materialize_local_token_kv_pages(
kv_cache=kv_cache,
@@ -2820,13 +3170,14 @@ def materialize_shared_token_kv_buffer(
),
)
dense_kv_cache = _all_reduce_materialized_buffer(
dense_kv_cache,
layout.cp_size,
nvtx_source=nvtx_source,
nvtx_layer_id=nvtx_layer_id,
nvtx_cp_rank=layout.cp_rank,
)
if not materialized_by_ipc:
dense_kv_cache = _all_reduce_materialized_buffer(
dense_kv_cache,
layout.cp_size,
nvtx_source=nvtx_source,
nvtx_layer_id=nvtx_layer_id,
nvtx_cp_rank=layout.cp_rank,
)
if cp_shared_kv_debug_enabled():
cp_shared_kv_debug_log(
@@ -2860,29 +3211,48 @@ def materialize_shared_paged_buffer(
layout=layout,
physical_page_capacity=page_buffer.shape[0],
)
tai_result = _try_tai_materialize_shared_pages(
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
)
if tai_result is not None:
dense_page_buffer, dense_pages = tai_result
materialized_logical_pages = logical_pages.reshape(-1)
elif slot_remap is not None:
materialized_by_ipc = False
if slot_remap is not None:
materialized_logical_pages = slot_remap.slot_logical_pages
dense_pages = slot_remap.dense_pages
dense_page_buffer = materialize_local_paged_buffer_page_slots(
dense_page_buffer = page_buffer.new_zeros(
(slot_remap.dense_num_pages, *page_buffer.shape[1:])
)
materialized_by_ipc = _try_tai_ipc_materialize_paged_buffer_page_slots_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=materialized_logical_pages,
layout=layout,
start_slot=0,
end_slot=int(materialized_logical_pages.numel()),
)
if not materialized_by_ipc:
materialize_local_paged_buffer_page_slots_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=materialized_logical_pages,
layout=layout,
start_slot=0,
end_slot=int(materialized_logical_pages.numel()),
)
else:
materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
dense_page_buffer = materialize_local_paged_buffer_page_slots(
tai_result = _try_tai_materialize_shared_pages(
page_buffer=page_buffer,
slot_logical_pages=materialized_logical_pages,
logical_pages=logical_pages,
layout=layout,
)
if tai_result is not None:
dense_page_buffer, dense_pages = tai_result
materialized_logical_pages = logical_pages.reshape(-1)
else:
materialized_logical_pages, dense_pages = build_slot_page_remap(
logical_pages
)
dense_page_buffer = materialize_local_paged_buffer_page_slots(
page_buffer=page_buffer,
slot_logical_pages=materialized_logical_pages,
layout=layout,
)
if cp_shared_kv_debug_enabled():
owned_pages = materialized_logical_pages[
@@ -2910,13 +3280,14 @@ def materialize_shared_paged_buffer(
),
)
dense_page_buffer = _all_reduce_materialized_buffer(
dense_page_buffer,
layout.cp_size,
nvtx_source=nvtx_source,
nvtx_layer_id=nvtx_layer_id,
nvtx_cp_rank=layout.cp_rank,
)
if not materialized_by_ipc:
dense_page_buffer = _all_reduce_materialized_buffer(
dense_page_buffer,
layout.cp_size,
nvtx_source=nvtx_source,
nvtx_layer_id=nvtx_layer_id,
nvtx_cp_rank=layout.cp_rank,
)
if cp_shared_kv_debug_enabled():
cp_shared_kv_debug_log(
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+221
View File
@@ -885,6 +885,125 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]])
def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_pages(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=4, cp_rank=0)
slot_logical_pages = torch.tensor(
[[1, 2, 0, 5], [-1, 8, 9, 17]], dtype=torch.int64
)
owner_ranks, src_page_indices = (
runtime.build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages,
layout,
physical_page_capacity=3,
)
)
self.assertEqual(owner_ranks.tolist(), [0, 1, -1, 0, -1, 3, -1, -1])
self.assertEqual(src_page_indices.tolist(), [1, 1, -1, 2, -1, 2, -1, -1])
def test_materialize_shared_token_kv_uses_ipc_without_all_reduce(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
page_size = 4
layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0)
kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32)
remap_logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=torch.tensor([4, 8, 20], dtype=torch.int64),
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 3)
kwargs["dense_kv_cache"][4:16] = torch.arange(
12, dtype=torch.float32
).view(12, 1, 1)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_token_kv_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer",
side_effect=AssertionError("IPC path must not all-reduce"),
):
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
kv_cache=kv_cache,
logical_locs=torch.tensor([4, 8, 20], dtype=torch.int64),
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
slot_remap=slot_remap,
)
self.assertEqual(dense_locs.tolist(), [4, 8, 12])
self.assertEqual(list(dense_kv[4:16].flatten().tolist()), list(range(12)))
def test_materialize_prefix_current_token_kv_uses_ipc_without_all_reduce(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
page_size = 4
layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0)
kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32)
logical_locs = torch.tensor([[4, 8, 20, 21, 22, 23]], dtype=torch.int64)
current_locs = torch.tensor([20, 21], dtype=torch.int64)
current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1)
remap_logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64)
slot_remap = runtime.build_shared_token_kv_slot_remap(
kv_cache=kv_cache,
logical_locs=logical_locs,
remap_logical_pages=remap_logical_pages,
layout=layout,
page_size=page_size,
)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 2)
kwargs["dense_kv_cache"][4:12] = torch.arange(
10, 18, dtype=torch.float32
).view(8, 1, 1)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_token_kv_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer_range",
side_effect=AssertionError("IPC path must not range all-reduce"),
):
mixed_kv, mixed_locs = (
runtime.materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_locs,
current_kv_cache=current_kv,
current_locs=current_locs,
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
prefix_pages=2,
)
)
self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]])
self.assertEqual(mixed_kv[4].item(), 10)
self.assertEqual(mixed_kv[8].item(), 14)
self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -1195,6 +1314,108 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
)
)
def test_materialize_shared_index_uses_ipc_without_all_reduce(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)
page_buffer = torch.zeros((8, 3), dtype=torch.uint8)
logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64)
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
page_buffer,
logical_pages,
layout,
)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 3)
kwargs["dense_page_buffer"][1:4] = torch.tensor(
[[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.uint8
)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_paged_buffer_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer",
side_effect=AssertionError("IPC path must not all-reduce"),
):
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
page_buffer=page_buffer,
logical_pages=logical_pages,
layout=layout,
slot_remap=slot_remap,
)
self.assertEqual(dense_pages.tolist(), [[1, 2, 3]])
self.assertEqual(
dense_page_buffer[1:4].tolist(),
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
)
def test_materialize_prefix_current_index_uses_ipc_without_all_reduce(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
page_size = 4
index_head_dim = 4
scale_bytes = 4
page_bytes = page_size * index_head_dim + page_size * scale_bytes
layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0)
page_buffer = torch.zeros((8, page_bytes), dtype=torch.uint8)
logical_pages = torch.tensor([[1, 2]], dtype=torch.int64)
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
page_buffer,
logical_pages,
layout,
)
current_k = torch.tensor(
[[1, 2, 3, 4], [5, 6, 7, 8]],
dtype=torch.uint8,
)
current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32)
def fake_ipc_into(**kwargs):
self.assertEqual(kwargs["start_slot"], 0)
self.assertEqual(kwargs["end_slot"], 1)
kwargs["dense_page_buffer"][1] = torch.arange(
page_bytes, dtype=torch.uint8
)
return True
with patch.object(
runtime,
"_try_tai_ipc_materialize_paged_buffer_page_slots_into",
side_effect=fake_ipc_into,
), patch.object(
runtime,
"_all_reduce_materialized_buffer_range",
side_effect=AssertionError("IPC path must not range all-reduce"),
):
dense_page_buffer, dense_pages = (
runtime.materialize_prefix_and_reuse_current_index_page_slots(
page_buffer=page_buffer,
current_index_k=current_k,
current_index_scale=current_scale,
current_locs=torch.tensor([8, 9], dtype=torch.int64),
slot_remap=slot_remap,
layout=layout,
page_size=page_size,
index_head_dim=index_head_dim,
prefix_pages=1,
layer_id=2,
)
)
self.assertEqual(dense_pages.tolist(), [[1, 2]])
self.assertEqual(dense_page_buffer[1].tolist(), list(range(page_bytes)))
self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0]))
self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1]))
def test_index_prefetch_partial_current_compose_fills_current_page_slots(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch