A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 12:30 UTC
[Feature]: IndexCache support for DSA models (#37735)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Chauncey committed
Apr 30, 2026 at 03:15 UTC
faab18955407f256c7ced2d227ce097f472db16d
3 files changed
+83
-5
docs/features/index_cache.md
new
+54
@@ -0,0 +1,54 @@
1
+# IndexCache
2
+
3
+IndexCache reduces redundant top-k computation in DeepSeek-V3.2 (DSA) models by caching and reusing top-k indices across layers.
4
+
5
+## Background
6
+
7
+DeepSeek-V3.2 uses a DeepSeek Sparse Attention (DSA) mechanism where top-k token selection is computed per layer. For deep models with many layers, this computation can be expensive. IndexCache allows skipping redundant top-k computations by reusing indices from previous layers.
8
+
9
+See: [IndexCache Paper](https://arxiv.org/abs/2603.12201)
10
+
11
+## Usage
12
+
13
+### CLI
14
+
15
+```bash
16
+vllm serve deepseek-ai/DeepSeek-V3.2 \
17
+ --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
18
+```
19
+
20
+### Configuration Reference
21
+
22
+| Parameter | Type | Default | Description |
23
+|----------------------|------|---------|--------------------------------------------------------------------------------------------------------------------------------------------------|
24
+| `use_index_cache` | bool | false | Enable IndexCache. Must be set to true to use this feature |
25
+| `index_topk_freq` | int | 1 | Frequency (in layers) at which top-k is computed. 1 = compute on every layer (disabled), 4 = compute on 1/4 of layers |
26
+| `index_topk_pattern` | str | null | Per-layer F/S pattern. Overrides index_topk_freq if set. Each character maps to one DSA layer: F = Full, S = Shared |
27
+
28
+### Configuration Examples
29
+
30
+**Using `index_topk_freq`** (compute every N layers):
31
+
32
+```bash
33
+vllm serve deepseek-ai/DeepSeek-V3.2 \
34
+ --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
35
+```
36
+
37
+**Using `index_topk_pattern`** (explicit per-layer control):
38
+
39
+```bash
40
+# custom pattern for 61 layers: F = compute, S = reuse
41
+vllm serve deepseek-ai/DeepSeek-V3.2 \
42
+ --hf-overrides '{"use_index_cache": true, "index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSF"}'
43
+```
44
+
45
+## How It Works
46
+
47
+1. When IndexCache is enabled, layers marked with `"F"` (Full) calculate and store top-k indices
48
+2. Subsequent layers marked with `"S"` (Shared) receive the cached indices from the previous layer instead of recomputing
49
+3. The cached indices are passed through the layer stack, reducing total computation
50
+
51
+## Requirements
52
+
53
+- DeepSeek-V3.2 or compatible DSA model
54
+- `use_index_cache: true` via `--hf-overrides`
vllm/model_executor/layers/mla.py
+8
-4
@@ -64,6 +64,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
64
cache_config: CacheConfig | None = None,
65
quant_config: QuantizationConfig | None = None,
66
prefix: str = "",
67
+ skip_topk: bool = False,
68
) -> None:
69
super().__init__()
70
self.hidden_size = hidden_size
@@ -87,6 +88,11 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
88
self.indexer_rope_emb = mla_modules.indexer_rotary_emb
89
self.is_sparse = mla_modules.is_sparse
90
91
+ # Whether to skip top-k token selection computation in this layer.
92
+ # When True, the indexer will not be called, and the layer will reuse
93
+ # the topk_tokens buffer written by a previous layer in the same pass.
94
+ # Refer: https://arxiv.org/abs/2603.12201 for more details.
95
+ self.skip_topk = skip_topk
96
if self.indexer is not None:
97
assert hasattr(self.indexer, "topk_tokens")
98
self.topk_tokens = self.indexer.topk_tokens
@@ -159,10 +165,8 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
165
positions, q[..., self.qk_nope_head_dim :], k_pe
166
)
167
162
- if self.indexer and self.is_sparse:
163
- _topk_indices = self.indexer(
164
- hidden_states, q_c, positions, self.indexer_rope_emb
165
- )
168
+ if self.indexer and self.is_sparse and not self.skip_topk:
169
+ self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)
170
171
if llama_4_scaling is not None:
172
q *= llama_4_scaling
vllm/model_executor/models/deepseek_v2.py
+21
-1
@@ -82,7 +82,10 @@ from vllm.model_executor.model_loader.weight_utils import (
82
default_weight_loader,
83
maybe_remap_kv_scale_name,
84
)
85
-from vllm.model_executor.models.utils import sequence_parallel_chunk
85
+from vllm.model_executor.models.utils import (
86
+ extract_layer_index,
87
+ sequence_parallel_chunk,
88
+)
89
from vllm.platforms import current_platform
90
from vllm.sequence import IntermediateTensors
91
from vllm.utils.torch_utils import direct_register_custom_op
@@ -963,6 +966,7 @@ class DeepseekV2MLAAttention(nn.Module):
966
967
self.is_v32 = hasattr(config, "index_topk")
968
969
+ _skip_topk = False
970
if self.is_v32:
971
self.indexer_rope_emb = get_rope(
972
qk_rope_head_dim,
@@ -980,6 +984,21 @@ class DeepseekV2MLAAttention(nn.Module):
984
topk_indices_buffer,
985
f"{prefix}.indexer",
986
)
987
+
988
+ # Enable IndexCache for DeepSeek models to reduce redundant top-k
989
+ # token selection computations in sparse attention.
990
+ use_index_cache = getattr(config, "use_index_cache", False)
991
+ if use_index_cache:
992
+ # IndexCache config
993
+ # Refer: https://arxiv.org/abs/2603.12201 for more details.
994
+ _index_topk_freq = getattr(config, "index_topk_freq", 1)
995
+ _index_topk_pattern = getattr(config, "index_topk_pattern", None)
996
+ layer_id = extract_layer_index(prefix)
997
+ if _index_topk_pattern is None:
998
+ _skip_topk = max(layer_id - 1, 0) % _index_topk_freq != 0
999
+ elif 0 <= layer_id < len(_index_topk_pattern):
1000
+ _skip_topk = _index_topk_pattern[layer_id] == "S"
1001
+
1002
else:
1003
self.indexer_rope_emb = None
1004
self.indexer = None
@@ -1017,6 +1036,7 @@ class DeepseekV2MLAAttention(nn.Module):
1036
cache_config,
1037
quant_config,
1038
prefix,
1039
+ skip_topk=_skip_topk,
1040
)
1041
1042
def forward(