A high-throughput and memory-efficient inference and serving engine for LLMs
Revert "[DSV4 Perf] Optimize dsv4 cudagraph by reducing `eager_break_during_capture`" (#45309) (#45972)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Woosuk Kwon committed
Jun 17, 2026 at 17:20 UTC
1797576237cc05d9c5b6789f35f09d13253db17a
1 file changed
+30
-37
vllm/models/deepseek_v4/attention.py
+30
-37
@@ -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 BreakableCUDAGraphCapture
17
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
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. For C4A layers, the inner sparse_attn_indexer custom op
335
- # runs in the eager break.
334
+ # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer,
335
+ # compressor, MLA attention) 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,6 +345,9 @@ 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).
351
self.attention_impl(
352
hidden_states,
353
qr,
@@ -420,6 +423,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
423
424
return qr_kv, kv_score, indexer_kv_score, indexer_weights
425
426
+ @eager_break_during_capture
427
def attention_impl(
428
self,
429
hidden_states: torch.Tensor,
@@ -447,42 +451,31 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
451
452
def wq_b_kv_insert() -> torch.Tensor:
453
q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
450
- return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
454
+ q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
455
+ return q
456
452
- run_indexer = lambda: indexer(
453
- hidden_states,
454
- qr,
455
- indexer_kv_score,
456
- indexer_weights,
457
- positions,
458
- self.indexer_rotary_emb,
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,
478
)
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
- )
479
elif self.compressor is not None:
480
# wq_b + kv_insert on default, compressor on aux.
481
aux_stream = (