update

Signed-off-by: yewentao256 <zhyanwentao@126.com>

vllm/models/deepseek_v4/attention.py

@@ -333,8 +333,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         )
 
         # Metadata-independent input GEMMs + RMSNorm stay in the captured
-        # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer,
-        # compressor, MLA attention) runs in the eager break.
+        # 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,6 +346,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             self.kv_norm.weight.data,
             self.eps,
         )
+        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
@@ -353,6 +356,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         self.attention_impl(
             hidden_states,
             qr,
+            q,
             kv,
             kv_score,
             indexer_kv_score,
@@ -430,6 +434,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         self,
         hidden_states: torch.Tensor,
         qr: torch.Tensor,
+        q: torch.Tensor | None,
         kv: torch.Tensor,
         kv_score: torch.Tensor,
         indexer_kv_score: torch.Tensor,
@@ -499,7 +504,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             )
         else:
             # SWA-only layer: no compressor, no overlap.
-            q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
+            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