karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 2 Packages Activity

update

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
yewentao256
2026-07-02 15:16:21 +00:00
parent 7c43124d85
commit ce89bc6358
1 changed files with 22 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/deepseek_v4/attention.py
+22 -15
View File
@@ -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