From 4342de04639c105a6561a33042b6b828d18e2d33 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 31 May 2026 23:59:23 +0800 Subject: [PATCH] Avoid redundant CP collectives on sync shared-KV materialize 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 --- ..._prefill_cp_page_aligned_cache_contract.md | 202 ++++++++ .../attention/nsa/cp_shared_kv_runtime.py | 471 ++++++++++++++++-- .../mem_cache/test_cp_shared_kv_runtime.py | 221 ++++++++ 3 files changed, 844 insertions(+), 50 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index d903be5bd..422624c2a 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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. 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 7eb21b872..2c531ee42 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 @@ -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( 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 3ecaa2ffb..53dc3eca9 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 @@ -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