[DSv4 Perf] Skip topk and router when not needed, 3.4% E2E TTFT improvement for Decode case (#49486)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 23, 2026 at 13:08 UTC
b0cb1da1bde62a738baba33f1fbb1fcf906d29bc
1 file changed
+43
vllm/models/deepseek_v4/attention.py
+43
@@ -47,6 +47,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
47
from vllm.model_executor.models.utils import extract_layer_index
48
from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope
49
from vllm.models.deepseek_v4.compressor import DeepseekCompressor
50
+from vllm.triton_utils import tl, triton
51
from vllm.utils.multi_stream_utils import (
52
execute_in_parallel,
53
maybe_execute_in_parallel,
@@ -66,6 +67,25 @@ from vllm.v1.kv_cache_interface import (
67
logger = init_logger(__name__)
68
69
70
+@triton.jit
71
+def _fill_short_context_topk_indices(
72
+ output,
73
+ positions,
74
+ TOP_K: tl.constexpr,
75
+ COMPRESS_RATIO: tl.constexpr,
76
+ PADDED_TOP_K: tl.constexpr,
77
+):
78
+ # small triton kernel that selects every candidate, -1 otherwise
79
+ row = tl.program_id(0)
80
+ offsets = tl.arange(0, PADDED_TOP_K)
81
+ num_compressed = (tl.load(positions + row) + 1) // COMPRESS_RATIO
82
+ tl.store(
83
+ output + row * TOP_K + offsets,
84
+ tl.where(offsets < num_compressed, offsets, -1),
85
+ mask=offsets < TOP_K,
86
+ )
87
+
88
+
89
def _resolve_dsv4_kv_cache_dtype(
90
use_fp8_ds_mla_layout: bool,
91
kv_cache_dtype: str,
@@ -787,6 +807,29 @@ class DeepseekV4Indexer(nn.Module):
807
) -> torch.Tensor:
808
compressor = self.compressor
809
810
+ attn_metadata = get_forward_context().attn_metadata
811
+ if isinstance(attn_metadata, dict):
812
+ indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
813
+ if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens:
814
+ # candidates num smaller than topk, every candidate is selected
815
+ # but we still need to build k cache
816
+ compressor(compressed_kv_score, positions, rotary_emb)
817
+ assert self.topk_indices_buffer is not None
818
+ num_tokens = (
819
+ indexer_metadata.num_decode_tokens
820
+ + indexer_metadata.num_prefill_tokens
821
+ )
822
+ if num_tokens > 0:
823
+ _fill_short_context_topk_indices[(num_tokens,)](
824
+ self.topk_indices_buffer,
825
+ positions,
826
+ TOP_K=self.topk_tokens,
827
+ COMPRESS_RATIO=self.compress_ratio,
828
+ PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens),
829
+ num_warps=8,
830
+ )
831
+ return self.topk_indices_buffer
832
+
833
def wq_b_and_q_quant():
834
# ReplicatedLinear returns (output, bias); bias is None.
835
q, _ = self.wq_b(qr)