From 4a11ed5bc4c2f4e818e1e33010da160449d59159 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 7 Aug 2026 16:30:42 -0400 Subject: [PATCH] [DSv32/GLM Perf] Skip short prefill topk for dense mha layer, 97.9% kernel level latency reduction (#51298) Signed-off-by: yewentao256 --- vllm/models/deepseek_v32/attention.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index 705db45f59..6ad6ccba5d 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -32,6 +32,7 @@ from vllm.model_executor.models.deepseek_v2 import ( ) from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q +from vllm.platforms import current_platform from vllm.utils.torch_utils import is_quantized_kv_cache @@ -259,6 +260,15 @@ class DeepseekV32Attention(MLAAttention): self.topk_indices_buffer = topk_indices_buffer self.skip_topk = False + enable_short_prefill_scoring_skip = ( + not is_mtp_layer + and not skip_topk + and not self.use_pcp + and current_platform.is_cuda() + ) + self._dense_mha_metadata_layer_name = ( + self.layer_name if enable_short_prefill_scoring_skip else "" + ) if self.require_fp8_kv_cache: assert is_quantized_kv_cache(self.kv_cache_dtype), ( @@ -470,7 +480,7 @@ class DeepseekV32Attention(MLAAttention): self.topk_indices_buffer, skip_k_cache_insert=True, use_pcp=False, - dense_mha_metadata_layer_name="", + dense_mha_metadata_layer_name=self._dense_mha_metadata_layer_name, use_fp4_cache=False, # fused_norm_rope already cleared the topk buffer this forward. skip_topk_buffer_clear=True,