From ce89bc6358ab58732652084f33f9ccdaedb9dd6b Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Thu, 2 Jul 2026 15:16:21 +0000 Subject: [PATCH] update Signed-off-by: yewentao256 --- vllm/models/deepseek_v4/attention.py | 37 +++++++++++++++++----------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 41e6a597470..58781a658dd 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -332,9 +332,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): device=hidden_states.device, ) - # Metadata-independent input GEMMs, RMSNorm, and q up-proj stay in the - # captured graph; the metadata-dependent rest (kv-insert, indexer, - # compressor, MLA attention) runs in the eager break. + # Metadata-independent input GEMMs + RMSNorm stay in the captured + # graph. For SWA-only layers, q up-proj also stays in the graph because + # there is no compressor/indexer overlap to preserve. qr_kv, kv_score, indexer_kv_score, indexer_weights = ( self.attn_gemm_parallel_execute(hidden_states) ) @@ -346,7 +346,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.kv_norm.weight.data, self.eps, ) - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = None + if self.compressor is None: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) # attention_impl is wrapped with @eager_break_during_capture: this is # where the breakable cudagraph capture breaks (the attention op runs @@ -432,7 +434,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self, hidden_states: torch.Tensor, qr: torch.Tensor, - q: torch.Tensor, + q: torch.Tensor | None, kv: torch.Tensor, kv_score: torch.Tensor, indexer_kv_score: torch.Tensor, @@ -443,10 +445,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata - # Q norm + kv_insert (+ MLA compressor when an indexer is present) ride + # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride # on the default stream so q stays on its consumer stream (forward_mqa # downstream reads q on default). Indexer/compressor go on aux for - # overlap with default's cache write. + # overlap with default's GEMM + cache write. if self.indexer is not None: aux_streams = self.aux_stream_list indexer = self.indexer @@ -454,15 +456,17 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): assert self.compressor is not None compressor = self.compressor - def qnorm_rope_kv_insert() -> torch.Tensor: - return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs - # qnorm+kv_insert; slot [0] runs the full indexer; slot [1] runs - # the MLA compressor. Slot [2] is reserved for the indexer's inner + # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the + # MLA compressor. Slot [2] is reserved for the indexer's inner # overlap. ROCm (aux_streams is None) falls back to sequential. q, _ = execute_in_parallel( - qnorm_rope_kv_insert, + wq_b_kv_insert, [ lambda: indexer( hidden_states, @@ -486,11 +490,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) compressor = self.compressor - def qnorm_rope_kv_insert() -> torch.Tensor: - return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q q, _ = maybe_execute_in_parallel( - qnorm_rope_kv_insert, + wq_b_kv_insert, lambda: compressor(kv_score, positions, self.rotary_emb), self.ln_events[0], self.ln_events[1], @@ -498,6 +504,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) else: # SWA-only layer: no compressor, no overlap. + assert q is not None q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) # MLA attention writes into the pre-allocated `out` buffer