A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf] Optimize DSv4 prefill chunk planning, 4.0% E2E Throughput Improvement (#45061)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 15, 2026 at 15:50 UTC
e18fe932ca61fbdcf9575989c75fefa8ff8d701b
3 files changed
+133
-23
tests/kernels/attention/test_flashmla_sparse.py
+21
@@ -122,3 +122,24 @@ def test_sparse_flashmla_prefill_smoke():
122
assert out.shape == (s_q, h_q, d_v)
123
assert max_logits.shape == (s_q, h_q)
124
assert lse.shape == (s_q, h_q)
125
+
126
+
127
+def test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences():
128
+ from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
129
+
130
+ metadata = DeepseekSparseSWAMetadata(
131
+ block_table=torch.empty(0, dtype=torch.int32),
132
+ slot_mapping=torch.empty(0, dtype=torch.int32),
133
+ block_size=64,
134
+ num_prefills=5,
135
+ prefill_seq_lens_cpu=torch.tensor([80, 96, 112, 128, 144], dtype=torch.int32),
136
+ prefill_query_lens_cpu=torch.tensor([4, 4, 4, 4, 4], dtype=torch.int32),
137
+ prefill_window_size=64,
138
+ prefill_max_model_len=1024,
139
+ prefill_max_num_batched_tokens=128,
140
+ )
141
+
142
+ chunk_plan = metadata.get_prefill_chunk_plan(compress_ratio=4, prefill_chunk_size=4)
143
+
144
+ # the adaptive plan keeps all 5 in one chunk
145
+ assert chunk_plan == [(0, 5, 36, 103)]
vllm/models/deepseek_v4/nvidia/flashmla.py
+12
-20
@@ -246,7 +246,6 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
246
) -> None:
247
swa_only = attn_metadata is None
248
249
- num_prefills = swa_metadata.num_prefills
249
num_prefill_tokens = swa_metadata.num_prefill_tokens
250
num_decodes = swa_metadata.num_decodes
251
num_decode_tokens = swa_metadata.num_decode_tokens
@@ -274,29 +273,22 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
273
assert attn_metadata is not None
274
topk_indices = attn_metadata.c128a_prefill_topk_indices
275
top_k = topk_indices.shape[-1]
277
- # Compressed region must fit the full compressed pool (seq_len //
278
- # compress_ratio), not just top_k. top_k bounds how many indices
279
- # the indexer selects, not the pool size it indexes into.
280
- N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio
276
else:
277
# NOTE(woosuk): topk_indices will not be used for SWA-only layers.
278
assert self.topk_indices_buffer is not None
279
topk_indices = self.topk_indices_buffer[num_decode_tokens:]
280
top_k = 0
286
- N = 0
287
-
288
- M = N + self.window_size + self.max_num_batched_tokens
289
- chunk_size_const = self.PREFILL_CHUNK_SIZE
290
- num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const
291
-
281
+ chunk_plan = swa_metadata.get_prefill_chunk_plan(
282
+ compress_ratio=self.compress_ratio,
283
+ prefill_chunk_size=self.PREFILL_CHUNK_SIZE,
284
+ )
285
+ assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
286
workspace_manager = current_workspace_manager()
293
- kv = workspace_manager.get_simultaneous(
294
- ((chunk_size_const, M, q.shape[-1]), torch.bfloat16),
295
- )[0]
296
- for chunk_idx in range(num_chunks):
297
- chunk_start = chunk_idx * chunk_size_const
298
- chunk_end = min(chunk_start + chunk_size_const, num_prefills)
287
+ for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
288
chunk_size = chunk_end - chunk_start
289
+ kv = workspace_manager.get_simultaneous(
290
+ ((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16),
291
+ )[0]
292
if not swa_only:
293
# Gather compressed KV
294
assert attn_metadata is not None
@@ -320,7 +312,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
312
gather_lens=gather_lens[chunk_start:chunk_end],
313
block_table=swa_block_table[chunk_start:chunk_end],
314
block_size=swa_metadata.block_size,
323
- offset=N,
315
+ offset=chunk_N,
316
)
317
318
# Combine the topk indices and SWA indices for gathered KV cache
@@ -341,8 +333,8 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
333
self.window_size,
334
self.compress_ratio,
335
top_k,
344
- M,
345
- N,
336
+ chunk_M,
337
+ chunk_N,
338
)
339
flash_mla_sparse_fwd(
340
q=q[query_start:query_end],
vllm/v1/attention/backends/mla/sparse_swa.py
+100
-3
@@ -9,6 +9,7 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
9
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10
from vllm.platforms import current_platform
11
from vllm.triton_utils import tl, triton
12
+from vllm.utils.math_utils import cdiv
13
from vllm.v1.attention.backend import (
14
AttentionBackend,
15
AttentionCGSupport,
@@ -172,7 +173,12 @@ class DeepseekSparseSWAMetadata:
173
174
# Pre-computed prefill metadata shared across all DeepseekV4 attention layers.
175
prefill_seq_lens: torch.Tensor | None = None
176
+ prefill_seq_lens_cpu: torch.Tensor | None = None
177
prefill_gather_lens: torch.Tensor | None = None
178
+ prefill_query_lens_cpu: torch.Tensor | None = None
179
+ prefill_window_size: int = 0
180
+ prefill_max_model_len: int = 0
181
+ prefill_max_num_batched_tokens: int = 0
182
183
# Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta
184
# per present DeepseekV4 layer type, shared across all ~60 layers of that type
@@ -188,6 +194,79 @@ class DeepseekSparseSWAMetadata:
194
tile_sched_c4a: "FlashMLASchedMeta | None" = None
195
tile_sched_c128a: "FlashMLASchedMeta | None" = None
196
197
+ def get_prefill_chunk_plan(
198
+ self, compress_ratio: int, prefill_chunk_size: int
199
+ ) -> list[tuple[int, int, int, int]]:
200
+ if self.num_prefills == 0:
201
+ return []
202
+
203
+ assert self.prefill_seq_lens_cpu is not None
204
+ assert self.prefill_query_lens_cpu is not None
205
+
206
+ # query_len <= max_num_batched_tokens and
207
+ # gather_len = query_len + min(prefix_len, window_size - 1), so the
208
+ # worst-case gathered width is bounded by
209
+ # max_num_batched_tokens + window_size - 1. The compressed prefix pool
210
+ # is bounded by ceil(max_model_len / compress_ratio).
211
+ max_workspace_area = prefill_chunk_size * (
212
+ (
213
+ 0
214
+ if compress_ratio <= 1
215
+ else cdiv(self.prefill_max_model_len, compress_ratio)
216
+ )
217
+ + self.prefill_window_size
218
+ + self.prefill_max_num_batched_tokens
219
+ )
220
+ prefix_lens_cpu = self.prefill_seq_lens_cpu - self.prefill_query_lens_cpu
221
+ gather_lens_cpu = self.prefill_query_lens_cpu + torch.clamp(
222
+ prefix_lens_cpu, min=0, max=self.prefill_window_size - 1
223
+ )
224
+ compressed_lens_cpu = (
225
+ torch.zeros_like(self.prefill_seq_lens_cpu)
226
+ if compress_ratio <= 1
227
+ else torch.div(
228
+ self.prefill_seq_lens_cpu,
229
+ compress_ratio,
230
+ rounding_mode="floor",
231
+ )
232
+ )
233
+
234
+ chunk_plan: list[tuple[int, int, int, int]] = []
235
+ chunk_start = 0
236
+ while chunk_start < self.num_prefills:
237
+ chunk_max_compressed = int(compressed_lens_cpu[chunk_start].item())
238
+ chunk_max_gather = int(gather_lens_cpu[chunk_start].item())
239
+ chunk_end = chunk_start + 1
240
+
241
+ while chunk_end < self.num_prefills:
242
+ candidate_max_compressed = max(
243
+ chunk_max_compressed,
244
+ int(compressed_lens_cpu[chunk_end].item()),
245
+ )
246
+ candidate_max_gather = max(
247
+ chunk_max_gather,
248
+ int(gather_lens_cpu[chunk_end].item()),
249
+ )
250
+ candidate_width = candidate_max_compressed + candidate_max_gather
251
+ candidate_area = (chunk_end - chunk_start + 1) * candidate_width
252
+ if candidate_area > max_workspace_area:
253
+ break
254
+ chunk_max_compressed = candidate_max_compressed
255
+ chunk_max_gather = candidate_max_gather
256
+ chunk_end += 1
257
+
258
+ chunk_plan.append(
259
+ (
260
+ chunk_start,
261
+ chunk_end,
262
+ chunk_max_compressed,
263
+ chunk_max_compressed + chunk_max_gather,
264
+ )
265
+ )
266
+ chunk_start = chunk_end
267
+
268
+ return chunk_plan
269
+
270
271
class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
272
"""Builds metadata for DeepseekV4 SWA cache.
@@ -213,6 +292,10 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
292
self.head_size = mla_spec.head_size # Already considered quantization.
293
self.compress_ratio = mla_spec.compress_ratio
294
self.block_size = mla_spec.block_size
295
+ self.max_model_len = self.vllm_config.model_config.max_model_len
296
+ self.max_num_batched_tokens = (
297
+ self.vllm_config.scheduler_config.max_num_batched_tokens
298
+ )
299
300
# Handle MTP: adjust decode_threshold like the indexer does
301
self.num_speculative_tokens = (
@@ -279,6 +362,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
362
"""
363
num_reqs = common_attn_metadata.num_reqs
364
seq_lens = common_attn_metadata.seq_lens
365
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
366
query_start_loc = common_attn_metadata.query_start_loc
367
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
368
block_table = common_attn_metadata.block_table_tensor
@@ -323,7 +407,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
407
num_decodes,
408
num_prefills,
409
seq_lens,
410
+ seq_lens_cpu,
411
query_start_loc,
412
+ query_start_loc_cpu,
413
)
414
415
# Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta
@@ -350,7 +436,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
436
tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY],
437
tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A],
438
tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A],
353
- **deepseek_v4_fields,
439
+ **deepseek_v4_fields, # type: ignore[arg-type]
440
)
441
442
def build_tile_scheduler(
@@ -391,8 +477,10 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
477
num_decodes: int,
478
num_prefills: int,
479
seq_lens: torch.Tensor,
480
+ seq_lens_cpu: torch.Tensor | None,
481
query_start_loc: torch.Tensor,
395
- ) -> dict[str, torch.Tensor | None]:
482
+ query_start_loc_cpu: torch.Tensor,
483
+ ) -> dict[str, torch.Tensor | int | None]:
484
"""Pre-compute DeepseekV4 prefill metadata during the metadata build phase.
485
486
Returns a dict of keyword arguments to pass to the
@@ -401,10 +489,11 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
489
Note: C128A topk indices are computed by the FlashMLASparse builder
490
(which owns the C128A block_table), not here.
491
"""
404
- result: dict[str, torch.Tensor | None] = {}
492
+ result: dict[str, torch.Tensor | int | None] = {}
493
494
# --- Prefill query metadata (single Triton kernel + CPU slicing) ---
495
if num_prefills > 0:
496
+ assert seq_lens_cpu is not None
497
pfx_gather_lens = torch.empty(
498
num_prefills, dtype=torch.int32, device=seq_lens.device
499
)
@@ -419,7 +508,15 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
508
)
509
510
result["prefill_seq_lens"] = seq_lens[num_decodes:]
511
+ result["prefill_seq_lens_cpu"] = seq_lens_cpu[num_decodes:]
512
result["prefill_gather_lens"] = pfx_gather_lens
513
+ result["prefill_query_lens_cpu"] = (
514
+ query_start_loc_cpu[num_decodes + 1 : num_decodes + num_prefills + 1]
515
+ - query_start_loc_cpu[num_decodes : num_decodes + num_prefills]
516
+ ).to(dtype=torch.int32)
517
+ result["prefill_window_size"] = self.window_size
518
+ result["prefill_max_model_len"] = self.max_model_len
519
+ result["prefill_max_num_batched_tokens"] = self.max_num_batched_tokens
520
521
return result
522