diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 805ac6967..d46ce03b4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -205,6 +205,9 @@ class Envs: SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False) SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False) SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False) + SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False) + SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False) + SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py new file mode 100644 index 000000000..d335d6f8d --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -0,0 +1,431 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + _all_reduce_materialized_buffer_async, + _all_reduce_materialized_buffer_range, + build_shared_token_kv_slot_remap, + cp_shared_kv_debug_enabled, + cp_shared_kv_mla_prefetch_enabled, + cp_shared_kv_mla_prefetch_log, + cp_shared_kv_mla_prefetch_should_log_layer, + cp_shared_kv_mla_prefetch_wait_after_attention_enabled, + filter_locs_mappable_to_physical_pool, + materialize_local_token_kv_page_slots_into, + remap_logical_locs_to_slot_dense_locs_optimized, + slot_range_to_token_slice, +) +from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split +from sglang.srt.layers.dp_attention import get_attention_cp_group +from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + +logger = logging.getLogger(__name__) + + +def _prefetch_log(message: str, *args) -> None: + cp_shared_kv_mla_prefetch_log(message, *args) + + +def _is_cuda_stream_capturing() -> bool: + if not torch.cuda.is_available(): + return False + try: + return torch.cuda.is_current_stream_capturing() + except RuntimeError: + return False + + +@dataclass +class CpSharedKVMlaPrefetchHandle: + layer_id: int + dense_kv_cache: torch.Tensor + event: torch.cuda.Event + + +class CpSharedKVMlaPrefetcher: + """One-layer-ahead MLA prefix materialize prefetch for CP shared KV. + + This object is per-forward-batch. It only materializes historical prefix + pages, because current/suffix pages for layer L+1 are not written until that + layer's MLA prepare has run. + """ + + def __init__( + self, + *, + layout: CpSharedKVLayout, + page_size: int, + prefix_pages: int, + slot_logical_pages: torch.Tensor, + page_inverse: torch.Tensor, + dense_num_pages: int, + ) -> None: + self.layout = layout + self.page_size = page_size + self.prefix_pages = prefix_pages + self.slot_logical_pages = slot_logical_pages + self.page_inverse = page_inverse + self.dense_num_pages = dense_num_pages + self.total_slots = int(slot_logical_pages.numel()) + self.stream = torch.cuda.Stream() + self.handles: dict[int, CpSharedKVMlaPrefetchHandle] = {} + self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None + self.disabled = False + + @classmethod + def maybe_create( + cls, + *, + forward_batch: Any, + metadata: Any, + topk_transform_is_paged: bool, + ) -> Optional["CpSharedKVMlaPrefetcher"]: + if not cp_shared_kv_mla_prefetch_enabled(): + return None + if cp_shared_kv_debug_enabled(): + _prefetch_log("create_skip reason=debug_enabled") + return None + if not torch.cuda.is_available() or _is_cuda_stream_capturing(): + _prefetch_log( + "create_skip reason=cuda_unavailable_or_stream_capturing cuda_available=%s", + torch.cuda.is_available(), + ) + return None + if not getattr(forward_batch, "uses_cp_shared_kv", False): + _prefetch_log("create_skip reason=not_cp_shared_kv") + return None + if getattr(forward_batch, "hisparse_coordinator", None) is not None: + _prefetch_log("create_skip reason=hisparse") + return None + forward_mode = getattr(forward_batch, "forward_mode", None) + if forward_mode is None or not forward_mode.is_context_parallel_extend(): + _prefetch_log("create_skip reason=not_context_parallel_extend") + return None + if not is_nsa_prefill_cp_in_seq_split(): + _prefetch_log("create_skip reason=not_in_seq_split") + return None + if not topk_transform_is_paged: + _prefetch_log("create_skip reason=not_paged_topk") + return None + if int(getattr(forward_batch, "batch_size", 0)) != 1: + _prefetch_log( + "create_skip reason=batch_size batch_size=%s", + getattr(forward_batch, "batch_size", None), + ) + return None + + token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None) + if token_to_kv_pool is None: + _prefetch_log("create_skip reason=missing_token_to_kv_pool") + return None + if getattr(token_to_kv_pool, "layer_transfer_counter", None) is not None: + _prefetch_log("create_skip reason=layer_transfer_active") + return None + + layout = getattr(forward_batch, "cp_shared_kv_layout", None) + if layout is None: + _prefetch_log("create_skip reason=missing_layout") + return None + + extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + if extend_prefix_lens_cpu is None or len(extend_prefix_lens_cpu) != 1: + _prefetch_log("create_skip reason=bad_prefix_lens_metadata") + return None + page_size = int(getattr(token_to_kv_pool, "page_size", 1)) + if page_size <= 1: + _prefetch_log("create_skip reason=bad_page_size page_size=%s", page_size) + return None + extend_prefix_len = int(extend_prefix_lens_cpu[0]) + if extend_prefix_len <= 0 or extend_prefix_len % page_size != 0: + _prefetch_log( + "create_skip reason=prefix_not_page_aligned prefix_len=%s page_size=%s", + extend_prefix_len, + page_size, + ) + return None + prefix_pages = extend_prefix_len // page_size + + real_page_table = getattr(metadata, "real_page_table", None) + page_table_1 = getattr(metadata, "page_table_1", None) + if real_page_table is None or page_table_1 is None: + _prefetch_log("create_skip reason=missing_page_tables") + return None + if prefix_pages <= 0 or prefix_pages > int(real_page_table.numel()): + _prefetch_log( + "create_skip reason=prefix_pages_out_of_range prefix_pages=%s real_pages=%s", + prefix_pages, + int(real_page_table.numel()), + ) + return None + + cp_group = get_attention_cp_group() + if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: + _prefetch_log( + "create_skip reason=missing_pynccl cp_rank=%s cp_size=%s", + layout.cp_rank, + layout.cp_size, + ) + return None + + try: + first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) + kv_cache = token_to_kv_pool.get_key_buffer(first_layer_id) + remap = build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=None, + remap_logical_pages=real_page_table, + layout=layout, + page_size=page_size, + ) + except Exception: + logger.exception("Failed to initialize CP shared KV MLA prefetcher.") + return None + + _prefetch_log( + "create cp_rank=%s cp_size=%s prefix_pages=%s total_slots=%s dense_pages=%s page_size=%s", + layout.cp_rank, + layout.cp_size, + prefix_pages, + int(remap.slot_logical_pages.numel()), + remap.dense_num_pages, + page_size, + ) + + return cls( + layout=layout, + page_size=page_size, + prefix_pages=prefix_pages, + slot_logical_pages=remap.slot_logical_pages, + page_inverse=remap.page_inverse, + dense_num_pages=remap.dense_num_pages, + ) + + def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool: + start_layer = int(getattr(token_to_kv_pool, "start_layer", 0)) + kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None) + if kv_buffer is None: + return layer_id >= start_layer + return start_layer <= layer_id < start_layer + len(kv_buffer) + + def consume( + self, + *, + layer_id: int, + kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if self.disabled: + self._log_layer( + layer_id, + "consume_skip reason=disabled layer=%s", + layer_id, + ) + return None + + handle = self.handles.pop(layer_id, None) + if handle is None: + self._log_layer(layer_id, "consume_miss layer=%s", layer_id) + return None + if self.pending_attention_handle is handle: + self.pending_attention_handle = None + if handle.layer_id != layer_id: + self.disabled = True + self._log_layer( + layer_id, + "consume_skip reason=layer_mismatch expected=%s actual=%s", + layer_id, + handle.layer_id, + ) + return None + + torch.cuda.current_stream().wait_event(handle.event) + dense_kv_cache = handle.dense_kv_cache + suffix_slots = self.total_slots - self.prefix_pages + + if self.prefix_pages < self.total_slots: + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + start_slot=self.prefix_pages, + end_slot=self.total_slots, + ) + suffix_rows = slot_range_to_token_slice( + self.page_size, + self.prefix_pages, + self.total_slots, + ) + _all_reduce_materialized_buffer_range( + dense_kv_cache, + self.layout.cp_size, + suffix_rows.start, + suffix_rows.stop, + ) + + self._log_layer( + layer_id, + "consume_hit layer=%s prefix_pages=%s suffix_slots=%s dense_rows=%s", + layer_id, + self.prefix_pages, + suffix_slots, + int(dense_kv_cache.shape[0]), + ) + + logical_locs = filter_locs_mappable_to_physical_pool( + logical_locs=logical_locs, + layout=self.layout, + physical_token_capacity=kv_cache.shape[0], + ) + dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=self.page_inverse, + page_size=self.page_size, + ) + return dense_kv_cache, dense_locs + + def start_next_layer_prefix( + self, + *, + next_layer_id: int, + token_to_kv_pool: Any, + ) -> None: + if self.disabled: + self._log_next_layer( + next_layer_id, + "start_skip reason=disabled next_layer=%s", + next_layer_id, + ) + return + if next_layer_id in self.handles: + self._log_next_layer( + next_layer_id, + "start_skip reason=already_started next_layer=%s", + next_layer_id, + ) + return + if not self._layer_in_pool(token_to_kv_pool, next_layer_id): + self._log_next_layer( + next_layer_id, + "start_skip reason=layer_out_of_pool next_layer=%s", + next_layer_id, + ) + return + + try: + kv_cache = token_to_kv_pool.get_key_buffer(next_layer_id) + except Exception: + logger.exception( + "Failed to get next-layer KV cache for CP shared KV MLA prefetch." + ) + self.disabled = True + self._log_next_layer( + next_layer_id, + "start_disable reason=get_kv_failed next_layer=%s", + next_layer_id, + ) + return + + current_stream = torch.cuda.current_stream() + self.stream.wait_stream(current_stream) + try: + with torch.cuda.stream(self.stream): + dense_kv_cache = kv_cache.new_zeros( + (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) + ) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + start_slot=0, + end_slot=self.prefix_pages, + ) + prefix_rows = slot_range_to_token_slice( + self.page_size, + 0, + self.prefix_pages, + ) + event = _all_reduce_materialized_buffer_async( + dense_kv_cache[prefix_rows], + cp_size=self.layout.cp_size, + stream=self.stream, + ) + if event is None: + self.disabled = True + self._log_next_layer( + next_layer_id, + "start_disable reason=async_reduce_unavailable next_layer=%s", + next_layer_id, + ) + return + except Exception: + logger.exception("Failed to start CP shared KV MLA prefix prefetch.") + self.disabled = True + self._log_next_layer( + next_layer_id, + "start_disable reason=start_exception next_layer=%s", + next_layer_id, + ) + return + + handle = CpSharedKVMlaPrefetchHandle( + layer_id=next_layer_id, + dense_kv_cache=dense_kv_cache, + event=event, + ) + self.handles[next_layer_id] = handle + self.pending_attention_handle = handle + self._log_next_layer( + next_layer_id, + "start next_layer=%s prefix_pages=%s prefix_rows=%s dense_rows=%s", + next_layer_id, + self.prefix_pages, + prefix_rows.stop - prefix_rows.start, + int(dense_kv_cache.shape[0]), + ) + + def wait_attention_window(self) -> None: + if not cp_shared_kv_mla_prefetch_wait_after_attention_enabled(): + handle = self.pending_attention_handle + if handle is not None: + self._log_next_layer( + handle.layer_id, + "attention_wait_deferred next_layer=%s", + handle.layer_id, + ) + return + + handle = self.pending_attention_handle + self.pending_attention_handle = None + if handle is None: + return + torch.cuda.current_stream().wait_event(handle.event) + self._log_next_layer( + handle.layer_id, + "attention_wait next_layer=%s", + handle.layer_id, + ) + + def _log_layer(self, layer_id: int, message: str, *args) -> None: + if cp_shared_kv_mla_prefetch_should_log_layer(layer_id): + self._log(message, *args) + + def _log_next_layer(self, next_layer_id: int, message: str, *args) -> None: + if cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id): + self._log(message, *args) + + def _log(self, message: str, *args) -> None: + _prefetch_log( + "cp_rank=%s cp_size=%s " + message, + self.layout.cp_rank, + self.layout.cp_size, + *args, + ) 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 d868529d2..6e8e4fb72 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 @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from dataclasses import dataclass from functools import lru_cache import torch @@ -13,6 +14,7 @@ logger = logging.getLogger(__name__) _DEBUG_LOG_COUNTS: dict[str, int] = {} _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {} +_MLA_PREFETCH_LOG_PROBE_LAYER = 2 def cp_shared_kv_debug_enabled() -> bool: @@ -27,6 +29,36 @@ def cp_shared_kv_tai_materialize_enabled() -> bool: return envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get() +def cp_shared_kv_mla_prefetch_enabled() -> bool: + return envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get() + + +def cp_shared_kv_mla_prefetch_log_enabled() -> bool: + return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get() + + +def cp_shared_kv_mla_prefetch_wait_after_attention_enabled() -> bool: + return envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.get() + + +def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None: + if cp_shared_kv_mla_prefetch_log_enabled(): + logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args) + + +def cp_shared_kv_mla_prefetch_should_log_layer(layer_id: int) -> bool: + return int(layer_id) == _MLA_PREFETCH_LOG_PROBE_LAYER + + +@dataclass(frozen=True) +class SharedTokenKVSlotRemap: + slot_logical_pages: torch.Tensor + page_inverse: torch.Tensor + dense_locs: torch.Tensor | None + logical_page_capacity: int + dense_num_pages: int + + @lru_cache(maxsize=1) def _load_tai_materialize_kernels(): try: @@ -142,6 +174,157 @@ def _try_tai_materialize_token_kv_pages_and_locs( return None +def _try_tai_build_slot_page_inverse( + slot_logical_pages: torch.Tensor, + logical_page_capacity: int, +) -> torch.Tensor | None: + if not _tai_materialize_runtime_enabled(): + return None + + kernels = _load_tai_materialize_kernels() + if kernels is None: + return None + + try: + return kernels.build_slot_page_inverse( + _contiguous_for_tai(slot_logical_pages.reshape(-1)), + logical_page_capacity, + ) + except Exception as exc: + _log_tai_materialize_fallback( + "page_inverse_failed", + "CP shared KV tai page inverse build failed; falling back to torch " + "remap. error=%s", + exc, + ) + return None + + +def build_slot_page_inverse_optimized( + slot_logical_pages: torch.Tensor, + logical_page_capacity: int, +) -> torch.Tensor: + tai_result = _try_tai_build_slot_page_inverse( + slot_logical_pages, + logical_page_capacity, + ) + if tai_result is not None: + return tai_result + return build_slot_page_inverse( + slot_logical_pages, + logical_page_capacity=logical_page_capacity, + ) + + +def remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs: torch.Tensor, + page_inverse: torch.Tensor, + page_size: int, +) -> torch.Tensor: + if _tai_materialize_runtime_enabled(): + kernels = _load_tai_materialize_kernels() + if kernels is not None: + try: + return kernels.remap_logical_locs_to_slot_dense_locs( + _contiguous_for_tai(logical_locs), + _contiguous_for_tai(page_inverse), + page_size=page_size, + ) + except Exception as exc: + _log_tai_materialize_fallback( + "loc_remap_failed", + "CP shared KV tai loc remap failed; falling back to torch " + "remap. error=%s", + exc, + ) + return remap_logical_locs_to_slot_dense_locs( + logical_locs, + page_inverse=page_inverse, + page_size=page_size, + ) + + +def _copy_tai_dense_slot_range_body( + *, + tai_dense_kv_cache: torch.Tensor, + dense_kv_cache: torch.Tensor, + page_size: int, + start_slot: int, + end_slot: int, +) -> None: + if start_slot == end_slot: + return + dst_rows = slot_range_to_token_slice(page_size, start_slot, end_slot) + src_rows = slot_range_to_token_slice(page_size, 0, end_slot - start_slot) + dense_kv_cache[dst_rows].copy_(tai_dense_kv_cache[src_rows]) + + +def _try_tai_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 not _tai_materialize_runtime_enabled(): + return False + + kernels = _load_tai_materialize_kernels() + if kernels is None: + return False + + 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: + materialize_into = getattr( + kernels, + "materialize_shared_token_kv_pages_into", + None, + ) + if materialize_into is not None: + materialize_into( + kv_cache, + slot_logical_pages_range, + dense_kv_cache, + page_size=page_size, + start_slot=start_slot, + cp_rank=layout.cp_rank, + cp_size=layout.cp_size, + ) + else: + tai_dense_kv_cache = kernels.materialize_shared_token_kv_pages( + kv_cache, + slot_logical_pages_range, + page_size=page_size, + cp_rank=layout.cp_rank, + cp_size=layout.cp_size, + ) + _copy_tai_dense_slot_range_body( + tai_dense_kv_cache=tai_dense_kv_cache, + dense_kv_cache=dense_kv_cache, + page_size=page_size, + start_slot=start_slot, + end_slot=end_slot, + ) + return True + except Exception as exc: + _log_tai_materialize_fallback( + "token_range_failed", + "CP shared KV tai token range materialize failed; falling back to " + "torch materialize. error=%s", + exc, + ) + return False + + def is_current_only_extend_batch(forward_batch) -> bool: """Return whether an extend batch has no cached/history tokens. @@ -472,6 +655,59 @@ def remap_logical_locs_to_slot_dense_locs( return torch.where(mapped, dense_values, dense_locs) +def build_shared_token_kv_slot_remap( + kv_cache: torch.Tensor, + logical_locs: torch.Tensor | None, + remap_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + page_size: int, +) -> SharedTokenKVSlotRemap: + """Build the fixed slot-layout remap used by shared token KV materialize. + + The slot layout is intentionally the same as `build_slot_page_remap`: dense + page 0 is the dummy page and dense page `slot + 1` corresponds to + `remap_logical_pages.reshape(-1)[slot]`. Phase 8 uses the same remap to + materialize prefix/suffix ranges into one dense buffer without changing + attention page-table semantics. + """ + + _debug_assert_no_negative_tensor_values( + remap_logical_pages, + context="CP shared KV token materialize page remap", + tensor_name="remap_logical_pages", + ) + remap_logical_pages = filter_pages_mappable_to_physical_pool( + logical_pages=remap_logical_pages, + layout=layout, + physical_page_capacity=kv_cache.shape[0] // page_size, + ) + logical_page_capacity = _logical_page_capacity_from_physical_page_capacity( + kv_cache.shape[0] // page_size, + layout, + ) + slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages) + page_inverse = build_slot_page_inverse_optimized( + slot_logical_pages, + logical_page_capacity=logical_page_capacity, + ) + dense_locs = ( + remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=page_inverse, + page_size=page_size, + ) + if logical_locs is not None + else None + ) + return SharedTokenKVSlotRemap( + slot_logical_pages=slot_logical_pages, + page_inverse=page_inverse, + dense_locs=dense_locs, + logical_page_capacity=logical_page_capacity, + dense_num_pages=int(slot_logical_pages.numel()) + 1, + ) + + def remap_logical_locs_to_dense_locs( logical_locs: torch.Tensor, unique_logical_pages: torch.Tensor, @@ -675,7 +911,65 @@ def materialize_local_token_kv_page_slots( if slot_logical_pages.numel() == 0: return dense_kv_cache - logical_pages = slot_logical_pages.reshape(-1).to(torch.long) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=slot_logical_pages, + layout=layout, + page_size=page_size, + start_slot=0, + end_slot=int(slot_logical_pages.numel()), + ) + return dense_kv_cache + + +def materialize_local_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 | None = None, +) -> None: + """Materialize a slot range into an existing dense token KV buffer. + + `start_slot`/`end_slot` are page-table slots, not dense page ids. Dense + page 0 is the dummy page, so slot `i` writes dense token rows for page + `i + 1`. + """ + + flat_slot_logical_pages = slot_logical_pages.reshape(-1) + total_slots = int(flat_slot_logical_pages.numel()) + if end_slot is None: + end_slot = total_slots + if start_slot < 0 or end_slot < start_slot or end_slot > total_slots: + raise ValueError( + "Invalid CP shared KV slot materialize range: " + f"start_slot={start_slot} end_slot={end_slot} total_slots={total_slots}" + ) + if start_slot == end_slot: + return + + expected_rows = (total_slots + 1) * page_size + if dense_kv_cache.shape[0] < expected_rows: + raise ValueError( + "CP shared KV dense token buffer is too small for slot materialize: " + f"dense_rows={dense_kv_cache.shape[0]} expected_at_least={expected_rows}" + ) + + if _try_tai_materialize_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=flat_slot_logical_pages, + layout=layout, + page_size=page_size, + start_slot=start_slot, + end_slot=end_slot, + ): + return + + logical_pages = flat_slot_logical_pages[start_slot:end_slot].to(torch.long) owned_mask = layout.owned_pages_mask(logical_pages) physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long) safe_physical_pages = torch.where( @@ -686,16 +980,29 @@ def materialize_local_token_kv_page_slots( page_offsets = torch.arange(page_size, device=kv_cache.device, dtype=torch.long) src_tokens = (safe_physical_pages[:, None] * page_size + page_offsets).reshape(-1) - dense_body = dense_kv_cache[page_size:].view( - dense_num_pages - 1, + dense_body = dense_kv_cache[page_size:expected_rows].view( + total_slots, page_size, *kv_cache.shape[1:], ) - gathered = kv_cache[src_tokens].view_as(dense_body) + dense_range = dense_body[start_slot:end_slot] + gathered = kv_cache[src_tokens].view_as(dense_range) owned_view = owned_mask.view(-1, *([1] * (dense_body.ndim - 1))) zero = torch.zeros((), dtype=kv_cache.dtype, device=kv_cache.device) - dense_body.copy_(torch.where(owned_view, gathered, zero)) - return dense_kv_cache + dense_range.copy_(torch.where(owned_view, gathered, zero)) + + +def slot_range_to_token_slice( + page_size: int, + start_slot: int, + end_slot: int, +) -> slice: + if start_slot < 0 or end_slot < start_slot: + raise ValueError( + "Invalid CP shared KV slot token slice range: " + f"start_slot={start_slot} end_slot={end_slot}" + ) + return slice((start_slot + 1) * page_size, (end_slot + 1) * page_size) def token_page_copy_debug_checksum( @@ -824,6 +1131,65 @@ def _all_reduce_materialized_buffer(buffer: torch.Tensor, cp_size: int) -> torch return buffer +def _all_reduce_materialized_buffer_range( + buffer: torch.Tensor, + cp_size: int, + start_row: int, + end_row: int, +) -> torch.Tensor: + if start_row < 0 or end_row < start_row or end_row > buffer.shape[0]: + raise ValueError( + "Invalid CP shared KV materialize reduce row range: " + f"start_row={start_row} end_row={end_row} rows={buffer.shape[0]}" + ) + if start_row == end_row: + return buffer + _all_reduce_materialized_buffer(buffer[start_row:end_row], cp_size) + return buffer + + +def _all_reduce_materialized_buffer_async( + buffer: torch.Tensor, + cp_size: int, + stream: torch.cuda.Stream, +) -> torch.cuda.Event | None: + """Enqueue an in-place CP all-reduce on `stream`. + + Returns a CUDA event recorded after the collective, or `None` when the + async pynccl path is unavailable. Callers must fallback before launching + rank-divergent collectives if this returns `None`. + """ + + if not torch.cuda.is_available(): + return None + event = torch.cuda.Event() + if cp_size <= 1 or buffer.numel() == 0: + with torch.cuda.stream(stream): + event.record(stream) + return event + + cp_group = get_attention_cp_group() + pynccl_comm = getattr(cp_group, "pynccl_comm", None) + if pynccl_comm is None: + return None + + comm_buffer = _comm_view(buffer) + try: + with pynccl_comm.change_state(enable=True, stream=stream): + pynccl_comm.all_reduce(comm_buffer, stream=stream) + event.record(stream) + except Exception as exc: + _log_tai_materialize_fallback( + "prefetch_async_allreduce_failed", + "CP shared KV MLA prefetch async all-reduce is unavailable; " + "falling back to sync materialize. error=%s", + exc, + limit=4, + ) + return None + return event + + def materialize_shared_token_kv_buffer( kv_cache: torch.Tensor, logical_locs: torch.Tensor, @@ -896,11 +1262,11 @@ def materialize_shared_token_kv_buffer( ) if tai_result is None: materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages) - page_inverse = build_slot_page_inverse( + page_inverse = build_slot_page_inverse_optimized( materialized_logical_pages, logical_page_capacity=logical_page_capacity, ) - dense_locs = remap_logical_locs_to_slot_dense_locs( + dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( logical_locs, page_inverse=page_inverse, page_size=page_size, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index f5454d29a..41c272d07 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -9,11 +9,17 @@ import torch from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import ( + CpSharedKVMlaPrefetcher, +) from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_current_loc_remap, cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, cp_shared_kv_current_reuse_enabled, + cp_shared_kv_mla_prefetch_log, + cp_shared_kv_mla_prefetch_log_enabled, + cp_shared_kv_mla_prefetch_should_log_layer, filter_owned_logical_locs, is_current_only_extend_batch, materialize_shared_token_kv_buffer, @@ -585,6 +591,7 @@ class NativeSparseAttnBackend( """Init the metadata for a forward pass.""" batch_size = forward_batch.batch_size device = forward_batch.seq_lens.device + forward_batch.cp_shared_kv_mla_prefetcher = None if forward_batch.forward_mode.is_target_verify(): draft_token_num = self.speculative_num_draft_tokens @@ -857,6 +864,15 @@ class NativeSparseAttnBackend( token_to_batch_idx=token_to_batch_idx, ) self.forward_metadata = metadata + forward_batch.cp_shared_kv_mla_prefetcher = ( + CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=( + topk_transform_method == TopkTransformMethod.PAGED + ), + ) + ) def _cal_indexer_k_start_end( self, @@ -1594,6 +1610,9 @@ class NativeSparseAttnBackend( and topk_transform_method == TopkTransformMethod.PAGED ): assert forward_batch.cp_shared_kv_layout is not None + mla_prefetcher = getattr( + forward_batch, "cp_shared_kv_mla_prefetcher", None + ) can_reuse_current_kv = ( cp_shared_kv_current_reuse_enabled() and is_current_only_extend_batch(forward_batch) @@ -1602,6 +1621,39 @@ class NativeSparseAttnBackend( and k.shape[0] == forward_batch.out_cache_loc.numel() and k_rope.shape[0] == forward_batch.out_cache_loc.numel() ) + if cp_shared_kv_mla_prefetch_log_enabled(): + if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id): + prefix_lens_cpu = getattr( + forward_batch, "extend_prefix_lens_cpu", None + ) + extend_lens_cpu = getattr( + forward_batch, "extend_seq_lens_cpu", None + ) + prefix_lens = ( + [int(x) for x in prefix_lens_cpu] + if prefix_lens_cpu is not None + else None + ) + extend_lens = ( + [int(x) for x in extend_lens_cpu] + if extend_lens_cpu is not None + else None + ) + cp_shared_kv_mla_prefetch_log( + "forward_layer cp_rank=%s layer=%s cache_hit=%s " + "has_prefetcher=%s can_current_reuse=%s prefix_lens=%s " + "extend_lens=%s page_table_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + any(prefix_len > 0 for prefix_len in prefix_lens or []), + mla_prefetcher is not None, + can_reuse_current_kv, + prefix_lens, + extend_lens, + tuple(page_table_1.shape) + if page_table_1 is not None + else None, + ) if can_reuse_current_kv: logical_page_table_1 = page_table_1 current_mask, page_table_1 = build_current_loc_remap( @@ -1630,100 +1682,124 @@ class NativeSparseAttnBackend( ) kv_cache = _cat([k, k_rope], dim=-1) else: - kv_cache, page_table_1 = materialize_shared_token_kv_buffer( - kv_cache=kv_cache, - logical_locs=page_table_1, - remap_logical_locs=metadata.page_table_1, - remap_logical_pages=metadata.real_page_table, - layout=forward_batch.cp_shared_kv_layout, - page_size=forward_batch.token_to_kv_pool.page_size, - ) - - if nsa_impl == "tilelang": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_tilelang( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) - elif nsa_impl == "flashmla_sparse": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - - if topk_transform_method == TopkTransformMethod.RAGGED: - if any(forward_batch.extend_prefix_lens_cpu): - page_table_1_flattened = ( - self.forward_metadata.page_table_1_flattened - ) - assert page_table_1_flattened is not None - if forward_batch.uses_cp_shared_kv: - assert forward_batch.cp_shared_kv_layout is not None - kv_cache, page_table_1_flattened = ( - materialize_shared_token_kv_buffer( - kv_cache=kv_cache, - logical_locs=page_table_1_flattened, - layout=forward_batch.cp_shared_kv_layout, - page_size=forward_batch.token_to_kv_pool.page_size, - ) - ) - kv_cache = dequantize_k_cache_paged( - kv_cache, page_table_1_flattened + prefetched_kv = None + if mla_prefetcher is not None: + prefetched_kv = mla_prefetcher.consume( + layer_id=layer.layer_id, + kv_cache=kv_cache, + logical_locs=page_table_1, ) + if prefetched_kv is not None: + kv_cache, page_table_1 = prefetched_kv else: - kv_cache = _cat([k, k_rope], dim=-1) - page_table_1 = topk_indices + kv_cache, page_table_1 = materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=page_table_1, + remap_logical_locs=metadata.page_table_1, + remap_logical_pages=metadata.real_page_table, + layout=forward_batch.cp_shared_kv_layout, + page_size=forward_batch.token_to_kv_pool.page_size, + ) - return self._forward_flashmla_sparse( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) - elif nsa_impl == "flashmla_kv": - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - return self._forward_flashmla_kv( - q_all=q_all, - kv_cache=kv_cache, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - # TODO optimize args - layer=layer, - metadata=metadata, - page_table_1=page_table_1, - ) - elif nsa_impl == "fa3": - return self._forward_fa3( - q_rope=q_rope, - kv_cache=kv_cache, - v_head_dim=layer.v_head_dim, - q_nope=q_nope, - page_table=page_table_1, - cache_seqlens=metadata.nsa_cache_seqlens_int32, - cu_seqlens_q=metadata.nsa_cu_seqlens_q, - cu_seqlens_k=metadata.nsa_cu_seqlens_k, - max_seqlen_q=metadata.nsa_max_seqlen_q, - sm_scale=layer.scaling, - logit_cap=layer.logit_cap, - page_size=1, - ) - elif nsa_impl == "aiter": - if q_rope is not None: - q_all = torch.cat([q_nope, q_rope], dim=-1) - return self._forward_aiter_extend( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - layer=layer, - ) + if mla_prefetcher is not None: + mla_prefetcher.start_next_layer_prefix( + next_layer_id=layer.layer_id + 1, + token_to_kv_pool=forward_batch.token_to_kv_pool, + ) else: - raise ValueError( - f"Unsupported {nsa_impl = } for forward_extend. Consider using an other attention backend." - ) + mla_prefetcher = None + + try: + if nsa_impl == "tilelang": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + attn_output = self._forward_tilelang( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif nsa_impl == "flashmla_sparse": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + + if topk_transform_method == TopkTransformMethod.RAGGED: + if any(forward_batch.extend_prefix_lens_cpu): + page_table_1_flattened = ( + self.forward_metadata.page_table_1_flattened + ) + assert page_table_1_flattened is not None + if forward_batch.uses_cp_shared_kv: + assert forward_batch.cp_shared_kv_layout is not None + kv_cache, page_table_1_flattened = ( + materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=page_table_1_flattened, + layout=forward_batch.cp_shared_kv_layout, + page_size=forward_batch.token_to_kv_pool.page_size, + ) + ) + kv_cache = dequantize_k_cache_paged( + kv_cache, page_table_1_flattened + ) + else: + kv_cache = _cat([k, k_rope], dim=-1) + page_table_1 = topk_indices + + attn_output = self._forward_flashmla_sparse( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif nsa_impl == "flashmla_kv": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + attn_output = self._forward_flashmla_kv( + q_all=q_all, + kv_cache=kv_cache, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + # TODO optimize args + layer=layer, + metadata=metadata, + page_table_1=page_table_1, + ) + elif nsa_impl == "fa3": + attn_output = self._forward_fa3( + q_rope=q_rope, + kv_cache=kv_cache, + v_head_dim=layer.v_head_dim, + q_nope=q_nope, + page_table=page_table_1, + cache_seqlens=metadata.nsa_cache_seqlens_int32, + cu_seqlens_q=metadata.nsa_cu_seqlens_q, + cu_seqlens_k=metadata.nsa_cu_seqlens_k, + max_seqlen_q=metadata.nsa_max_seqlen_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + page_size=1, + ) + elif nsa_impl == "aiter": + if q_rope is not None: + q_all = torch.cat([q_nope, q_rope], dim=-1) + attn_output = self._forward_aiter_extend( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + layer=layer, + ) + else: + raise ValueError( + f"Unsupported {nsa_impl = } for forward_extend. Consider using an other attention backend." + ) + finally: + if mla_prefetcher is not None: + mla_prefetcher.wait_attention_window() + + return attn_output def forward_decode( self, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 0ef8700d4..b2c949414 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -32,7 +32,7 @@ from __future__ import annotations from dataclasses import dataclass from enum import IntEnum, auto from functools import total_ordering -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import torch import triton @@ -424,6 +424,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): cp_shared_kv_layout: Optional[CpSharedKVLayout] = None cp_local_out_cache_loc: Optional[torch.Tensor] = None cp_shared_mla_direct_write_done: bool = False + cp_shared_kv_mla_prefetcher: Optional[Any] = None # For hidden states before normal return_hidden_states_before_norm: bool = False 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 13005505c..d66d60683 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 @@ -182,6 +182,208 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(dense_kv[20:24], kv_cache[12:16])) self.assertEqual(float(dense_kv[4:8].abs().sum().item()), 0.0) + def test_materialize_local_token_kv_page_slots_into_matches_full_slots(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + materialize_local_token_kv_page_slots, + materialize_local_token_kv_page_slots_into, + ) + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1) + kv_cache = torch.arange(0, 32 * 2, dtype=torch.float32).view(32, 1, 2) + slot_logical_pages = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int64) + + full = materialize_local_token_kv_page_slots( + kv_cache=kv_cache, + slot_logical_pages=slot_logical_pages, + layout=layout, + page_size=4, + ) + ranged = kv_cache.new_zeros(full.shape) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=ranged, + slot_logical_pages=slot_logical_pages, + layout=layout, + page_size=4, + start_slot=0, + end_slot=3, + ) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=ranged, + slot_logical_pages=slot_logical_pages, + layout=layout, + page_size=4, + start_slot=3, + end_slot=slot_logical_pages.numel(), + ) + + self.assertTrue(torch.equal(ranged, full)) + + def test_slot_range_to_token_slice_preserves_dummy_page_offset(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + slot_range_to_token_slice, + ) + + self.assertEqual(slot_range_to_token_slice(4, 0, 2), slice(4, 12)) + self.assertEqual(slot_range_to_token_slice(4, 2, 6), slice(12, 28)) + + def test_build_shared_token_kv_slot_remap_reuses_slot_layout(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + build_shared_token_kv_slot_remap, + ) + 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((20, 1, 2), dtype=torch.float32) + logical_locs = torch.tensor([[4, 12, -1], [16, 0, 20]], dtype=torch.int64) + remap_logical_pages = torch.tensor([[1, 3, 0], [4, 5, 6]], dtype=torch.int64) + + remap = build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=logical_locs, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=4, + ) + + self.assertEqual(remap.slot_logical_pages.tolist(), [1, 3, 0, 4, 5, 6]) + self.assertEqual(remap.dense_num_pages, 7) + self.assertEqual(remap.dense_locs.tolist(), [[4, 8, -1], [16, 0, 20]]) + + def test_mla_prefetch_log_env_defaults_to_off_and_can_enable(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + cp_shared_kv_mla_prefetch_log_enabled, + cp_shared_kv_mla_prefetch_should_log_layer, + cp_shared_kv_mla_prefetch_wait_after_attention_enabled, + ) + + envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.clear() + envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.clear() + self.assertFalse(cp_shared_kv_mla_prefetch_log_enabled()) + self.assertFalse(cp_shared_kv_mla_prefetch_wait_after_attention_enabled()) + self.assertFalse(cp_shared_kv_mla_prefetch_should_log_layer(1)) + self.assertTrue(cp_shared_kv_mla_prefetch_should_log_layer(2)) + self.assertFalse(cp_shared_kv_mla_prefetch_should_log_layer(3)) + + with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True): + self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled()) + with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.override(True): + self.assertTrue(cp_shared_kv_mla_prefetch_wait_after_attention_enabled()) + + def test_token_range_materialize_uses_tai_kernel_when_enabled(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 + + class FakeTaiKernels: + def __init__(self): + self.token_calls = [] + + def materialize_shared_token_kv_pages( + self, + kv_cache, + slot_logical_pages, + *, + page_size, + cp_rank, + cp_size, + ): + self.token_calls.append( + (kv_cache, slot_logical_pages, page_size, cp_rank, cp_size) + ) + rows = (slot_logical_pages.numel() + 1) * page_size + return torch.arange( + rows * 2, + dtype=kv_cache.dtype, + device=kv_cache.device, + ).view(rows, 1, 2) + + fake_tai = FakeTaiKernels() + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1) + kv_cache = torch.zeros((32, 1, 2), dtype=torch.float32) + dense_kv_cache = torch.zeros((28, 1, 2), dtype=torch.float32) + slot_logical_pages = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int64) + + with patch( + "sglang.srt.environ.envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get", + return_value=True, + ), patch.object( + runtime, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + runtime, "_load_tai_materialize_kernels", return_value=fake_tai + ): + runtime.materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=slot_logical_pages, + layout=layout, + page_size=4, + start_slot=2, + end_slot=5, + ) + + self.assertEqual(len(fake_tai.token_calls), 1) + self.assertTrue( + torch.equal( + fake_tai.token_calls[0][1], + torch.tensor([3, 4, 5], dtype=torch.int64), + ) + ) + expected_tmp = torch.arange(32, dtype=torch.float32).view(16, 1, 2) + self.assertEqual(float(dense_kv_cache[:12].sum().item()), 0.0) + self.assertTrue(torch.equal(dense_kv_cache[12:24], expected_tmp[4:16])) + self.assertEqual(float(dense_kv_cache[24:].sum().item()), 0.0) + + def test_slot_remap_helpers_use_tai_when_enabled(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + class FakeTaiKernels: + def __init__(self): + self.inverse_calls = [] + self.remap_calls = [] + + def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity): + self.inverse_calls.append((slot_logical_pages, logical_page_capacity)) + return torch.tensor([0, 1, 2, -1], dtype=torch.long) + + def remap_logical_locs_to_slot_dense_locs( + self, + logical_locs, + page_inverse, + *, + page_size, + ): + self.remap_calls.append((logical_locs, page_inverse, page_size)) + return torch.tensor([4, -1], dtype=logical_locs.dtype) + + fake_tai = FakeTaiKernels() + slot_logical_pages = torch.tensor([1, 2, 3], dtype=torch.int64) + logical_locs = torch.tensor([4, 12], dtype=torch.int64) + + with patch( + "sglang.srt.environ.envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get", + return_value=True, + ), patch.object( + runtime, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + runtime, "_load_tai_materialize_kernels", return_value=fake_tai + ): + page_inverse = runtime.build_slot_page_inverse_optimized( + slot_logical_pages, + logical_page_capacity=4, + ) + dense_locs = runtime.remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=page_inverse, + page_size=4, + ) + + self.assertEqual(len(fake_tai.inverse_calls), 1) + self.assertEqual(len(fake_tai.remap_calls), 1) + self.assertEqual(dense_locs.tolist(), [4, -1]) + def test_materialize_local_paged_index_buffer(self): from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_dense_page_remap,