[ROCm][Refactor] Move DeepSeek-V4/V4.1 multi-stream overlap gate to ROCm platform (#58983)
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Shanshan Shen committed
Sep 28, 2026 at 20:46 UTC
7ba3df63cbe2e3e7aca19074bed2958311f46400
4 files changed
+40
-37
vllm/models/deepseek_v4/amd/rocm.py
+6
-19
@@ -649,23 +649,6 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
649
# Disable indexer inner overlap.
650
self.indexer.aux_stream = None
651
652
- def _enable_multi_stream_overlap(self) -> bool:
653
- """ROCm multi-stream gates: streams and capture region.
654
-
655
- Dict metadata marks piecewise cudagraph, whose eager breaks rebuild
656
- the attention inputs on the owning stream. Forking side streams
657
- there would rely on runtime HIP event sync, which is unreliable in
658
- this overlap on ROCm (event waits can hang), so multi-stream only
659
- runs where the fork/join becomes static graph edges: inside capture,
660
- or with non-dict metadata (full cudagraph or the profile run), which
661
- has no eager breaks. Covers both the HCA and CSA forks.
662
- """
663
- attn_metadata = get_forward_context().attn_metadata
664
- return self.aux_stream_list is not None and (
665
- torch.cuda.is_current_stream_capturing()
666
- or not isinstance(attn_metadata, dict)
667
- )
668
-
652
def _run_sequential_pipeline(
653
self,
654
hidden_states: torch.Tensor,
@@ -709,7 +692,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
692
device=hidden_states.device,
693
)
694
712
- if self._enable_multi_stream_overlap():
695
+ if current_platform.enable_multi_stream_overlap(
696
+ self.aux_stream_list, get_forward_context().attn_metadata
697
+ ):
698
# The ROCm override consumes these sentinels inside the capture
699
# boundary, moving the stream fan-out ahead of the projections.
700
self._prepare_and_attn_fn(
@@ -769,7 +754,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
754
755
# Re-check: forward's gate ran inside a captured segment that
756
# _prepare_and_attn_eager (MRV1) then broke, making this region eager.
772
- if not self._enable_multi_stream_overlap():
757
+ if not current_platform.enable_multi_stream_overlap(
758
+ self.aux_stream_list, get_forward_context().attn_metadata
759
+ ):
760
self._run_sequential_pipeline(hidden_states, positions, o_padded)
761
return
762
vllm/models/deepseek_v41/amd/rocm.py
+3
-18
@@ -615,23 +615,6 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
615
self.aux_stream_list = aux_streams
616
return super()._run_parallel_input_projections(hidden_states)
617
618
- def _enable_multi_stream_overlap(self) -> bool:
619
- """ROCm multi-stream gates: streams and capture region.
620
-
621
- Dict metadata marks piecewise cudagraph, whose eager breaks rebuild
622
- the attention inputs on the owning stream. Forking side streams
623
- there would rely on runtime HIP event sync, which is unreliable in
624
- this overlap on ROCm (event waits can hang), so multi-stream only
625
- runs where the fork/join becomes static graph edges: inside capture,
626
- or with non-dict metadata (full cudagraph or the profile run), which
627
- has no eager breaks.
628
- """
629
- attn_metadata = get_forward_context().attn_metadata
630
- return self.aux_stream_list is not None and (
631
- torch.cuda.is_current_stream_capturing()
632
- or not isinstance(attn_metadata, dict)
633
- )
634
-
618
def forward(
619
self,
620
positions: torch.Tensor,
@@ -641,7 +624,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
624
if (
625
self.indexer is None
626
or self._prepare_and_attn_fn == self._prepare_and_attn_eager
644
- or not self._enable_multi_stream_overlap()
627
+ or not current_platform.enable_multi_stream_overlap(
628
+ self.aux_stream_list, get_forward_context().attn_metadata
629
+ )
630
):
631
# Sequential fallback: no forks outside the capture region, where
632
# HIP event sync is unreliable; MRV1 keeps the input prep in its
vllm/platforms/interface.py
+10
@@ -1120,6 +1120,16 @@ class Platform:
1120
"""
1121
return False
1122
1123
+ @classmethod
1124
+ def enable_multi_stream_overlap(
1125
+ cls,
1126
+ aux_stream_list: list[torch.cuda.Stream] | None,
1127
+ attn_metadata: object,
1128
+ ) -> bool:
1129
+ """Whether the current platform should enable multi-stream overlap
1130
+ for the given aux streams and attention metadata."""
1131
+ return False
1132
+
1133
@classmethod
1134
def supports_mx(cls) -> bool:
1135
"""Returns whether the current platform supports MX types."""
vllm/platforms/rocm.py
+21
@@ -1106,6 +1106,27 @@ class RocmPlatform(Platform):
1106
def is_navi(cls) -> bool:
1107
return "gfx1" in _GCN_ARCH
1108
1109
+ @classmethod
1110
+ def enable_multi_stream_overlap(
1111
+ cls,
1112
+ aux_stream_list: list[torch.cuda.Stream] | None,
1113
+ attn_metadata: object,
1114
+ ) -> bool:
1115
+ """ROCm multi-stream gates: streams and capture region.
1116
+
1117
+ Dict metadata marks piecewise cudagraph, whose eager breaks rebuild
1118
+ the attention inputs on the owning stream. Forking side streams
1119
+ there would rely on runtime HIP event sync, which is unreliable in
1120
+ this overlap on ROCm (event waits can hang), so multi-stream only
1121
+ runs where the fork/join becomes static graph edges: inside capture,
1122
+ or with non-dict metadata (full cudagraph or the profile run), which
1123
+ has no eager breaks.
1124
+ """
1125
+ return aux_stream_list is not None and (
1126
+ torch.cuda.is_current_stream_capturing()
1127
+ or not isinstance(attn_metadata, dict)
1128
+ )
1129
+
1130
@classmethod
1131
def get_static_graph_wrapper_cls(cls) -> str:
1132
return "vllm.compilation.cuda_graph.CUDAGraphWrapper"