A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 21:15 UTC
[DSV4] Move more ops out of eager breakpoint (#44561)
Woosuk Kwon committed
Jun 5, 2026 at 06:42 UTC
02d2da0748a11de735df0a65f8e7dae5bb99ec96
1 file changed
+30
-14
vllm/models/deepseek_v4/attention.py
+30
-14
@@ -330,10 +330,34 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
330
device=hidden_states.device,
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.
336
+ qr_kv, kv_score, indexer_kv_score, indexer_weights = (
337
+ self.attn_gemm_parallel_execute(hidden_states)
338
+ )
339
+ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
340
+ qr, kv = fused_q_kv_rmsnorm(
341
+ qr,
342
+ kv,
343
+ self.q_norm.weight.data,
344
+ self.kv_norm.weight.data,
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).
336
- self.attention_impl(hidden_states, positions, o_padded)
351
+ self.attention_impl(
352
+ hidden_states,
353
+ qr,
354
+ kv,
355
+ kv_score,
356
+ indexer_kv_score,
357
+ indexer_weights,
358
+ positions,
359
+ o_padded,
360
+ )
361
o = o_padded[:, : self.n_local_heads, :]
362
363
# Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
@@ -403,25 +427,17 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
427
def attention_impl(
428
self,
429
hidden_states: torch.Tensor,
430
+ qr: torch.Tensor,
431
+ kv: torch.Tensor,
432
+ kv_score: torch.Tensor,
433
+ indexer_kv_score: torch.Tensor,
434
+ indexer_weights: torch.Tensor,
435
positions: torch.Tensor,
436
out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place
437
) -> None:
438
forward_context = get_forward_context()
439
attn_metadata = forward_context.attn_metadata
440
412
- qr_kv, kv_score, indexer_kv_score, indexer_weights = (
413
- self.attn_gemm_parallel_execute(hidden_states)
414
- )
415
-
416
- qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
417
- qr, kv = fused_q_kv_rmsnorm(
418
- qr,
419
- kv,
420
- self.q_norm.weight.data,
421
- self.kv_norm.weight.data,
422
- self.eps,
423
- )
424
-
441
# wq_b + kv_insert (+ MLA compressor when an indexer is present) ride
442
# on the default stream so q stays on its consumer stream (forward_mqa
443
# downstream reads q on default). Indexer/compressor go on aux for