[Perf] Narrow DeepSeek V4 eager CUDA graph region (#51430)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Woosuk Kwon committed
Aug 10, 2026 at 10:52 UTC
79c865b838e34f7a98a936771284773819d79c8f
3 files changed
+90
-97
vllm/models/deepseek_v4/amd/model.py
+1
-1
@@ -551,7 +551,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
551
self.rms_norm_eps = config.rms_norm_eps
552
553
# Three aux streams: one per non-default input GEMM in
554
- # DeepseekV4Attention.attn_gemm_parallel_execute
554
+ # DeepseekV4Attention._run_parallel_input_projections
555
# (compressor kv_score, indexer.weights_proj, indexer.compressor
556
# kv_score). fused_wqa_wkv stays on the default stream.
557
# Disable them on ROCm because of hang issues.
vllm/models/deepseek_v4/attention.py
+88
-95
@@ -362,11 +362,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
362
device=hidden_states.device,
363
)
364
365
- # Metadata-independent input GEMMs + RMSNorm stay in the captured
366
- # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer,
367
- # compressor, MLA attention) runs in the eager break.
365
+ # Keep the attention input preparation in the captured graph. Only the
366
+ # sparse indexer and MLA attention run in the eager break below.
367
qr_kv, kv_score, indexer_kv_score, indexer_weights = (
369
- self.attn_gemm_parallel_execute(hidden_states)
368
+ self._run_parallel_input_projections(hidden_states)
369
)
370
qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
371
qr, kv = fused_q_kv_rmsnorm(
@@ -377,16 +376,62 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
376
self.eps,
377
)
378
380
- # attention_impl is wrapped with @eager_break_during_capture: this is
381
- # where the breakable cudagraph capture breaks (the attention op runs
382
- # eagerly between captured graph segments).
383
- self.attention_impl(
379
+ attn_metadata = get_forward_context().attn_metadata
380
+ indexer = self.indexer
381
+ compressor = self.compressor
382
+ aux_streams = self.aux_stream_list
383
+
384
+ def project_query_and_cache_kv() -> torch.Tensor:
385
+ q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
386
+ return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
387
+
388
+ index_q: torch.Tensor | None = None
389
+ index_q_scale: torch.Tensor | None = None
390
+ index_weights_out: torch.Tensor | None = None
391
+
392
+ # Keep Q projection and KV insertion on the default stream. The indexer
393
+ # and MLA compressor use aux streams 0 and 1; aux 2 is internal to the
394
+ # indexer. ROCm runs the same work sequentially without aux streams.
395
+ if indexer is not None:
396
+ assert compressor is not None
397
+ q, (indexer_inputs, _) = execute_in_parallel(
398
+ project_query_and_cache_kv,
399
+ [
400
+ lambda: indexer(
401
+ hidden_states,
402
+ qr,
403
+ indexer_kv_score,
404
+ indexer_weights,
405
+ positions,
406
+ self.indexer_rotary_emb,
407
+ ),
408
+ lambda: compressor(kv_score, positions, self.rotary_emb),
409
+ ],
410
+ self.ln_events[0],
411
+ [self.ln_events[1], self.ln_events[2]],
412
+ [aux_streams[0], aux_streams[1]] if aux_streams is not None else None,
413
+ enable=aux_streams is not None,
414
+ )
415
+ index_q, index_q_scale, index_weights_out = indexer_inputs
416
+ elif compressor is not None:
417
+ aux_stream = aux_streams[0] if aux_streams is not None else None
418
+ q, _ = maybe_execute_in_parallel(
419
+ project_query_and_cache_kv,
420
+ lambda: compressor(kv_score, positions, self.rotary_emb),
421
+ self.ln_events[0],
422
+ self.ln_events[1],
423
+ aux_stream,
424
+ )
425
+ else:
426
+ q = project_query_and_cache_kv()
427
+
428
+ self._sparse_indexer_and_attn(
429
hidden_states,
385
- qr,
430
+ index_q,
431
+ index_q_scale,
432
+ index_weights_out,
433
+ q,
434
kv,
387
- kv_score,
388
- indexer_kv_score,
389
- indexer_weights,
435
positions,
436
o_padded,
437
)
@@ -395,12 +440,14 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
440
# Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
441
return self._o_proj(o, positions)
442
398
- def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
399
- # MergedColumnParallelLinear returns (output, bias); bias is None.
400
- qr_kv, _ = self.fused_wqa_wkv(hidden_states)
401
- return qr_kv
402
-
403
- def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
443
+ def _run_parallel_input_projections(
444
+ self, hidden_states: torch.Tensor
445
+ ) -> tuple[
446
+ torch.Tensor,
447
+ torch.Tensor | None,
448
+ torch.Tensor | None,
449
+ torch.Tensor | None,
450
+ ]:
451
aux_streams = self.aux_stream_list
452
if aux_streams is not None:
453
assert len(aux_streams) >= 3
@@ -443,11 +490,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
490
aux_fns[1] = indexer_weights_proj
491
aux_fns[2] = indexer_compressor_kv_score
492
446
- def fused_wqa_wkv() -> torch.Tensor:
447
- return self._fused_wqa_wkv_gemm(hidden_states)
448
-
493
qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel(
450
- fused_wqa_wkv,
494
+ lambda: self.fused_wqa_wkv(hidden_states)[0],
495
aux_fns,
496
self.ln_events[0],
497
self.ln_events[1:4],
@@ -459,81 +503,26 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
503
return qr_kv, kv_score, indexer_kv_score, indexer_weights
504
505
@eager_break_during_capture
462
- def attention_impl(
506
+ def _sparse_indexer_and_attn(
507
self,
508
hidden_states: torch.Tensor,
465
- qr: torch.Tensor,
509
+ index_q: torch.Tensor | None,
510
+ index_q_scale: torch.Tensor | None,
511
+ index_weights: torch.Tensor | None,
512
+ q: torch.Tensor,
513
kv: torch.Tensor,
467
- kv_score: torch.Tensor,
468
- indexer_kv_score: torch.Tensor,
469
- indexer_weights: torch.Tensor,
514
positions: torch.Tensor,
471
- out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place
515
+ out: torch.Tensor,
516
) -> None:
473
- forward_context = get_forward_context()
474
- attn_metadata = forward_context.attn_metadata
475
-
476
- # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride
477
- # on the default stream so q stays on its consumer stream (forward_mqa
478
- # downstream reads q on default). Indexer/compressor go on aux for
479
- # overlap with default's GEMM + cache write.
480
- if self.indexer is not None:
481
- aux_streams = self.aux_stream_list
482
- indexer = self.indexer
483
- # Local ref so the closure keeps a non-None type for mypy.
484
- assert self.compressor is not None
485
- compressor = self.compressor
486
-
487
- def wq_b_kv_insert() -> torch.Tensor:
488
- q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
489
- q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
490
- return q
491
-
492
- # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
493
- # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the
494
- # MLA compressor. Slot [2] is reserved for the indexer's inner
495
- # overlap. ROCm (aux_streams is None) falls back to sequential.
496
- q, _ = execute_in_parallel(
497
- wq_b_kv_insert,
498
- [
499
- lambda: indexer(
500
- hidden_states,
501
- qr,
502
- indexer_kv_score,
503
- indexer_weights,
504
- positions,
505
- self.indexer_rotary_emb,
506
- ),
507
- lambda: compressor(kv_score, positions, self.rotary_emb),
508
- ],
509
- self.ln_events[0],
510
- [self.ln_events[1], self.ln_events[2]],
511
- [aux_streams[0], aux_streams[1]] if aux_streams is not None else None,
512
- enable=aux_streams is not None,
517
+ if self.indexer is not None and index_q is not None:
518
+ assert index_weights is not None
519
+ q_quant = (index_q, index_q_scale) if index_q_scale is not None else index_q
520
+ self.indexer.indexer_op(
521
+ hidden_states,
522
+ q_quant,
523
+ None,
524
+ index_weights,
525
)
514
- elif self.compressor is not None:
515
- # wq_b + kv_insert on default, compressor on aux.
516
- aux_stream = (
517
- self.aux_stream_list[0] if self.aux_stream_list is not None else None
518
- )
519
- compressor = self.compressor
520
-
521
- def wq_b_kv_insert() -> torch.Tensor:
522
- q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
523
- q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
524
- return q
525
-
526
- q, _ = maybe_execute_in_parallel(
527
- wq_b_kv_insert,
528
- lambda: compressor(kv_score, positions, self.rotary_emb),
529
- self.ln_events[0],
530
- self.ln_events[1],
531
- aux_stream,
532
- )
533
- else:
534
- # SWA-only layer: no compressor, no overlap.
535
- q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
536
- q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
526
527
# MLA attention writes into the pre-allocated `out` buffer
528
# ([num_tokens, padded_heads, head_dim]).
@@ -836,7 +825,7 @@ class DeepseekV4Indexer(nn.Module):
825
indexer_weights: torch.Tensor,
826
positions: torch.Tensor,
827
rotary_emb: nn.Module,
839
- ) -> torch.Tensor:
828
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
829
compressor = self.compressor
830
831
attn_metadata = get_forward_context().attn_metadata
@@ -860,7 +849,7 @@ class DeepseekV4Indexer(nn.Module):
849
PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens),
850
num_warps=8,
851
)
863
- return self.topk_indices_buffer
852
+ return None, None, None
853
854
def wq_b_and_q_quant():
855
# ReplicatedLinear returns (output, bias); bias is None.
@@ -882,11 +871,15 @@ class DeepseekV4Indexer(nn.Module):
871
872
# compressor returns None and writes K to the indexer KV cache; the
873
# join orders that write before indexer_op (skip_k_cache_insert=True).
885
- (q_quant, weights), k = maybe_execute_in_parallel(
874
+ (q_quant, weights), _ = maybe_execute_in_parallel(
875
wq_b_and_q_quant,
876
lambda: compressor(compressed_kv_score, positions, rotary_emb),
877
self.ln_events[0],
878
self.ln_events[1],
879
self.aux_stream,
880
)
892
- return self.indexer_op(hidden_states, q_quant, k, weights)
881
+ if isinstance(q_quant, tuple):
882
+ q, q_scale = q_quant
883
+ else:
884
+ q, q_scale = q_quant, None
885
+ return q, q_scale, weights
vllm/models/deepseek_v4/nvidia/model.py
+1
-1
@@ -1018,7 +1018,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1018
self.rms_norm_eps = config.rms_norm_eps
1019
1020
# Three aux streams: one per non-default input GEMM in
1021
- # DeepseekV4Attention.attn_gemm_parallel_execute
1021
+ # DeepseekV4Attention._run_parallel_input_projections
1022
# (compressor kv_score, indexer.weights_proj, indexer.compressor
1023
# kv_score). fused_wqa_wkv stays on the default stream.
1024
aux_stream_list = [torch.cuda.Stream() for _ in range(3)]