[DSv4 Perf] Adaptive topk width for dsv4, making #50004 back (#52823)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 21, 2026 at 15:11 UTC
e6f35d3c69b23b2ed35afebde31a0567c40bc661
2 files changed
+78
-3
tests/kernels/attention/test_flashmla_sparse.py
+57
@@ -4,6 +4,63 @@ import pytest
4
import torch
5
6
7
+def test_deepseek_v4_c128a_adaptive_width_has_capture_stable_stride():
8
+ from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata
9
+
10
+ device = torch.device("cuda")
11
+ capacity_width = 512
12
+ global_decode_buffer = torch.empty(
13
+ (2, capacity_width), dtype=torch.int32, device=device
14
+ )
15
+ prefill_buffer = torch.empty_like(global_decode_buffer)
16
+ kwargs = dict(
17
+ positions=torch.tensor([255, 511, 383, 639], device=device),
18
+ compress_ratio=128,
19
+ num_decode_tokens=2,
20
+ token_to_req_indices=torch.tensor(
21
+ [0, 1, 0, 1], dtype=torch.int32, device=device
22
+ ),
23
+ block_table=torch.tensor([[3], [5]], dtype=torch.int32, device=device),
24
+ block_size=capacity_width,
25
+ slot_mapping=torch.arange(4, dtype=torch.int64, device=device),
26
+ global_decode_buffer=global_decode_buffer,
27
+ decode_lens_buffer=torch.empty(2, dtype=torch.int32, device=device),
28
+ prefill_buffer=prefill_buffer,
29
+ )
30
+ captured_decode, _, captured_prefill = build_c128a_topk_metadata(
31
+ max_compressed_tokens=256,
32
+ **kwargs,
33
+ )
34
+ assert captured_decode.shape == captured_prefill.shape == (2, 256)
35
+ assert captured_decode.stride(0) == captured_prefill.stride(0) == capacity_width
36
+
37
+ captured_rows = torch.empty((4, 4), dtype=torch.int32, device=device)
38
+ captured_rows[:2].copy_(captured_decode[:, :4])
39
+ captured_rows[2:].copy_(captured_prefill[:, :4])
40
+ torch.accelerator.synchronize()
41
+ graph = torch.cuda.CUDAGraph()
42
+ with torch.cuda.graph(graph):
43
+ captured_rows[:2].copy_(captured_decode[:, :4])
44
+ captured_rows[2:].copy_(captured_prefill[:, :4])
45
+
46
+ global_decode_buffer.fill_(-99)
47
+ prefill_buffer.fill_(-99)
48
+ build_c128a_topk_metadata(
49
+ max_compressed_tokens=128,
50
+ **kwargs,
51
+ )
52
+ graph.replay()
53
+
54
+ assert captured_rows.cpu().tolist() == [
55
+ [1536, 1537, -1, -1],
56
+ [2560, 2561, 2562, 2563],
57
+ [0, 1, 2, -1],
58
+ [0, 1, 2, 3],
59
+ ]
60
+ assert torch.all(global_decode_buffer[:, 128:] == -99)
61
+ assert torch.all(prefill_buffer[:, 128:] == -99)
62
+
63
+
64
def test_sparse_flashmla_metadata_smoke():
65
import vllm.v1.attention.ops.flashmla as fm
66
vllm/models/deepseek_v4/sparse_mla.py
+21
-3
@@ -257,6 +257,15 @@ class DeepseekV4SparseMLAMetadataBuilder(
257
assert cm.positions is not None, (
258
"positions is required for C128A metadata build"
259
)
260
+ active_topk_width = min(
261
+ max(
262
+ triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)),
263
+ _C128A_TOPK_ALIGNMENT,
264
+ ),
265
+ self.c128a_max_compressed,
266
+ )
267
+ assert active_topk_width >= cm.max_seq_len // self.compress_ratio
268
+ assert active_topk_width % _C128A_TOPK_ALIGNMENT == 0
269
block_size = self.kv_cache_spec.block_size // self.compress_ratio
270
global_decode, decode_lens, prefill_local = build_c128a_topk_metadata(
271
cm.positions[:num_total],
@@ -269,7 +278,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
278
self.c128a_global_decode_buffer,
279
self.c128a_decode_lens_buffer,
280
self.c128a_prefill_buffer,
272
- max_compressed_tokens=self.c128a_max_compressed,
281
+ max_compressed_tokens=active_topk_width,
282
)
283
284
result: dict[str, torch.Tensor | None] = {}
@@ -320,10 +329,19 @@ def build_c128a_topk_metadata(
329
"""
330
num_tokens = positions.shape[0]
331
num_prefill_tokens = num_tokens - num_decode_tokens
332
+ assert max_compressed_tokens % _C128A_TOPK_ALIGNMENT == 0
333
+ assert (
334
+ 0
335
+ < max_compressed_tokens
336
+ <= min(global_decode_buffer.shape[1], prefill_buffer.shape[1])
337
+ )
338
+ assert global_decode_buffer.stride(-1) == prefill_buffer.stride(-1) == 1
339
324
- global_decode = global_decode_buffer[:num_decode_tokens]
340
+ global_decode = global_decode_buffer[:num_decode_tokens, :max_compressed_tokens]
341
decode_lens = decode_lens_buffer[:num_decode_tokens]
326
- prefill_local = prefill_buffer[:num_prefill_tokens]
342
+ prefill_local = prefill_buffer[:num_prefill_tokens, :max_compressed_tokens]
343
+ assert global_decode.stride(0) == global_decode_buffer.stride(0)
344
+ assert prefill_local.stride(0) == prefill_buffer.stride(0)
345
346
if num_tokens == 0:
347
return global_decode, decode_lens, prefill_local