[DSv32/GLM Perf] Skip short prefill topk for dense mha layer, 97.9% kernel level latency reduction (#51298)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 7, 2026 at 16:30 UTC
4a11ed5bc4c2f4e818e1e33010da160449d59159
1 file changed
+11
-1
vllm/models/deepseek_v32/attention.py
+11
-1
@@ -32,6 +32,7 @@ from vllm.model_executor.models.deepseek_v2 import (
32
)
33
from vllm.model_executor.models.utils import extract_layer_index
34
from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q
35
+from vllm.platforms import current_platform
36
from vllm.utils.torch_utils import is_quantized_kv_cache
37
38
@@ -259,6 +260,15 @@ class DeepseekV32Attention(MLAAttention):
260
self.topk_indices_buffer = topk_indices_buffer
261
262
self.skip_topk = False
263
+ enable_short_prefill_scoring_skip = (
264
+ not is_mtp_layer
265
+ and not skip_topk
266
+ and not self.use_pcp
267
+ and current_platform.is_cuda()
268
+ )
269
+ self._dense_mha_metadata_layer_name = (
270
+ self.layer_name if enable_short_prefill_scoring_skip else ""
271
+ )
272
273
if self.require_fp8_kv_cache:
274
assert is_quantized_kv_cache(self.kv_cache_dtype), (
@@ -470,7 +480,7 @@ class DeepseekV32Attention(MLAAttention):
480
self.topk_indices_buffer,
481
skip_k_cache_insert=True,
482
use_pcp=False,
473
- dense_mha_metadata_layer_name="",
483
+ dense_mha_metadata_layer_name=self._dense_mha_metadata_layer_name,
484
use_fp4_cache=False,
485
# fused_norm_rope already cleared the topk buffer this forward.
486
skip_topk_buffer_clear=True,