[ROCm][Perf] Enable CSA multi-stream overlap for DeepSeek-V4 (#51794)
Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Shanshan Shen committed
Sep 14, 2026 at 19:21 UTC
a6c5d6d0fcd7dd5b49e3146991c201b16caebef2
4 files changed
+214
-25
vllm/models/deepseek_v4/amd/model.py
+1
-10
@@ -957,16 +957,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
957
self.hc_dim = self.hc_mult * config.hidden_size
958
self.rms_norm_eps = config.rms_norm_eps
959
960
- # Three aux streams: one per non-default input GEMM in
961
- # DeepseekV4Attention._run_parallel_input_projections
962
- # (compressor kv_score, indexer.weights_proj, indexer.compressor
963
- # kv_score). fused_wqa_wkv stays on the default stream.
964
- # Disable them on ROCm because of hang issues.
965
- aux_stream_list = (
966
- None
967
- if current_platform.is_rocm()
968
- else [torch.cuda.Stream() for _ in range(3)]
969
- )
960
+ aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
961
962
self.device = current_platform.device_type
963
# Reserved topk indices buffer for all Indexer layers to reuse.
vllm/models/deepseek_v4/amd/rocm.py
+194
@@ -23,6 +23,7 @@ from vllm.models.deepseek_v4.sparse_mla import (
23
from vllm.platforms import current_platform
24
from vllm.platforms.rocm import _ON_GFX950
25
from vllm.triton_utils import tl, triton
26
+from vllm.utils.multi_stream_utils import execute_in_parallel
27
from vllm.v1.attention.backend import (
28
CommonAttentionMetadata,
29
)
@@ -571,6 +572,199 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
572
self.register_buffer("_fused_compressor_weight", None, persistent=False)
573
self._fused_compressor_split_sizes: tuple[int, int] | None = None
574
575
+ if self.indexer is None:
576
+ # Only enable multi-stream overlap for CSA layer now.
577
+ self.aux_stream_list = None
578
+ else:
579
+ # Disable indexer inner overlap.
580
+ self.indexer.aux_stream = None
581
+
582
+ def _enable_csa_multi_stream(self) -> bool:
583
+ """All CSA multi-stream gates: env var, streams, and capture region.
584
+
585
+ Dict metadata marks piecewise cudagraph, whose eager breaks rebuild
586
+ the attention inputs on the owning stream. Forking side streams
587
+ there would rely on runtime HIP event sync, which is unreliable in
588
+ this overlap on ROCm (event waits can hang), so multi-stream only
589
+ runs where the fork/join becomes static graph edges: inside capture,
590
+ or with non-dict metadata (full cudagraph or the profile run), which
591
+ has no eager breaks.
592
+ """
593
+ attn_metadata = get_forward_context().attn_metadata
594
+ return self.aux_stream_list is not None and (
595
+ torch.cuda.is_current_stream_capturing()
596
+ or not isinstance(attn_metadata, dict)
597
+ )
598
+
599
+ def _run_sequential_pipeline(
600
+ self,
601
+ hidden_states: torch.Tensor,
602
+ positions: torch.Tensor,
603
+ o_padded: torch.Tensor,
604
+ ) -> None:
605
+ """Disable ROCm streams when the current execution region cannot overlap."""
606
+ aux_streams = self.aux_stream_list
607
+ self.aux_stream_list = None
608
+ try:
609
+ qr_kv, kv_score, indexer_kv_score, indexer_weights = (
610
+ self._run_parallel_input_projections(hidden_states)
611
+ )
612
+ qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
613
+ self._prepare_and_attn_fn(
614
+ hidden_states,
615
+ qr,
616
+ kv,
617
+ qr_scale,
618
+ kv_score,
619
+ indexer_kv_score,
620
+ indexer_weights,
621
+ positions,
622
+ o_padded,
623
+ )
624
+ finally:
625
+ self.aux_stream_list = aux_streams
626
+
627
+ def forward(
628
+ self,
629
+ positions: torch.Tensor,
630
+ hidden_states: torch.Tensor,
631
+ llama_4_scaling: torch.Tensor | None = None,
632
+ ) -> torch.Tensor:
633
+ # Pre-allocate attention output with FlashMLA-padded head count.
634
+ # The op writes into `o_padded`; we slice to n_local_heads after.
635
+ num_tokens = hidden_states.shape[0]
636
+ o_padded = torch.empty(
637
+ (num_tokens, self.padded_heads, self.head_dim),
638
+ dtype=hidden_states.dtype,
639
+ device=hidden_states.device,
640
+ )
641
+
642
+ if self._enable_csa_multi_stream():
643
+ # The ROCm override consumes these sentinels inside the capture
644
+ # boundary, moving the stream fan-out ahead of the projections.
645
+ self._prepare_and_attn_fn(
646
+ hidden_states,
647
+ None,
648
+ None,
649
+ None,
650
+ None,
651
+ None,
652
+ None,
653
+ positions,
654
+ o_padded,
655
+ )
656
+ else:
657
+ self._run_sequential_pipeline(hidden_states, positions, o_padded)
658
+
659
+ o = o_padded[:, : self.n_local_heads, :]
660
+
661
+ # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
662
+ return self._o_proj(o, positions)
663
+
664
+ def _prepare_and_attn(
665
+ self,
666
+ hidden_states: torch.Tensor,
667
+ qr: torch.Tensor | None,
668
+ kv: torch.Tensor | None,
669
+ qr_scale: torch.Tensor | None,
670
+ kv_score: torch.Tensor | None,
671
+ indexer_kv_score: torch.Tensor | None,
672
+ indexer_weights: torch.Tensor | None,
673
+ positions: torch.Tensor,
674
+ o_padded: torch.Tensor,
675
+ ) -> None:
676
+ """Run the ROCm CSA fork/join inside the graph capture boundary."""
677
+ aux_streams = self.aux_stream_list
678
+ # The sequential pipeline disables aux_stream_list before calling
679
+ # back with real projection inputs; aux_streams is None ends that
680
+ # recursion here.
681
+ if aux_streams is None:
682
+ saved_streams = self.aux_stream_list
683
+ self.aux_stream_list = None
684
+ try:
685
+ super()._prepare_and_attn(
686
+ hidden_states,
687
+ cast(torch.Tensor, qr),
688
+ cast(torch.Tensor, kv),
689
+ qr_scale,
690
+ cast(torch.Tensor, kv_score),
691
+ cast(torch.Tensor, indexer_kv_score),
692
+ cast(torch.Tensor, indexer_weights),
693
+ positions,
694
+ o_padded,
695
+ )
696
+ finally:
697
+ self.aux_stream_list = saved_streams
698
+ return
699
+
700
+ # Re-check: forward's gate ran inside a captured segment that
701
+ # _prepare_and_attn_eager (MRV1) then broke, making this region eager.
702
+ if not self._enable_csa_multi_stream():
703
+ self._run_sequential_pipeline(hidden_states, positions, o_padded)
704
+ return
705
+
706
+ indexer = self.indexer
707
+ compressor = self.compressor
708
+ assert indexer is not None and compressor is not None
709
+
710
+ def default_chain():
711
+ qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
712
+ qr_out, qr_scale_out, kv_out = self._split_qkv_and_norm(qr_kv)
713
+ q = self._wq_b_proj(qr_out, qr_scale_out).view(
714
+ -1, self.n_local_heads, self.head_dim
715
+ )
716
+ attn_metadata = get_forward_context().attn_metadata
717
+ q = self._fused_qnorm_rope_kv_insert(q, kv_out, positions, attn_metadata)
718
+ return q, qr_out, qr_scale_out, kv_out
719
+
720
+ def main_compressor_chain() -> None:
721
+ score = torch.mm(
722
+ hidden_states,
723
+ compressor.fused_wkv_wgate.weight.T,
724
+ out_dtype=torch.float32,
725
+ )
726
+ compressor(score, positions, self.rotary_emb)
727
+
728
+ def indexer_compressor_chain() -> None:
729
+ score = torch.mm(
730
+ hidden_states,
731
+ indexer.compressor.fused_wkv_wgate.weight.T,
732
+ out_dtype=torch.float32,
733
+ )
734
+ indexer.compressor(score, positions, self.indexer_rotary_emb)
735
+
736
+ (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel(
737
+ default_chain,
738
+ [main_compressor_chain, indexer_compressor_chain],
739
+ self.ln_events[0],
740
+ self.ln_events[1:3],
741
+ aux_streams[:2],
742
+ enable=True,
743
+ )
744
+
745
+ indexer_weights_out, _ = indexer.weights_proj(hidden_states)
746
+ # The indexer compressor already ran on aux stream 1; build queries only.
747
+ index_q, index_q_scale, weights = indexer(
748
+ hidden_states,
749
+ qr_out,
750
+ None,
751
+ indexer_weights_out,
752
+ positions,
753
+ self.indexer_rotary_emb,
754
+ qr_scale_out,
755
+ skip_compressor=True,
756
+ )
757
+ self._sparse_indexer_and_attn(
758
+ hidden_states,
759
+ index_q,
760
+ index_q_scale,
761
+ weights,
762
+ q,
763
+ kv_out,
764
+ positions,
765
+ o_padded,
766
+ )
767
+
768
@classmethod
769
def get_padded_num_q_heads(cls, num_heads: int) -> int:
770
return num_heads
vllm/models/deepseek_v4/attention.py
+17
-14
@@ -318,7 +318,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
318
# graph and MRV1 produces garbage (#51430).
319
self._prepare_and_attn_fn = self._prepare_and_attn_eager
320
321
- # Will be None on ROCm for now.
321
self.aux_stream_list = aux_stream_list
322
# [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
323
# [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
@@ -584,7 +583,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
583
584
# Keep Q projection and KV insertion on the default stream. The indexer
585
# and MLA compressor use aux streams 0 and 1; aux 2 is internal to the
587
- # indexer. ROCm runs the same work sequentially without aux streams.
586
+ # indexer.
587
if indexer is not None:
588
assert compressor is not None
589
q, (indexer_inputs, _) = execute_in_parallel(
@@ -671,7 +670,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
670
# fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs
671
# on aux streams 0..2 when their owning module exists. ln_events[0]
672
# is the fan-out start event; ln_events[1..3] are per-aux done events.
674
- # On ROCm, aux_streams is None and execute_in_parallel runs serially.
673
aux_fns: list[Callable[[], Any] | None] = [None, None, None]
674
675
if self.compressor is not None:
@@ -1055,11 +1053,12 @@ class DeepseekV4Indexer(nn.Module):
1053
self,
1054
hidden_states: torch.Tensor,
1055
qr: torch.Tensor,
1058
- compressed_kv_score: torch.Tensor,
1056
+ compressed_kv_score: torch.Tensor | None,
1057
indexer_weights: torch.Tensor,
1058
positions: torch.Tensor,
1059
rotary_emb: nn.Module,
1060
qr_scale: torch.Tensor | None = None,
1061
+ skip_compressor: bool = False,
1062
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
1063
compressor = self.compressor
1064
@@ -1072,7 +1071,8 @@ class DeepseekV4Indexer(nn.Module):
1071
):
1072
# candidates num smaller than topk, every candidate is selected
1073
# but we still need to build k cache
1075
- compressor(compressed_kv_score, positions, rotary_emb)
1074
+ if not skip_compressor:
1075
+ compressor(compressed_kv_score, positions, rotary_emb)
1076
assert self.topk_indices_buffer is not None
1077
num_tokens = (
1078
indexer_metadata.num_decode_tokens
@@ -1102,15 +1102,18 @@ class DeepseekV4Indexer(nn.Module):
1102
use_fp4=self.use_fp4_kv,
1103
)
1104
1105
- # compressor returns None and writes K to the indexer KV cache; the
1106
- # join orders that write before indexer_op (skip_k_cache_insert=True).
1107
- (q_quant, weights), _ = maybe_execute_in_parallel(
1108
- wq_b_and_q_quant,
1109
- lambda: compressor(compressed_kv_score, positions, rotary_emb),
1110
- self.ln_events[0],
1111
- self.ln_events[1],
1112
- self.aux_stream,
1113
- )
1105
+ if not skip_compressor:
1106
+ # compressor returns None and writes K to the indexer KV cache; the
1107
+ # join orders that write before indexer_op (skip_k_cache_insert=True).
1108
+ (q_quant, weights), _ = maybe_execute_in_parallel(
1109
+ wq_b_and_q_quant,
1110
+ lambda: compressor(compressed_kv_score, positions, rotary_emb),
1111
+ self.ln_events[0],
1112
+ self.ln_events[1],
1113
+ self.aux_stream,
1114
+ )
1115
+ else:
1116
+ q_quant, weights = wq_b_and_q_quant()
1117
if isinstance(q_quant, tuple):
1118
q, q_scale = q_quant
1119
else:
vllm/models/deepseek_v4/cpu/cpu_sparse.py
+2
-1
@@ -198,11 +198,12 @@ class DeepseekV4CPUIndexer(DeepseekV4Indexer):
198
self,
199
hidden_states: torch.Tensor,
200
qr: torch.Tensor,
201
- compressed_kv_score: torch.Tensor,
201
+ compressed_kv_score: torch.Tensor | None,
202
indexer_weights: torch.Tensor,
203
positions: torch.Tensor,
204
rotary_emb: nn.Module,
205
qr_scale: torch.Tensor | None = None,
206
+ skip_compressor: bool = False,
207
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
208
"""CPU override: no aux streams, so wq_b_and_q_quant and the
209
compressor run straight-line instead of through