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
python/sglang/srt/models/deepseek_v2.py
+4 -38
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,