diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 349ff52679..c1a6903937 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -957,16 +957,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): self.hc_dim = self.hc_mult * config.hidden_size self.rms_norm_eps = config.rms_norm_eps - # Three aux streams: one per non-default input GEMM in - # DeepseekV4Attention._run_parallel_input_projections - # (compressor kv_score, indexer.weights_proj, indexer.compressor - # kv_score). fused_wqa_wkv stays on the default stream. - # Disable them on ROCm because of hang issues. - aux_stream_list = ( - None - if current_platform.is_rocm() - else [torch.cuda.Stream() for _ in range(3)] - ) + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] self.device = current_platform.device_type # Reserved topk indices buffer for all Indexer layers to reuse. diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index b39ad953e7..3cc9827c1f 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -23,6 +23,7 @@ from vllm.models.deepseek_v4.sparse_mla import ( from vllm.platforms import current_platform from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton +from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( CommonAttentionMetadata, ) @@ -571,6 +572,199 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): self.register_buffer("_fused_compressor_weight", None, persistent=False) self._fused_compressor_split_sizes: tuple[int, int] | None = None + if self.indexer is None: + # Only enable multi-stream overlap for CSA layer now. + self.aux_stream_list = None + else: + # Disable indexer inner overlap. + self.indexer.aux_stream = None + + def _enable_csa_multi_stream(self) -> bool: + """All CSA multi-stream gates: env var, streams, and capture region. + + Dict metadata marks piecewise cudagraph, whose eager breaks rebuild + the attention inputs on the owning stream. Forking side streams + there would rely on runtime HIP event sync, which is unreliable in + this overlap on ROCm (event waits can hang), so multi-stream only + runs where the fork/join becomes static graph edges: inside capture, + or with non-dict metadata (full cudagraph or the profile run), which + has no eager breaks. + """ + attn_metadata = get_forward_context().attn_metadata + return self.aux_stream_list is not None and ( + torch.cuda.is_current_stream_capturing() + or not isinstance(attn_metadata, dict) + ) + + def _run_sequential_pipeline( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Disable ROCm streams when the current execution region cannot overlap.""" + aux_streams = self.aux_stream_list + self.aux_stream_list = None + try: + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self._run_parallel_input_projections(hidden_states) + ) + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + finally: + self.aux_stream_list = aux_streams + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + if self._enable_csa_multi_stream(): + # The ROCm override consumes these sentinels inside the capture + # boundary, moving the stream fan-out ahead of the projections. + self._prepare_and_attn_fn( + hidden_states, + None, + None, + None, + None, + None, + None, + positions, + o_padded, + ) + else: + self._run_sequential_pipeline(hidden_states, positions, o_padded) + + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor | None, + kv: torch.Tensor | None, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor | None, + indexer_kv_score: torch.Tensor | None, + indexer_weights: torch.Tensor | None, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Run the ROCm CSA fork/join inside the graph capture boundary.""" + aux_streams = self.aux_stream_list + # The sequential pipeline disables aux_stream_list before calling + # back with real projection inputs; aux_streams is None ends that + # recursion here. + if aux_streams is None: + saved_streams = self.aux_stream_list + self.aux_stream_list = None + try: + super()._prepare_and_attn( + hidden_states, + cast(torch.Tensor, qr), + cast(torch.Tensor, kv), + qr_scale, + cast(torch.Tensor, kv_score), + cast(torch.Tensor, indexer_kv_score), + cast(torch.Tensor, indexer_weights), + positions, + o_padded, + ) + finally: + self.aux_stream_list = saved_streams + return + + # Re-check: forward's gate ran inside a captured segment that + # _prepare_and_attn_eager (MRV1) then broke, making this region eager. + if not self._enable_csa_multi_stream(): + self._run_sequential_pipeline(hidden_states, positions, o_padded) + return + + indexer = self.indexer + compressor = self.compressor + assert indexer is not None and compressor is not None + + def default_chain(): + qr_kv = self._fused_wqa_wkv_gemm(hidden_states) + qr_out, qr_scale_out, kv_out = self._split_qkv_and_norm(qr_kv) + q = self._wq_b_proj(qr_out, qr_scale_out).view( + -1, self.n_local_heads, self.head_dim + ) + attn_metadata = get_forward_context().attn_metadata + q = self._fused_qnorm_rope_kv_insert(q, kv_out, positions, attn_metadata) + return q, qr_out, qr_scale_out, kv_out + + def main_compressor_chain() -> None: + score = torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + compressor(score, positions, self.rotary_emb) + + def indexer_compressor_chain() -> None: + score = torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + indexer.compressor(score, positions, self.indexer_rotary_emb) + + (q, qr_out, qr_scale_out, kv_out), _ = execute_in_parallel( + default_chain, + [main_compressor_chain, indexer_compressor_chain], + self.ln_events[0], + self.ln_events[1:3], + aux_streams[:2], + enable=True, + ) + + indexer_weights_out, _ = indexer.weights_proj(hidden_states) + # The indexer compressor already ran on aux stream 1; build queries only. + index_q, index_q_scale, weights = indexer( + hidden_states, + qr_out, + None, + indexer_weights_out, + positions, + self.indexer_rotary_emb, + qr_scale_out, + skip_compressor=True, + ) + self._sparse_indexer_and_attn( + hidden_states, + index_q, + index_q_scale, + weights, + q, + kv_out, + positions, + o_padded, + ) + @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: return num_heads diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index d32a48cbe4..9c77980965 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -318,7 +318,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # graph and MRV1 produces garbage (#51430). self._prepare_and_attn_fn = self._prepare_and_attn_eager - # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins @@ -584,7 +583,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # Keep Q projection and KV insertion on the default stream. The indexer # and MLA compressor use aux streams 0 and 1; aux 2 is internal to the - # indexer. ROCm runs the same work sequentially without aux streams. + # indexer. if indexer is not None: assert compressor is not None q, (indexer_inputs, _) = execute_in_parallel( @@ -671,7 +670,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs # on aux streams 0..2 when their owning module exists. ln_events[0] # is the fan-out start event; ln_events[1..3] are per-aux done events. - # On ROCm, aux_streams is None and execute_in_parallel runs serially. aux_fns: list[Callable[[], Any] | None] = [None, None, None] if self.compressor is not None: @@ -1055,11 +1053,12 @@ class DeepseekV4Indexer(nn.Module): self, hidden_states: torch.Tensor, qr: torch.Tensor, - compressed_kv_score: torch.Tensor, + compressed_kv_score: torch.Tensor | None, indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + skip_compressor: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: compressor = self.compressor @@ -1072,7 +1071,8 @@ class DeepseekV4Indexer(nn.Module): ): # candidates num smaller than topk, every candidate is selected # but we still need to build k cache - compressor(compressed_kv_score, positions, rotary_emb) + if not skip_compressor: + compressor(compressed_kv_score, positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( indexer_metadata.num_decode_tokens @@ -1102,15 +1102,18 @@ class DeepseekV4Indexer(nn.Module): use_fp4=self.use_fp4_kv, ) - # compressor returns None and writes K to the indexer KV cache; the - # join orders that write before indexer_op (skip_k_cache_insert=True). - (q_quant, weights), _ = maybe_execute_in_parallel( - wq_b_and_q_quant, - lambda: compressor(compressed_kv_score, positions, rotary_emb), - self.ln_events[0], - self.ln_events[1], - self.aux_stream, - ) + if not skip_compressor: + # compressor returns None and writes K to the indexer KV cache; the + # join orders that write before indexer_op (skip_k_cache_insert=True). + (q_quant, weights), _ = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + else: + q_quant, weights = wq_b_and_q_quant() if isinstance(q_quant, tuple): q, q_scale = q_quant else: diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py index d51e589652..a2c3b00493 100644 --- a/vllm/models/deepseek_v4/cpu/cpu_sparse.py +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -198,11 +198,12 @@ class DeepseekV4CPUIndexer(DeepseekV4Indexer): self, hidden_states: torch.Tensor, qr: torch.Tensor, - compressed_kv_score: torch.Tensor, + compressed_kv_score: torch.Tensor | None, indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + skip_compressor: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """CPU override: no aux streams, so wq_b_and_q_quant and the compressor run straight-line instead of through