A high-throughput and memory-efficient inference and serving engine for LLMs
[DSV4 Perf] Optimize dsv4 cudagraph by reducing `eager_break_during_capture`, 26.8% ~ 27.9% E2E TTFT improvement (#45309)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 17, 2026 at 12:34 UTC
2a47a9ff0f4f302ee64915dfea64642757e31ee4
1 file changed
+37
-30
vllm/models/deepseek_v4/attention.py
+37
-30
@@ -14,7 +14,7 @@ import torch.nn.functional as F
14
from transformers import DeepseekV2Config, DeepseekV3Config
15
16
import vllm.envs as envs
17
-from vllm.compilation.breakable_cudagraph import eager_break_during_capture
17
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
20
MergedColumnParallelLinear,
@@ -331,8 +331,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
331
)
332
333
# Metadata-independent input GEMMs + RMSNorm stay in the captured
334
- # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer,
335
- # compressor, MLA attention) runs in the eager break.
334
+ # graph. For C4A layers, the inner sparse_attn_indexer custom op
335
+ # runs in the eager break.
336
qr_kv, kv_score, indexer_kv_score, indexer_weights = (
337
self.attn_gemm_parallel_execute(hidden_states)
338
)
@@ -345,9 +345,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
345
self.eps,
346
)
347
348
- # attention_impl is wrapped with @eager_break_during_capture: this is
349
- # where the breakable cudagraph capture breaks (the attention op runs
350
- # eagerly between captured graph segments).
348
self.attention_impl(
349
hidden_states,
350
qr,
@@ -423,7 +420,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
420
421
return qr_kv, kv_score, indexer_kv_score, indexer_weights
422
426
- @eager_break_during_capture
423
def attention_impl(
424
self,
425
hidden_states: torch.Tensor,
@@ -451,31 +447,42 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
447
448
def wq_b_kv_insert() -> torch.Tensor:
449
q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
454
- q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
455
- return q
450
+ return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
451
457
- # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
458
- # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the
459
- # MLA compressor. Slot [2] is reserved for the indexer's inner
460
- # overlap. ROCm (aux_streams is None) falls back to sequential.
461
- q, _ = execute_in_parallel(
462
- wq_b_kv_insert,
463
- [
464
- lambda: indexer(
465
- hidden_states,
466
- qr,
467
- indexer_kv_score,
468
- indexer_weights,
469
- positions,
470
- self.indexer_rotary_emb,
471
- ),
472
- lambda: compressor(kv_score, positions, self.rotary_emb),
473
- ],
474
- self.ln_events[0],
475
- [self.ln_events[1], self.ln_events[2]],
476
- [aux_streams[0], aux_streams[1]] if aux_streams is not None else None,
477
- enable=aux_streams is not None,
452
+ run_indexer = lambda: indexer(
453
+ hidden_states,
454
+ qr,
455
+ indexer_kv_score,
456
+ indexer_weights,
457
+ positions,
458
+ self.indexer_rotary_emb,
459
)
460
+ run_compressor = lambda: compressor(kv_score, positions, self.rotary_emb)
461
+
462
+ if BreakableCUDAGraphCapture.is_active():
463
+ q, _ = maybe_execute_in_parallel(
464
+ wq_b_kv_insert,
465
+ run_compressor,
466
+ self.ln_events[0],
467
+ self.ln_events[1],
468
+ aux_streams[1] if aux_streams is not None else None,
469
+ )
470
+ run_indexer()
471
+ else:
472
+ # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
473
+ # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the
474
+ # MLA compressor. Slot [2] is reserved for the indexer's inner
475
+ # overlap. ROCm (aux_streams is None) falls back to sequential.
476
+ q, _ = execute_in_parallel(
477
+ wq_b_kv_insert,
478
+ [run_indexer, run_compressor],
479
+ self.ln_events[0],
480
+ [self.ln_events[1], self.ln_events[2]],
481
+ [aux_streams[0], aux_streams[1]]
482
+ if aux_streams is not None
483
+ else None,
484
+ enable=aux_streams is not None,
485
+ )
486
elif self.compressor is not None:
487
# wq_b + kv_insert on default, compressor on aux.
488
aux_stream = (