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

Reduce inactive NSA index-cache transfer safely

Browse Source 
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.
This commit is contained in:
laoyao0822
2026-06-10 04:28:26 +08:00
parent 6229c7da60
commit d21952b903
8 changed files with 567 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
python/sglang/srt/configs/nsa_index_layers.py Normal file
View File

@@ -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)
},
)

6
python/sglang/srt/mem_cache/hisparse_memory_pool.py
View File

@@ -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

65
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -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
)

193
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -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:

8
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -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

42
python/sglang/srt/models/deepseek_v2.py
View File

@@ -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,