From d21952b90347150945e2002a5d237b59b21bb5a9 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 10 Jun 2026 04:28:26 +0800 Subject: [PATCH] Reduce inactive NSA index-cache transfer safely Centralize the IndexCache skip formula and thread the resulting active logical index layers into NSA KV pools. HiCache now skips only the indexer H2D/D2H payload for inactive target layers while preserving per-layer MLA KV transfer, keeping allocation shape unchanged for this phase. Constraint: P0-P2 must not compact device or host allocation yet; prefill/decode state transfer still has no logical layer-id metadata. Rejected: Recompute the skip formula separately in mem_cache | formula drift would corrupt cache or waste transfers when offset/pattern settings change. Rejected: Skip whole-layer HiCache load/backup | MLA KV remains required for every attention layer. Confidence: medium Scope-risk: moderate Directive: Before enabling compact state buffers or compact allocation, add layer-id metadata validation to PD transfer. Tested: Local py_compile for touched files; remote pytest in g0034 container: test_nsa_index_layers.py and TestNSAIndexerPageIndices, 20 passed. Not-tested: ETE replay/GSM8K with --nsa-index-topk-freq 4; PD state-transfer compaction remains unimplemented. --- python/sglang/srt/configs/nsa_index_layers.py | 113 ++++++++++ .../srt/mem_cache/hisparse_memory_pool.py | 6 +- python/sglang/srt/mem_cache/memory_pool.py | 65 +++++- .../sglang/srt/mem_cache/memory_pool_host.py | 193 ++++++++++++++---- .../model_runner_kv_cache_mixin.py | 8 + python/sglang/srt/models/deepseek_v2.py | 42 +--- .../unit/configs/test_nsa_index_layers.py | 83 ++++++++ .../unit/mem_cache/test_nsa_pool_host_unit.py | 147 +++++++++++++ 8 files changed, 567 insertions(+), 90 deletions(-) create mode 100644 python/sglang/srt/configs/nsa_index_layers.py create mode 100644 test/registered/unit/configs/test_nsa_index_layers.py diff --git a/python/sglang/srt/configs/nsa_index_layers.py b/python/sglang/srt/configs/nsa_index_layers.py new file mode 100644 index 000000000..c485684e8 --- /dev/null +++ b/python/sglang/srt/configs/nsa_index_layers.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict + + +@dataclass(frozen=True) +class NSAIndexLayerPlan: + """Logical NSA index-cache layers used by IndexCache/top-k sharing.""" + + start_layer: int + end_layer: int + active_layer_ids: tuple[int, ...] + layer_to_slot: Dict[int, int] + + def is_active(self, layer_id: int) -> bool: + return layer_id in self.layer_to_slot + + def slot_for_layer(self, layer_id: int) -> int: + try: + return self.layer_to_slot[layer_id] + except KeyError as exc: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_cache_layer] " + f"inactive index layer requested: layer_id={layer_id} " + f"active_layer_ids={list(self.active_layer_ids)}" + ) from exc + + +def nsa_index_skip_flags( + config, layer_id: int, *, is_nextn: bool = False +) -> tuple[bool, bool]: + """Return `(skip_topk, next_skip_topk)` for one logical layer. + + This intentionally mirrors the historical DeepseekV2AttentionMLA formula. + Keep this helper as the single source of truth for model-forward and cache + layer planning. + """ + + if is_nextn: + return True, True + + index_topk_freq = getattr(config, "index_topk_freq", 1) + if index_topk_freq is None: + index_topk_freq = 1 + if index_topk_freq < 1: + raise ValueError(f"index_topk_freq must be >= 1, got {index_topk_freq}") + + index_topk_pattern = getattr(config, "index_topk_pattern", None) + index_skip_topk_offset = getattr(config, "index_skip_topk_offset", None) + + if index_topk_pattern is None and index_skip_topk_offset is not None: + if index_skip_topk_offset <= 0: + raise ValueError( + "index_skip_topk_offset must be positive when configured; " + f"got {index_skip_topk_offset}" + ) + skip_topk = ( + max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0 + ) + next_skip_topk = ( + max(layer_id - index_skip_topk_offset + 2, 0) % index_topk_freq != 0 + ) + return skip_topk, next_skip_topk + + if index_topk_pattern is None: + skip_topk = max(layer_id - 1, 0) % index_topk_freq != 0 + next_skip_topk = layer_id % index_topk_freq != 0 + return skip_topk, next_skip_topk + + if layer_id < 0 or layer_id >= len(index_topk_pattern): + raise ValueError( + f"layer_id={layer_id} outside index_topk_pattern " + f"length={len(index_topk_pattern)}" + ) + skip_topk = index_topk_pattern[layer_id] == "S" + next_skip_topk = ( + layer_id < len(index_topk_pattern) - 1 + and index_topk_pattern[layer_id + 1] == "S" + ) + return skip_topk, next_skip_topk + + +def build_nsa_index_layer_plan( + config, start_layer: int, end_layer: int, *, is_nextn: bool = False +) -> NSAIndexLayerPlan: + """Build logical-layer to active-index-slot metadata. + + `end_layer` is exclusive, matching model-runner layer ranges. + Draft/nextn pools intentionally keep all local layers active for state + safety; top-k skip inside the draft forward is a separate model behavior. + """ + + if end_layer < start_layer: + raise ValueError(f"end_layer={end_layer} must be >= start_layer={start_layer}") + + if is_nextn: + active_layer_ids = tuple(range(start_layer, end_layer)) + else: + active_layer_ids = tuple( + layer_id + for layer_id in range(start_layer, end_layer) + if not nsa_index_skip_flags(config, layer_id, is_nextn=False)[0] + ) + + return NSAIndexLayerPlan( + start_layer=start_layer, + end_layer=end_layer, + active_layer_ids=active_layer_ids, + layer_to_slot={ + layer_id: slot for slot, layer_id in enumerate(active_layer_ids) + }, + ) diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index d1ada2f17..c7f806a59 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -1,7 +1,7 @@ # mapping on device memory, host memory and memory allocator import weakref -from typing import Optional +from typing import Optional, Sequence import torch from sgl_kernel.kvcacheio import transfer_kv_all_layer_mla @@ -30,6 +30,8 @@ class HiSparseNSATokenToKVPool(NSATokenToKVPool): start_layer: Optional[int] = None, end_layer: Optional[int] = None, host_to_device_ratio: int = 2, + index_active_layer_ids: Optional[Sequence[int]] = None, + compact_index_layers: bool = False, ): super().__init__( size=size, @@ -45,6 +47,8 @@ class HiSparseNSATokenToKVPool(NSATokenToKVPool): start_layer=start_layer, end_layer=end_layer, index_buf_size=size * host_to_device_ratio, + index_active_layer_ids=index_active_layer_ids, + compact_index_layers=compact_index_layers, ) self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 3026bc720..2c05f9acc 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -30,7 +30,7 @@ import logging from collections import deque from contextlib import contextmanager, nullcontext from dataclasses import dataclass, fields -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union import numpy as np import torch @@ -1870,6 +1870,8 @@ class NSATokenToKVPool(MLATokenToKVPool): start_layer: Optional[int] = None, end_layer: Optional[int] = None, index_buf_size: Optional[int] = None, + index_active_layer_ids: Optional[Sequence[int]] = None, + compact_index_layers: bool = False, ): override_dim = ( @@ -1893,6 +1895,10 @@ class NSATokenToKVPool(MLATokenToKVPool): # self.index_k_dtype = torch.float8_e4m3fn # self.index_k_scale_dtype = torch.float32 self.index_head_dim = index_head_dim + self._init_index_layer_metadata( + index_active_layer_ids=index_active_layer_ids, + compact_index_layers=compact_index_layers, + ) if index_buf_size is None: index_buf_size = size # num head == 1 and head dim == 128 for index_k in NSA @@ -1929,8 +1935,55 @@ class NSATokenToKVPool(MLATokenToKVPool): ] self._finalize_allocation_log(size) + def _init_index_layer_metadata( + self, + index_active_layer_ids: Optional[Sequence[int]], + compact_index_layers: bool, + ) -> None: + local_layer_ids = tuple( + range(self.start_layer, self.start_layer + self.layer_num) + ) + local_layer_set = frozenset(local_layer_ids) + if index_active_layer_ids is None: + active_layer_ids = local_layer_ids + else: + active_layer_ids = tuple(int(layer_id) for layer_id in index_active_layer_ids) + invalid_layer_ids = [ + layer_id for layer_id in active_layer_ids if layer_id not in local_layer_set + ] + if invalid_layer_ids: + raise ValueError( + "index_active_layer_ids must be local to this KV pool: " + f"invalid={invalid_layer_ids} local_layers={list(local_layer_ids)}" + ) + + self.index_active_layer_ids = active_layer_ids + self.index_active_layer_id_set = frozenset(active_layer_ids) + self.index_compact_layers = bool(compact_index_layers) + if self.index_compact_layers: + self.index_layer_to_slot = { + layer_id: slot for slot, layer_id in enumerate(active_layer_ids) + } + else: + self.index_layer_to_slot = { + layer_id: layer_id - self.start_layer for layer_id in active_layer_ids + } + + def is_index_layer_active(self, layer_id: int) -> bool: + return layer_id in self.index_active_layer_id_set + + def get_index_layer_slot(self, layer_id: int) -> int: + try: + return self.index_layer_to_slot[layer_id] + except KeyError as exc: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_cache_layer] " + f"inactive index layer requested: layer_id={layer_id} " + f"active_layer_ids={list(self.index_active_layer_ids)}" + ) from exc + def _get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: - return self.index_k_with_scale_buffer[layer_id - self.start_layer] + return self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)] def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: if self.layer_transfer_counter is not None: @@ -1943,7 +1996,7 @@ class NSATokenToKVPool(MLATokenToKVPool): seq_len: int, page_indices: torch.Tensor, ): - buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + buf = self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)] return index_buf_accessor.GetK.execute( self, buf, seq_len=seq_len, page_indices=page_indices ) @@ -1954,7 +2007,7 @@ class NSATokenToKVPool(MLATokenToKVPool): seq_len: int, page_indices: torch.Tensor, ): - buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + buf = self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)] return index_buf_accessor.GetS.execute( self, buf, seq_len=seq_len, page_indices=page_indices ) @@ -1978,7 +2031,7 @@ class NSATokenToKVPool(MLATokenToKVPool): k_fp8: (seq_len, index_head_dim), uint8 k_scale: (seq_len, 4), uint8 """ - buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + buf = self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)] return index_buf_accessor.GetKAndS.execute( self, buf, @@ -1995,7 +2048,7 @@ class NSATokenToKVPool(MLATokenToKVPool): index_k: torch.Tensor, index_k_scale: torch.Tensor, ) -> None: - buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + buf = self.index_k_with_scale_buffer[self.get_index_layer_slot(layer_id)] index_buf_accessor.SetKAndS.execute( pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale ) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 7a20a1c96..4032d1b32 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1945,6 +1945,32 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): ) return host_page_indices, device_page_indices + def _is_device_index_layer_active(self, device_pool, layer_id: int) -> bool: + is_active = getattr(device_pool, "is_index_layer_active", None) + if is_active is None: + return True + return bool(is_active(layer_id)) + + def _device_index_layer_slot(self, device_pool, layer_id: int) -> int: + slot_for_layer = getattr(device_pool, "get_index_layer_slot", None) + if slot_for_layer is not None: + return int(slot_for_layer(layer_id)) + return int(layer_id - getattr(device_pool, "start_layer", 0)) + + def _host_index_layer_slot(self, layer_id: int) -> int: + return int(layer_id - getattr(self, "start_layer", 0)) + + def _active_index_layer_ids_for_transfer(self, device_pool): + active_layer_ids = getattr(device_pool, "index_active_layer_ids", None) + if active_layer_ids is None: + start_layer = getattr(self, "start_layer", 0) + active_layer_ids = range(start_layer, start_layer + self.layer_num) + return tuple( + int(layer_id) + for layer_id in active_layer_ids + if self._is_device_index_layer_active(device_pool, int(layer_id)) + ) + def begin_load_to_device_op( self, host_indices: torch.Tensor, device_indices: torch.Tensor, io_backend: str ) -> None: @@ -1959,6 +1985,10 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): def _load_indexer_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): + if not self._is_device_index_layer_active(device_pool, layer_id): + return + device_layer_slot = self._device_index_layer_slot(device_pool, layer_id) + host_layer_slot = self._host_index_layer_slot(layer_id) prepared_indices = getattr(self, "_active_load_indexer_page_indices", None) if prepared_indices is None: host_page_indices, device_page_indices = self._get_indexer_page_indices( @@ -1970,8 +2000,8 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): if use_kernel: if self.layout == "layer_first": transfer_kv_per_layer_mla( - src=self.index_k_with_scale_buffer[layer_id], - dst=device_pool.index_k_with_scale_buffer[layer_id], + src=self.index_k_with_scale_buffer[host_layer_slot], + dst=device_pool.index_k_with_scale_buffer[device_layer_slot], src_indices=host_page_indices, dst_indices=device_page_indices, item_size=self.indexer_page_stride_size, @@ -1979,10 +2009,10 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): elif self.layout == "page_first": transfer_kv_per_layer_mla_pf_lf( src=self.index_k_with_scale_buffer, - dst=device_pool.index_k_with_scale_buffer[layer_id], + dst=device_pool.index_k_with_scale_buffer[device_layer_slot], src_indices=host_page_indices, dst_indices=device_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, item_size=self.indexer_page_stride_size, src_layout_dim=self.indexer_layout_dim, ) @@ -1991,8 +2021,10 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( - src_layers=[self.index_k_with_scale_buffer[layer_id]], - dst_layers=[device_pool.index_k_with_scale_buffer[layer_id]], + src_layers=[self.index_k_with_scale_buffer[host_layer_slot]], + dst_layers=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], src_indices=host_page_indices, dst_indices=device_page_indices, page_size=1, @@ -2000,19 +2032,23 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): elif self.layout == "page_first_direct": _load_tai_transfer_kv_per_layer_direct_pf_lf()( src_ptrs=[self.index_k_with_scale_buffer], - dst_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]], + dst_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], src_indices=host_page_indices, dst_indices=device_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, page_size=1, ) elif self.layout == "layer_page_first": _load_tai_transfer_kv_per_layer_direct_lpf_lf()( src_ptrs=[self.index_k_with_scale_buffer], - dst_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]], + dst_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], src_indices=host_page_indices, dst_indices=device_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, page_size=1, ) else: @@ -2026,34 +2062,79 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): host_page_indices, device_page_indices = self._get_indexer_page_indices( host_indices, device_indices ) + active_layer_ids = self._active_index_layer_ids_for_transfer(device_pool) use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0 if use_kernel: if self.layout == "layer_first": - transfer_kv_all_layer_mla( - src_layers=self.index_k_device_ptrs, - dst_layers=self.index_k_data_ptrs, - src_indices=device_page_indices, - dst_indices=host_page_indices, - item_size=self.indexer_page_stride_size, - num_layers=self.layer_num, - ) + if len(active_layer_ids) == self.layer_num: + transfer_kv_all_layer_mla( + src_layers=self.index_k_device_ptrs, + dst_layers=self.index_k_data_ptrs, + src_indices=device_page_indices, + dst_indices=host_page_indices, + item_size=self.indexer_page_stride_size, + num_layers=self.layer_num, + ) + else: + for layer_id in active_layer_ids: + device_layer_slot = self._device_index_layer_slot( + device_pool, layer_id + ) + host_layer_slot = self._host_index_layer_slot(layer_id) + transfer_kv_per_layer_mla( + src=device_pool.index_k_with_scale_buffer[ + device_layer_slot + ], + dst=self.index_k_with_scale_buffer[host_layer_slot], + src_indices=device_page_indices, + dst_indices=host_page_indices, + item_size=self.indexer_page_stride_size, + ) elif self.layout == "page_first": - transfer_kv_all_layer_mla_lf_pf( - src_layers=self.index_k_device_ptrs, - dst=self.index_k_with_scale_buffer, - src_indices=device_page_indices, - dst_indices=host_page_indices, - item_size=self.indexer_page_stride_size, - dst_layout_dim=self.indexer_layout_dim, - num_layers=self.layer_num, - ) + if len(active_layer_ids) == self.layer_num: + transfer_kv_all_layer_mla_lf_pf( + src_layers=self.index_k_device_ptrs, + dst=self.index_k_with_scale_buffer, + src_indices=device_page_indices, + dst_indices=host_page_indices, + item_size=self.indexer_page_stride_size, + dst_layout_dim=self.indexer_layout_dim, + num_layers=self.layer_num, + ) + else: + for layer_id in active_layer_ids: + device_layer_slot = self._device_index_layer_slot( + device_pool, layer_id + ) + host_layer_slot = self._host_index_layer_slot(layer_id) + _load_tai_transfer_kv_per_layer_mla_lf_pf()( + device_pool.index_k_with_scale_buffer[ + device_layer_slot + ], + self.index_k_with_scale_buffer, + device_page_indices, + host_page_indices, + layer_id=host_layer_slot, + item_size=self.indexer_page_stride_size, + dst_layout_dim=self.indexer_layout_dim, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( - src_layers=device_pool.index_k_with_scale_buffer, - dst_layers=self.index_k_with_scale_buffer, + src_layers=[ + device_pool.index_k_with_scale_buffer[ + self._device_index_layer_slot(device_pool, layer_id) + ] + for layer_id in active_layer_ids + ], + dst_layers=[ + self.index_k_with_scale_buffer[ + self._host_index_layer_slot(layer_id) + ] + for layer_id in active_layer_ids + ], src_indices=device_page_indices, dst_indices=host_page_indices, page_size=1, @@ -2065,24 +2146,36 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): # page-first-direct all-layer backup through the TAI per-layer op, # which owns the CUDA-version-specific memcpy-batch ABI handling. tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf() - for layer_id in range(self.layer_num): + for layer_id in active_layer_ids: + device_layer_slot = self._device_index_layer_slot( + device_pool, layer_id + ) + host_layer_slot = self._host_index_layer_slot(layer_id) tai_transfer( - src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]], + src_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], dst_ptrs=[self.index_k_with_scale_buffer], src_indices=device_page_indices, dst_indices=host_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, page_size=1, ) elif self.layout == "layer_page_first": tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() - for layer_id in range(self.layer_num): + for layer_id in active_layer_ids: + device_layer_slot = self._device_index_layer_slot( + device_pool, layer_id + ) + host_layer_slot = self._host_index_layer_slot(layer_id) tai_transfer( - src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]], + src_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], dst_ptrs=[self.index_k_with_scale_buffer], src_indices=device_page_indices, dst_indices=host_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, page_size=1, ) else: @@ -2093,6 +2186,10 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): def _backup_indexer_from_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): + if not self._is_device_index_layer_active(device_pool, layer_id): + return + device_layer_slot = self._device_index_layer_slot(device_pool, layer_id) + host_layer_slot = self._host_index_layer_slot(layer_id) host_page_indices, device_page_indices = self._get_indexer_page_indices( host_indices, device_indices ) @@ -2100,19 +2197,19 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): if use_kernel: if self.layout == "layer_first": transfer_kv_per_layer_mla( - src=device_pool.index_k_with_scale_buffer[layer_id], - dst=self.index_k_with_scale_buffer[layer_id], + src=device_pool.index_k_with_scale_buffer[device_layer_slot], + dst=self.index_k_with_scale_buffer[host_layer_slot], src_indices=device_page_indices, dst_indices=host_page_indices, item_size=self.indexer_page_stride_size, ) elif self.layout == "page_first": _load_tai_transfer_kv_per_layer_mla_lf_pf()( - device_pool.index_k_with_scale_buffer[layer_id], + device_pool.index_k_with_scale_buffer[device_layer_slot], self.index_k_with_scale_buffer, device_page_indices, host_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, item_size=self.indexer_page_stride_size, dst_layout_dim=self.indexer_layout_dim, ) @@ -2121,28 +2218,34 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( - src_layers=[device_pool.index_k_with_scale_buffer[layer_id]], - dst_layers=[self.index_k_with_scale_buffer[layer_id]], + src_layers=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], + dst_layers=[self.index_k_with_scale_buffer[host_layer_slot]], src_indices=device_page_indices, dst_indices=host_page_indices, page_size=1, ) elif self.layout == "page_first_direct": _load_tai_transfer_kv_per_layer_direct_lf_pf()( - src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]], + src_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], dst_ptrs=[self.index_k_with_scale_buffer], src_indices=device_page_indices, dst_indices=host_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, page_size=1, ) elif self.layout == "layer_page_first": _load_tai_transfer_kv_per_layer_direct_lf_lpf()( - src_ptrs=[device_pool.index_k_with_scale_buffer[layer_id]], + src_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], dst_ptrs=[self.index_k_with_scale_buffer], src_indices=device_page_indices, dst_indices=host_page_indices, - layer_id=layer_id, + layer_id=host_layer_slot, page_size=1, ) else: diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 42c650ded..ab70da97b 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Optional, Tuple import torch +from sglang.srt.configs.nsa_index_layers import build_nsa_index_layer_plan from sglang.srt.configs.model_config import get_nsa_index_head_dim, is_deepseek_nsa from sglang.srt.distributed.parallel_state import get_world_group from sglang.srt.layers.dp_attention import get_attention_cp_rank, get_attention_tp_size @@ -494,6 +495,12 @@ class ModelRunnerKVCacheMixin: if self.server_args.enable_nsa_prefill_cp_shared_kv else self.max_total_num_tokens ) + index_layer_plan = build_nsa_index_layer_plan( + self.model_config.hf_config, + self.start_layer, + self.end_layer, + is_nextn=self.is_draft_worker, + ) nsa_pool_kwargs = dict( size=physical_kv_pool_size, page_size=self.page_size, @@ -507,6 +514,7 @@ class ModelRunnerKVCacheMixin: start_layer=self.start_layer, end_layer=self.end_layer, index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config), + index_active_layer_ids=index_layer_plan.active_layer_ids, ) if self.enable_hisparse: from sglang.srt.mem_cache.sparsity import parse_hisparse_config diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index a6ecbf5cc..01609d1d5 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -32,6 +32,7 @@ from sglang.srt.batch_overlap.two_batch_overlap import ( MaybeTboDeepEPDispatcher, model_forward_maybe_tbo, ) +from sglang.srt.configs.nsa_index_layers import nsa_index_skip_flags from sglang.srt.configs.model_config import ( compute_mla_mscale_scaling, get_nsa_index_head_dim, @@ -1214,44 +1215,9 @@ class DeepseekV2AttentionMLA( # Refer: https://arxiv.org/abs/2603.12201 for more details. # skip_topk: when True, this layer will skip computation and reuse previous layer's topk indices. # next_skip_topk: when True, the next layer will skip computation and reuse this layer's topk indices. - if is_nextn: - self.skip_topk = True - self.next_skip_topk = True - else: - self.index_topk_freq = getattr(config, "index_topk_freq", 1) - self.index_topk_pattern = getattr(config, "index_topk_pattern", None) - self.index_skip_topk_offset = getattr( - config, "index_skip_topk_offset", None - ) - if ( - self.index_topk_pattern is None - and self.index_skip_topk_offset is not None - ): - assert self.index_skip_topk_offset > 0, ( - "index_skip_topk_offset must be positive; offset <= 0 " - "marks layer 0 as skip_topk with no prior topk to reuse" - ) - self.skip_topk = ( - max(layer_id - self.index_skip_topk_offset + 1, 0) - % self.index_topk_freq - != 0 - ) - self.next_skip_topk = ( - max(layer_id - self.index_skip_topk_offset + 2, 0) - % self.index_topk_freq - != 0 - ) - elif self.index_topk_pattern is None: - self.skip_topk = max(layer_id - 1, 0) % self.index_topk_freq != 0 - self.next_skip_topk = layer_id % self.index_topk_freq != 0 - else: - self.skip_topk = self.index_topk_pattern[layer_id] == "S" - if layer_id < len(self.index_topk_pattern) - 1: - self.next_skip_topk = ( - self.index_topk_pattern[layer_id + 1] == "S" - ) - else: - self.next_skip_topk = False + self.skip_topk, self.next_skip_topk = nsa_index_skip_flags( + config, layer_id, is_nextn=is_nextn + ) self.kv_b_proj = ColumnParallelLinear( self.kv_lora_rank, diff --git a/test/registered/unit/configs/test_nsa_index_layers.py b/test/registered/unit/configs/test_nsa_index_layers.py new file mode 100644 index 000000000..fca72d1c9 --- /dev/null +++ b/test/registered/unit/configs/test_nsa_index_layers.py @@ -0,0 +1,83 @@ +from types import SimpleNamespace + +import pytest + +from sglang.srt.configs.nsa_index_layers import ( + build_nsa_index_layer_plan, + nsa_index_skip_flags, +) + + +def test_default_freq_one_marks_all_target_layers_active(): + cfg = SimpleNamespace(index_topk_freq=1) + plan = build_nsa_index_layer_plan(cfg, 0, 6) + assert plan.active_layer_ids == (0, 1, 2, 3, 4, 5) + assert [nsa_index_skip_flags(cfg, i)[0] for i in range(6)] == [False] * 6 + + +def test_freq_four_without_offset_matches_current_model_formula(): + cfg = SimpleNamespace(index_topk_freq=4) + plan = build_nsa_index_layer_plan(cfg, 0, 12) + assert plan.active_layer_ids == (0, 1, 5, 9) + assert [nsa_index_skip_flags(cfg, i)[0] for i in range(10)] == [ + False, + False, + True, + True, + True, + False, + True, + True, + True, + False, + ] + + +def test_freq_four_with_offset_one_uses_layers_zero_four_eight(): + cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1) + plan = build_nsa_index_layer_plan(cfg, 0, 12) + assert plan.active_layer_ids == (0, 4, 8) + assert [nsa_index_skip_flags(cfg, i)[0] for i in range(9)] == [ + False, + True, + True, + True, + False, + True, + True, + True, + False, + ] + + +def test_pattern_marks_non_shared_layers_active(): + cfg = SimpleNamespace(index_topk_freq=1, index_topk_pattern="CSSSCSS") + plan = build_nsa_index_layer_plan(cfg, 0, len(cfg.index_topk_pattern)) + assert plan.active_layer_ids == (0, 4) + + +def test_nextn_keeps_all_draft_layers_active_for_state_safety(): + cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1) + plan = build_nsa_index_layer_plan(cfg, 0, 1, is_nextn=True) + assert plan.active_layer_ids == (0,) + assert nsa_index_skip_flags(cfg, 0, is_nextn=True) == (True, True) + + +def test_nonzero_start_layer_preserves_logical_layer_ids(): + cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1) + plan = build_nsa_index_layer_plan(cfg, 4, 13) + assert plan.active_layer_ids == (4, 8, 12) + assert plan.slot_for_layer(8) == 1 + + +def test_inactive_slot_lookup_fails_fast(): + cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=1) + plan = build_nsa_index_layer_plan(cfg, 0, 8) + with pytest.raises(RuntimeError, match="inactive index layer requested"): + plan.slot_for_layer(1) + + +def test_invalid_offset_fails_before_layer_zero_can_skip_without_prior_topk(): + cfg = SimpleNamespace(index_topk_freq=4, index_skip_topk_offset=0) + with pytest.raises(ValueError, match="index_skip_topk_offset"): + nsa_index_skip_flags(cfg, 0) diff --git a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py index 5b4507080..da2db2f8c 100644 --- a/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py +++ b/test/registered/unit/mem_cache/test_nsa_pool_host_unit.py @@ -1684,6 +1684,41 @@ class TestNSAIndexerPageIndices(CustomTestCase): host_pool.page_size = page_size return host_pool + def test_nsa_device_pool_active_index_layers_use_full_allocation_slots(self): + pool = object.__new__(NSATokenToKVPool) + pool.start_layer = 4 + pool.end_layer = 12 + pool.layer_num = 8 + + pool._init_index_layer_metadata( + index_active_layer_ids=(4, 8), + compact_index_layers=False, + ) + + self.assertEqual(pool.index_active_layer_ids, (4, 8)) + self.assertTrue(pool.is_index_layer_active(4)) + self.assertTrue(pool.is_index_layer_active(8)) + self.assertFalse(pool.is_index_layer_active(5)) + self.assertEqual(pool.get_index_layer_slot(4), 0) + self.assertEqual(pool.get_index_layer_slot(8), 4) + with self.assertRaisesRegex(RuntimeError, "inactive index layer requested"): + pool.get_index_layer_slot(5) + + def test_nsa_device_pool_default_active_layers_cover_local_range(self): + pool = object.__new__(NSATokenToKVPool) + pool.start_layer = 4 + pool.end_layer = 12 + pool.layer_num = 8 + + pool._init_index_layer_metadata( + index_active_layer_ids=None, + compact_index_layers=False, + ) + + self.assertEqual(pool.index_active_layer_ids, tuple(range(4, 12))) + self.assertTrue(pool.is_index_layer_active(11)) + self.assertEqual(pool.get_index_layer_slot(11), 7) + def test_indexer_page_indices_accepts_valid_page_spans(self): host_pool = self.make_host_pool_stub(page_size=4) @@ -1773,6 +1808,118 @@ class TestNSAIndexerPageIndices(CustomTestCase): self.assertEqual(call["dst_indices"].tolist(), [0, 1]) self.assertEqual(call["page_size"], 1) + def test_page_first_direct_all_layer_indexer_backup_skips_inactive_layers(self): + host_pool = self.make_host_pool_stub(page_size=4) + host_pool.layout = "page_first_direct" + host_pool.indexer_page_stride_size = 8 + host_pool.layer_num = 4 + host_pool.index_k_with_scale_buffer = "host-page-first-indexer" + + class FakeDevicePool: + index_active_layer_ids = (0, 2) + index_k_with_scale_buffer = [ + "device-layer-0", + "device-layer-1", + "device-layer-2", + "device-layer-3", + ] + + def is_index_layer_active(self, layer_id): + return layer_id in self.index_active_layer_ids + + def get_index_layer_slot(self, layer_id): + return layer_id + + calls = [] + + def fake_tai_transfer(**kwargs): + calls.append(kwargs) + + with patch( + "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf", + return_value=fake_tai_transfer, + ): + host_pool._backup_indexer_from_device_all_layer( + FakeDevicePool(), + torch.tensor([0, 1, 2, 3], dtype=torch.int64), + torch.tensor([8, 9, 10, 11], dtype=torch.int64), + "direct", + ) + + self.assertEqual([call["layer_id"] for call in calls], [0, 2]) + self.assertEqual(calls[0]["src_ptrs"], ["device-layer-0"]) + self.assertEqual(calls[1]["src_ptrs"], ["device-layer-2"]) + + def test_per_layer_indexer_backup_and_load_skip_inactive_layers(self): + host_pool = self.make_host_pool_stub(page_size=4) + host_pool.layout = "page_first_direct" + host_pool.indexer_page_stride_size = 8 + host_pool.index_k_with_scale_buffer = "host-page-first-indexer" + + class FakeDevicePool: + index_active_layer_ids = (2,) + index_k_with_scale_buffer = [ + "device-layer-0", + "device-layer-1", + "device-layer-2", + ] + + def is_index_layer_active(self, layer_id): + return layer_id in self.index_active_layer_ids + + def get_index_layer_slot(self, layer_id): + return layer_id + + calls = [] + + def fake_backup_transfer(**kwargs): + calls.append(("backup", kwargs)) + + def fake_load_transfer(**kwargs): + calls.append(("load", kwargs)) + + with ( + patch( + "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_lf_pf", + return_value=fake_backup_transfer, + ), + patch( + "sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_direct_pf_lf", + return_value=fake_load_transfer, + ), + ): + host_pool._backup_indexer_from_device_per_layer( + FakeDevicePool(), + torch.tensor([0, 1, 2, 3], dtype=torch.int64), + torch.tensor([8, 9, 10, 11], dtype=torch.int64), + 1, + "direct", + ) + host_pool._load_indexer_to_device_per_layer( + FakeDevicePool(), + torch.tensor([0, 1, 2, 3], dtype=torch.int64), + torch.tensor([8, 9, 10, 11], dtype=torch.int64), + 1, + "direct", + ) + host_pool._backup_indexer_from_device_per_layer( + FakeDevicePool(), + torch.tensor([0, 1, 2, 3], dtype=torch.int64), + torch.tensor([8, 9, 10, 11], dtype=torch.int64), + 2, + "direct", + ) + host_pool._load_indexer_to_device_per_layer( + FakeDevicePool(), + torch.tensor([0, 1, 2, 3], dtype=torch.int64), + torch.tensor([8, 9, 10, 11], dtype=torch.int64), + 2, + "direct", + ) + + self.assertEqual([kind for kind, _ in calls], ["backup", "load"]) + self.assertEqual([kwargs["layer_id"] for _, kwargs in calls], [2, 2]) + def test_mla_layer_page_first_all_layer_backup_uses_tai_per_layer_route(self): host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) host_pool.layout = "layer_page_first"