[DSv4 Perf] Remove redundant full kernel for dsv4, 1.88x kernel performance improvement (#50298)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 30, 2026 at 11:39 UTC
837eae64580c885101ee95b073aafb27a485e7ce
3 files changed
+44
-23
tests/kernels/attention/test_flashmla_sparse.py
+16
-12
@@ -135,7 +135,8 @@ def test_sparse_flashmla_decode_smoke():
135
assert lse.shape[0] == batch_size
136
137
138
-def test_sparse_flashmla_prefill_smoke():
138
+@pytest.mark.parametrize("h_q", [64, 128])
139
+def test_sparse_flashmla_prefill_smoke(h_q: int):
140
import vllm.v1.attention.ops.flashmla as fm
141
142
ok, reason = fm.is_flashmla_sparse_supported()
@@ -143,22 +144,25 @@ def test_sparse_flashmla_prefill_smoke():
144
pytest.skip(reason)
145
146
device = torch.device("cuda")
147
+ torch.manual_seed(0)
148
s_q = 1
147
- s_kv = 1
148
- h_q = 64 # kernel expects multiple of 64
149
+ s_kv = 8
150
h_kv = 1
151
d_qk = 576
152
d_v = 512
153
topk = 128
153
-
154
- q = torch.zeros((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device)
155
- kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device)
156
- indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device)
157
-
158
- out, max_logits, lse = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v)
159
- assert out.shape == (s_q, h_q, d_v)
160
- assert max_logits.shape == (s_q, h_q)
161
- assert lse.shape == (s_q, h_q)
154
+ q = torch.randn((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device)
155
+ kv = torch.randn((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device)
156
+ indices = torch.randint(s_kv, (s_q, h_kv, topk), dtype=torch.int32, device=device)
157
+ reference_indices = indices.clone()
158
+ reference_indices[..., 1:] = -1
159
+ kwargs = {"topk_length": torch.ones(1, dtype=torch.int32, device=device)}
160
+ reference = fm.flash_mla_sparse_fwd(q, kv, reference_indices, 1.0, d_v, **kwargs)
161
+ actual = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v, **kwargs)
162
+
163
+ for actual_tensor, reference_tensor in zip(actual, reference):
164
+ torch.testing.assert_close(actual_tensor, reference_tensor, rtol=0, atol=0)
165
+ assert actual[0].shape == (s_q, h_q, d_v)
166
167
168
def test_deepseek_v4_prefill_chunk_planning_expands_for_short_sequences():
vllm/models/deepseek_v4/common/ops/cache_utils.py
+13
-9
@@ -535,6 +535,7 @@ def combine_topk_swa_indices(
535
topk: int,
536
M: int,
537
N: int,
538
+ out: tuple[torch.Tensor, torch.Tensor] | None = None,
539
) -> tuple[torch.Tensor, torch.Tensor]:
540
num_tokens = topk_indices.shape[0]
541
combined_topk = (
@@ -542,15 +543,18 @@ def combine_topk_swa_indices(
543
// _SPARSE_PREFILL_TOPK_ALIGNMENT
544
* _SPARSE_PREFILL_TOPK_ALIGNMENT
545
)
545
- combined_indices = torch.full(
546
- (num_tokens, combined_topk),
547
- fill_value=-1,
548
- dtype=torch.int32,
549
- device=topk_indices.device,
550
- )
551
- combined_lens = torch.empty(
552
- num_tokens, dtype=torch.int32, device=topk_indices.device
553
- )
546
+ if out is None:
547
+ combined_indices = torch.full(
548
+ (num_tokens, combined_topk),
549
+ fill_value=-1,
550
+ dtype=torch.int32,
551
+ device=topk_indices.device,
552
+ )
553
+ combined_lens = torch.empty(
554
+ num_tokens, dtype=torch.int32, device=topk_indices.device
555
+ )
556
+ else:
557
+ combined_indices, combined_lens = out
558
559
_COMBINE_TOPK_SWA_INDICES_KERNEL(
560
combined_indices,
vllm/models/deepseek_v4/nvidia/flashmla.py
+15
-2
@@ -20,6 +20,7 @@ from vllm.models.deepseek_v4.sparse_mla import (
20
DeepseekV4FlashMLABackend,
21
DeepseekV4FlashMLAMetadata,
22
)
23
+from vllm.utils.math_utils import round_up
24
from vllm.v1.attention.ops.flashmla import (
25
flash_mla_sparse_fwd,
26
flash_mla_with_kvcache,
@@ -95,8 +96,13 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
96
// self.compress_ratio
97
)
98
M = N + self.window_size + self.max_num_batched_tokens
99
+ assert self.topk_indices_buffer is not None
100
+ top_k = 0 if swa_only else self.topk_indices_buffer.shape[-1]
101
+ combined_topk = round_up(top_k + self.window_size, 128)
102
current_workspace_manager().get_simultaneous(
103
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
104
+ ((self.max_num_batched_tokens, combined_topk), torch.int32),
105
+ ((self.max_num_batched_tokens,), torch.int32),
106
)
107
output.zero_()
108
return
@@ -284,11 +290,15 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
290
)
291
assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
292
workspace_manager = current_workspace_manager()
293
+ combined_topk = round_up(top_k + self.window_size, 128)
294
for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
295
chunk_size = chunk_end - chunk_start
289
- kv = workspace_manager.get_simultaneous(
296
+ workspace = workspace_manager.get_simultaneous(
297
((chunk_size, chunk_M, q.shape[-1]), torch.bfloat16),
291
- )[0]
298
+ ((self.max_num_batched_tokens, combined_topk), torch.int32),
299
+ ((self.max_num_batched_tokens,), torch.int32),
300
+ )
301
+ kv, combined_indices_out, combined_lens_out = workspace
302
if not swa_only:
303
# Gather compressed KV
304
assert attn_metadata is not None
@@ -322,6 +332,8 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
332
query_end = (
333
query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
334
)
335
+ combined_indices_out = combined_indices_out[: query_end - query_start]
336
+ combined_lens_out = combined_lens_out[: query_end - query_start]
337
338
combined_indices, combined_lens = combine_topk_swa_indices(
339
topk_indices[query_start:query_end],
@@ -335,6 +347,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
347
top_k,
348
chunk_M,
349
chunk_N,
350
+ out=(combined_indices_out, combined_lens_out),
351
)
352
flash_mla_sparse_fwd(
353
q=q[query_start:query_end],