[DSv4 Perf] Adaptive topk width, 1.0% E2E throughput improvement (#50004)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 27, 2026 at 15:56 UTC
b2f9e4caa49425d93667e01be5c9ad4c45bf81df
2 files changed
+56
-7
tests/kernels/attention/test_flashmla_sparse.py
+37
@@ -4,6 +4,43 @@ import pytest
4
import torch
5
6
7
+def test_deepseek_v4_c128a_dynamic_topk_packed_buffers():
8
+ from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata
9
+
10
+ device = torch.device("cuda")
11
+ capacity_width = 256
12
+ active_width = 128
13
+ global_decode_buffer = torch.empty(
14
+ (2, capacity_width), dtype=torch.int32, device=device
15
+ )
16
+ decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device)
17
+ prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device)
18
+
19
+ global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
20
+ positions=torch.tensor([255, 511], dtype=torch.int64, device=device),
21
+ compress_ratio=128,
22
+ num_decode_tokens=1,
23
+ token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device),
24
+ block_table=torch.tensor([[3]], dtype=torch.int32, device=device),
25
+ block_size=capacity_width,
26
+ slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device),
27
+ global_decode_buffer=global_decode_buffer,
28
+ decode_lens_buffer=decode_lens_buffer,
29
+ prefill_buffer=prefill_buffer,
30
+ max_compressed_tokens=active_width,
31
+ )
32
+
33
+ assert global_decode.shape == (1, active_width)
34
+ assert prefill_local.shape == (1, active_width)
35
+ assert global_decode.stride() == (active_width, 1)
36
+ assert prefill_local.stride() == (active_width, 1)
37
+ assert global_decode[0, :2].cpu().tolist() == [768, 769]
38
+ assert decode_lens.cpu().tolist() == [2]
39
+ assert prefill_local[0, :4].cpu().tolist() == list(range(4))
40
+ assert torch.all(global_decode[0, 2:] == -1)
41
+ assert torch.all(prefill_local[0, 4:] == -1)
42
+
43
+
44
def test_sparse_flashmla_metadata_smoke():
45
import vllm.v1.attention.ops.flashmla as fm
46
vllm/models/deepseek_v4/sparse_mla.py
+19
-7
@@ -261,6 +261,13 @@ class DeepseekV4FlashMLAMetadataBuilder(
261
assert cm.positions is not None, (
262
"positions is required for C128A metadata build"
263
)
264
+ active_topk_width = min(
265
+ max(
266
+ triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)),
267
+ _C128A_TOPK_ALIGNMENT,
268
+ ),
269
+ self.c128a_max_compressed,
270
+ )
271
block_size = self.kv_cache_spec.block_size // self.compress_ratio
272
global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
273
cm.positions[:num_total],
@@ -273,7 +280,7 @@ class DeepseekV4FlashMLAMetadataBuilder(
280
self.c128a_global_decode_buffer,
281
self.c128a_decode_lens_buffer,
282
self.c128a_prefill_buffer,
276
- max_compressed_tokens=self.c128a_max_compressed,
283
+ max_compressed_tokens=active_topk_width,
284
)
285
286
result: dict[str, torch.Tensor | None] = {}
@@ -305,25 +312,30 @@ def build_c128a_topk_metadata(
312
Decode tokens: position → block_table lookup → global slot ids + topk_lens.
313
Prefill tokens: position → local indices [0, ..., n-1, -1, ...].
314
308
- Writes into pre-allocated buffers for CUDA graph address stability.
309
- Returns slices of the buffers.
315
+ Writes into packed views of pre-allocated buffers for CUDA graph stability.
316
"""
317
num_tokens = positions.shape[0]
318
num_prefill_tokens = num_tokens - num_decode_tokens
319
314
- global_decode = global_decode_buffer[:num_decode_tokens]
320
+ # view(-1) as 1-d array and then expanded to
321
+ # [num_decode_tokens, max_compressed_tokens]
322
+ global_decode = global_decode_buffer.view(-1)[
323
+ : num_decode_tokens * max_compressed_tokens
324
+ ].view(num_decode_tokens, max_compressed_tokens)
325
decode_lens = decode_lens_buffer[:num_decode_tokens]
316
- prefill_local = prefill_buffer[:num_prefill_tokens]
326
+ prefill_local = prefill_buffer.view(-1)[
327
+ : num_prefill_tokens * max_compressed_tokens
328
+ ].view(num_prefill_tokens, max_compressed_tokens)
329
330
if num_tokens == 0:
331
return global_decode, decode_lens, prefill_local
332
333
_build_c128a_topk_metadata_kernel[(num_tokens,)](
334
global_decode_buffer,
323
- global_decode_buffer.stride(0),
335
+ max_compressed_tokens,
336
decode_lens_buffer,
337
prefill_buffer,
326
- prefill_buffer.stride(0),
338
+ max_compressed_tokens,
339
positions,
340
compress_ratio,
341
max_compressed_tokens,