[Perf] Narrow DeepSeek V3.2 eager CUDA graph region (#51425)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Aug 7, 2026 at 14:37 UTC
46b58640541274334e2ee502a9c744841fb310b3
2 files changed
+61
-16
vllm/models/deepseek_v32/amd/rocm.py
+29
@@ -66,6 +66,35 @@ class DeepseekV32MLAAttention(DeepseekV32Attention):
66
self._fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype)
67
self._fp8_kv_needs_view = self._fp8_kv and self.kv_cache_dtype != "fp8_ds_mla"
68
69
+ def forward( # type: ignore[override]
70
+ self,
71
+ positions: torch.Tensor,
72
+ hidden_states: torch.Tensor,
73
+ ) -> torch.Tensor:
74
+ qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
75
+ q_c, kv_c, k_pe = qkv_lora.split(
76
+ [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
77
+ )
78
+
79
+ if self.indexer is not None and not self.skip_topk:
80
+ kw = self.indexer.wk_weights_proj(hidden_states)[0]
81
+ index_k = kw[:, : self.indexer.head_dim]
82
+ index_weights = kw[:, self.indexer.head_dim :]
83
+ else:
84
+ index_k = None
85
+ index_weights = None
86
+
87
+ num_tokens = hidden_states.shape[0]
88
+ output = torch.empty(
89
+ (num_tokens, self.num_local_heads * self.v_head_dim),
90
+ dtype=hidden_states.dtype,
91
+ device=hidden_states.device,
92
+ )
93
+ self._fused_attention(
94
+ positions, q_c, kv_c, k_pe, index_k, index_weights, output
95
+ )
96
+ return self.o_proj(output)[0]
97
+
98
def _compute_ql_nope(self, q_c: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
99
q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)
100
q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
vllm/models/deepseek_v32/attention.py
+32
-16
@@ -353,22 +353,6 @@ class DeepseekV32Attention(MLAAttention):
353
dtype=hidden_states.dtype,
354
device=hidden_states.device,
355
)
356
- self._fused_attention(
357
- positions, q_c, kv_c, k_pe, index_k, index_weights, output
358
- )
359
- return self.o_proj(output)[0]
360
-
361
- @eager_break_during_capture
362
- def _fused_attention(
363
- self,
364
- positions: torch.Tensor,
365
- q_c: torch.Tensor,
366
- kv_c: torch.Tensor,
367
- k_pe: torch.Tensor,
368
- index_k: torch.Tensor | None,
369
- index_weights: torch.Tensor | None,
370
- output: torch.Tensor,
371
- ) -> None:
356
forward_context = get_forward_context()
357
attn_metadata_raw = forward_context.attn_metadata
358
if isinstance(attn_metadata_raw, dict):
@@ -462,6 +446,26 @@ class DeepseekV32Attention(MLAAttention):
446
quantize_mqa=self._fp8_query,
447
)
448
449
+ self._sparse_indexer_and_attn(
450
+ q_c,
451
+ index_q_fp8,
452
+ index_weights_out,
453
+ ql_nope,
454
+ mqa_q,
455
+ output,
456
+ )
457
+ return self.o_proj(output)[0]
458
+
459
+ @eager_break_during_capture
460
+ def _sparse_indexer_and_attn(
461
+ self,
462
+ q_c: torch.Tensor,
463
+ index_q_fp8: torch.Tensor | None,
464
+ index_weights_out: torch.Tensor | None,
465
+ ql_nope: torch.Tensor,
466
+ mqa_q: torch.Tensor,
467
+ output: torch.Tensor,
468
+ ) -> None:
469
if self.indexer is not None and not self.skip_topk:
470
sparse_attn_indexer(
471
q_c,
@@ -486,6 +490,14 @@ class DeepseekV32Attention(MLAAttention):
490
skip_topk_buffer_clear=True,
491
)
492
493
+ attn_metadata_raw = get_forward_context().attn_metadata
494
+ if isinstance(attn_metadata_raw, dict):
495
+ attn_metadata = attn_metadata_raw.get(self.layer_name)
496
+ elif isinstance(attn_metadata_raw, list):
497
+ attn_metadata = attn_metadata_raw[0].get(self.layer_name)
498
+ else:
499
+ attn_metadata = attn_metadata_raw
500
+
501
if attn_metadata is None:
502
output.zero_()
503
return
@@ -504,6 +516,10 @@ class DeepseekV32Attention(MLAAttention):
516
attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]
517
mqa_q_arg, kv_cache, attn_metadata, self
518
)
519
+
520
+ # NOTE(woosuk): While the below does not need to be in the eager region,
521
+ # we put it here to avoid copying the attention output. Move this back to the
522
+ # captured region once forward_mqa supports `out` argument.
523
x = attn_out.view(
524
num_actual, self.num_local_heads, self.kv_lora_rank
525
).transpose(0, 1)