iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[DSV4 Perf] Optimize dsv4 cudagraph by reducing eager_break_during_capture, 26.8% ~ 27.9% E2E TTFT improvement (#45309)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-06-17 09:34:53 -07:00
committed by GitHub
parent 9c7c74bf10
commit 2a47a9ff0f
1 changed files with 37 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/deepseek_v4/attention.py
+37 -30
View File
@@ -14,7 +14,7 @@ import torch.nn.functional as F
from transformers import DeepseekV2Config, DeepseekV3Config
import vllm.envs as envs
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
from vllm.model_executor.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
@@ -331,8 +331,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 C4A layers, the inner sparse_attn_indexer custom op
# runs in the eager break.
qr_kv, kv_score, indexer_kv_score, indexer_weights = (
self.attn_gemm_parallel_execute(hidden_states)
)
@@ -345,9 +345,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
self.eps,
)
# attention_impl is wrapped with @eager_break_during_capture: this is
# where the breakable cudagraph capture breaks (the attention op runs
# eagerly between captured graph segments).
self.attention_impl(
hidden_states,
qr,
@@ -423,7 +420,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
return qr_kv, kv_score, indexer_kv_score, indexer_weights
@eager_break_during_capture
def attention_impl(
self,
hidden_states: torch.Tensor,
@@ -451,31 +447,42 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
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
return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
# 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
# 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(
wq_b_kv_insert,
[
lambda: indexer(
hidden_states,
qr,
indexer_kv_score,
indexer_weights,
positions,
self.indexer_rotary_emb,
),
lambda: compressor(kv_score, positions, self.rotary_emb),
],
self.ln_events[0],
[self.ln_events[1], self.ln_events[2]],
[aux_streams[0], aux_streams[1]] if aux_streams is not None else None,
enable=aux_streams is not None,
run_indexer = lambda: indexer(
hidden_states,
qr,
indexer_kv_score,
indexer_weights,
positions,
self.indexer_rotary_emb,
)
run_compressor = lambda: compressor(kv_score, positions, self.rotary_emb)
if BreakableCUDAGraphCapture.is_active():
q, _ = maybe_execute_in_parallel(
wq_b_kv_insert,
run_compressor,
self.ln_events[0],
self.ln_events[1],
aux_streams[1] if aux_streams is not None else None,
)
run_indexer()
else:
# 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
# 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(
wq_b_kv_insert,
[run_indexer, run_compressor],
self.ln_events[0],
[self.ln_events[1], self.ln_events[2]],
[aux_streams[0], aux_streams[1]]
if aux_streams is not None
else None,
enable=aux_streams is not None,
)
elif self.compressor is not None:
# wq_b + kv_insert on default, compressor on aux.
aux_stream = (