A high-throughput and memory-efficient inference and serving engine for LLMs
[DSv4 Perf] DSv4 flashinfer sparse index cache for metadata, 2%~4% TTFT improvement (#45863)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 17, 2026 at 10:55 UTC
0a7bacdcacc5f5c7e511f596b1f6b372da6029b6
3 files changed
+184
-18
tests/kernels/attention/test_flashmla_sparse.py
+147
@@ -143,3 +143,150 @@ def test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences():
143
144
# the adaptive plan keeps all 5 in one chunk
145
assert chunk_plan == [(0, 5, 36, 103)]
146
+
147
+
148
+def test_flashinfer_sparse_indices_cache(monkeypatch):
149
+ from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as flashinfer_mod
150
+ from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLAMetadata
151
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
152
+
153
+ builder_calls = 0
154
+
155
+ def fake_build(*args, **kwargs):
156
+ nonlocal builder_calls
157
+ builder_calls += 1
158
+ return (
159
+ torch.tensor([[builder_calls]], dtype=torch.int32),
160
+ torch.tensor([builder_calls], dtype=torch.int32),
161
+ )
162
+
163
+ monkeypatch.setattr(
164
+ flashinfer_mod, "build_flashinfer_mixed_sparse_indices", fake_build
165
+ )
166
+
167
+ def make_attn(compress_ratio: int, topk_width: int):
168
+ attn = object.__new__(flashinfer_mod.DeepseekV4FlashInferMLAAttention)
169
+ attn.compress_ratio = compress_ratio
170
+ attn.window_size = 4
171
+ attn.topk_indices_buffer = torch.tensor(
172
+ [[0, 1], [2, 3], [4, 5]], dtype=torch.int32
173
+ )[:, :topk_width]
174
+ return attn
175
+
176
+ def make_swa_metadata():
177
+ return DeepseekSparseSWAMetadata(
178
+ block_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32),
179
+ slot_mapping=torch.tensor([0, 1], dtype=torch.int64),
180
+ block_size=64,
181
+ seq_lens=torch.tensor([8, 10], dtype=torch.int32),
182
+ query_start_loc=torch.tensor([0, 1, 3], dtype=torch.int32),
183
+ query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.int32),
184
+ token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32),
185
+ decode_swa_indices=torch.tensor([[5, 6, -1, -1]], dtype=torch.int32),
186
+ decode_swa_lens=torch.tensor([2], dtype=torch.int32),
187
+ is_valid_token=torch.tensor([True], dtype=torch.bool),
188
+ num_decodes=1,
189
+ num_prefills=1,
190
+ num_decode_tokens=1,
191
+ num_prefill_tokens=2,
192
+ )
193
+
194
+ def make_flashmla_metadata():
195
+ return DeepseekV4FlashMLAMetadata(
196
+ num_reqs=2,
197
+ max_query_len=2,
198
+ max_seq_len=10,
199
+ num_actual_tokens=3,
200
+ query_start_loc=torch.tensor([0, 1, 3], dtype=torch.int32),
201
+ slot_mapping=torch.tensor([0, 1, 2], dtype=torch.int64),
202
+ block_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32),
203
+ req_id_per_token=torch.tensor([0, 1, 1], dtype=torch.int32),
204
+ block_size=256,
205
+ topk_tokens=2,
206
+ c128a_global_decode_topk_indices=torch.tensor(
207
+ [[[9, 10]]], dtype=torch.int32
208
+ ),
209
+ c128a_decode_topk_lens=torch.tensor([2], dtype=torch.int32),
210
+ c128a_prefill_topk_indices=torch.tensor(
211
+ [[0, 1], [1, 2]], dtype=torch.int32
212
+ ),
213
+ )
214
+
215
+ swa_attn = make_attn(1, 0)
216
+ swa_metadata = make_swa_metadata()
217
+ _, _, sparse_indices_first, sparse_lens_first = (
218
+ swa_attn._build_sparse_index_metadata(
219
+ kv_cache=None,
220
+ swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
221
+ swa_metadata=swa_metadata,
222
+ attn_metadata=None,
223
+ swa_only=True,
224
+ )
225
+ )
226
+ _, _, sparse_indices_second, sparse_lens_second = (
227
+ swa_attn._build_sparse_index_metadata(
228
+ kv_cache=None,
229
+ swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
230
+ swa_metadata=swa_metadata,
231
+ attn_metadata=None,
232
+ swa_only=True,
233
+ )
234
+ )
235
+ assert builder_calls == 1
236
+ assert sparse_indices_first is sparse_indices_second
237
+ assert sparse_lens_first is sparse_lens_second
238
+
239
+ c128a_attn = make_attn(128, 2)
240
+ c128a_metadata = make_swa_metadata()
241
+ c128a_flashmla_md = make_flashmla_metadata()
242
+ _, _, sparse_indices_first, sparse_lens_first = (
243
+ c128a_attn._build_sparse_index_metadata(
244
+ kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
245
+ swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
246
+ swa_metadata=c128a_metadata,
247
+ attn_metadata=c128a_flashmla_md,
248
+ swa_only=False,
249
+ )
250
+ )
251
+ _, _, sparse_indices_second, sparse_lens_second = (
252
+ c128a_attn._build_sparse_index_metadata(
253
+ kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
254
+ swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
255
+ swa_metadata=c128a_metadata,
256
+ attn_metadata=c128a_flashmla_md,
257
+ swa_only=False,
258
+ )
259
+ )
260
+
261
+ assert builder_calls == 2
262
+ assert sparse_indices_first is sparse_indices_second
263
+ assert sparse_lens_first is sparse_lens_second
264
+
265
+ c4a_attn = make_attn(4, 2)
266
+ c4a_metadata = make_swa_metadata()
267
+ c4a_flashmla_md = make_flashmla_metadata()
268
+ c4a_flashmla_md.c128a_global_decode_topk_indices = None
269
+ c4a_flashmla_md.c128a_decode_topk_lens = None
270
+ c4a_flashmla_md.c128a_prefill_topk_indices = None
271
+ _, _, sparse_indices_third, sparse_lens_third = (
272
+ c4a_attn._build_sparse_index_metadata(
273
+ kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
274
+ swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
275
+ swa_metadata=c4a_metadata,
276
+ attn_metadata=c4a_flashmla_md,
277
+ swa_only=False,
278
+ )
279
+ )
280
+ _, _, sparse_indices_fourth, sparse_lens_fourth = (
281
+ c4a_attn._build_sparse_index_metadata(
282
+ kv_cache=torch.empty((1, 2, 512), dtype=torch.bfloat16),
283
+ swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16),
284
+ swa_metadata=c4a_metadata,
285
+ attn_metadata=c4a_flashmla_md,
286
+ swa_only=False,
287
+ )
288
+ )
289
+
290
+ assert builder_calls == 4
291
+ assert sparse_indices_third is not sparse_indices_fourth
292
+ assert sparse_lens_third is not sparse_lens_fourth
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+33
-17
@@ -288,24 +288,40 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
288
query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1]
289
seq_lens = swa_metadata.seq_lens[:num_reqs]
290
assert seq_lens.dtype == torch.int32
291
- sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
292
- decode_swa_indices,
293
- decode_compressed_indices,
294
- decode_compressed_topk_lens,
295
- prefill_topk_indices[:num_prefill_tokens],
296
- query_start_loc,
297
- seq_lens,
298
- swa_metadata.token_to_req_indices[:num_tokens],
299
- swa_metadata.block_table[:num_reqs],
300
- swa_metadata.block_size,
301
- compressed_block_table,
302
- compressed_block_size,
303
- self.window_size,
304
- self.compress_ratio,
305
- top_k,
306
- decode_compressed_indices_are_local=decode_compressed_indices_are_local,
307
- decode_is_valid_token=decode_is_valid_token,
291
+ # cache for SWA-only and C128A that build the same mixed sparse indices
292
+ # C4A stays uncached.
293
+ cache_key = (
294
+ "swa_only"
295
+ if swa_only
296
+ else ("c128a" if self.compress_ratio == 128 else "c4a")
297
)
298
+ cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None)
299
+ if cached_sparse is None:
300
+ sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
301
+ decode_swa_indices,
302
+ decode_compressed_indices,
303
+ decode_compressed_topk_lens,
304
+ prefill_topk_indices[:num_prefill_tokens],
305
+ query_start_loc,
306
+ seq_lens,
307
+ swa_metadata.token_to_req_indices[:num_tokens],
308
+ swa_metadata.block_table[:num_reqs],
309
+ swa_metadata.block_size,
310
+ compressed_block_table,
311
+ compressed_block_size,
312
+ self.window_size,
313
+ self.compress_ratio,
314
+ top_k,
315
+ decode_compressed_indices_are_local=decode_compressed_indices_are_local,
316
+ decode_is_valid_token=decode_is_valid_token,
317
+ )
318
+ if cache_key != "c4a":
319
+ swa_metadata.flashinfer_sparse_index_cache[cache_key] = (
320
+ sparse_indices,
321
+ sparse_topk_lens,
322
+ )
323
+ else:
324
+ sparse_indices, sparse_topk_lens = cached_sparse
325
return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens
326
327
def _forward(
vllm/v1/attention/backends/mla/sparse_swa.py
+4
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
3
+from dataclasses import dataclass, field
4
from typing import ClassVar, cast
5
6
import torch
@@ -193,6 +193,9 @@ class DeepseekSparseSWAMetadata:
193
tile_sched_swaonly: "FlashMLASchedMeta | None" = None
194
tile_sched_c4a: "FlashMLASchedMeta | None" = None
195
tile_sched_c128a: "FlashMLASchedMeta | None" = None
196
+ flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field(
197
+ default_factory=dict
198
+ )
199
200
def get_prefill_chunk_plan(
201
self, compress_ratio: int, prefill_chunk_size: int