From 367cb81966f98f6ba73e5e3f529573d017c446d0 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Fri, 22 May 2026 21:22:27 -0400 Subject: [PATCH] [DSV4] More multi-stream enablement for c4a (#42925) Signed-off-by: Yongye Zhu --- vllm/models/deepseek_v4/nvidia/model.py | 7 ++ .../deepseek_v4/nvidia/ops/attention.py | 83 ++++++++++++------- 2 files changed, 61 insertions(+), 29 deletions(-) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 0a07e6f9fb..adce3e5295 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -935,6 +935,12 @@ class DeepseekV4Attention(nn.Module): self.indexer = None if self.compress_ratio == 4: # Only C4A uses sparse attention and hence has indexer. + # aux_stream_list[0] runs indexer.forward() in the wrapper; [2] is + # free here (outer GEMMs joined) for the inner overlap of + # wq_b+fused_indexer_q_rope_quant vs compressor. + indexer_aux_stream = ( + aux_stream_list[2] if aux_stream_list is not None else None + ) self.indexer = DeepseekV4Indexer( vllm_config, config=config, @@ -945,6 +951,7 @@ class DeepseekV4Attention(nn.Module): topk_indices_buffer=topk_indices_buffer, compress_ratio=self.compress_ratio, prefix=f"{prefix}.indexer", + aux_stream=indexer_aux_stream, ) mla_modules = DeepseekV4MLAModules( diff --git a/vllm/models/deepseek_v4/nvidia/ops/attention.py b/vllm/models/deepseek_v4/nvidia/ops/attention.py index d23b5a578b..5a98373788 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/attention.py +++ b/vllm/models/deepseek_v4/nvidia/ops/attention.py @@ -442,33 +442,38 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): # downstream reads q on default). Indexer/compressor go on aux for # overlap with default's GEMM + cache write. if self.indexer is not None: - aux_stream = ( - self.aux_stream_list[0] if self.aux_stream_list is not None else None - ) + aux_streams = self.aux_stream_list indexer = self.indexer # Local ref so the closure keeps a non-None type for mypy. assert self.compressor is not None compressor = self.compressor - def wq_b_kv_insert_and_compress() -> torch.Tensor: + def wq_b_kv_insert() -> torch.Tensor: q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - compressor(kv_score, positions, self.rotary_emb) return q - q, _ = maybe_execute_in_parallel( - wq_b_kv_insert_and_compress, - lambda: indexer( - hidden_states, - qr, - indexer_kv_score, - indexer_weights, - positions, - self.indexer_rotary_emb, - ), + # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs + # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the + # MLA compressor. Slot [2] is reserved for the indexer's inner + # overlap. ROCm (aux_streams is None) falls back to sequential. + q, _ = execute_in_parallel( + wq_b_kv_insert, + [ + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], self.ln_events[0], - self.ln_events[1], - aux_stream, + [self.ln_events[1], self.ln_events[2]], + [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, + enable=aux_streams is not None, ) elif self.compressor is not None: # wq_b + kv_insert on default, compressor on aux. @@ -792,6 +797,7 @@ class DeepseekV4Indexer(nn.Module): topk_indices_buffer: torch.Tensor | None, compress_ratio: int = 1, prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, ): super().__init__() self.vllm_config = vllm_config @@ -877,6 +883,13 @@ class DeepseekV4Indexer(nn.Module): use_fp4_cache=self.use_fp4_kv, ) + # None on ROCm — maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self.ln_events: list[torch.cuda.Event] = [ + torch.cuda.Event(), + torch.cuda.Event(), + ] + def forward( self, hidden_states: torch.Tensor, @@ -886,17 +899,29 @@ class DeepseekV4Indexer(nn.Module): positions: torch.Tensor, rotary_emb: nn.Module, ) -> torch.Tensor: - # ReplicatedLinear returns (output, bias); bias is None. - q, _ = self.wq_b(qr) - q = q.view(-1, self.n_head, self.head_dim) - k = self.compressor(compressed_kv_score, positions, rotary_emb) - q_quant, weights = fused_indexer_q_rope_quant( - positions, - q, - rotary_emb.cos_sin_cache, - indexer_weights, - self.softmax_scale, - self.n_head**-0.5, - use_fp4=self.use_fp4_kv, + compressor = self.compressor + + def wq_b_and_q_quant(): + # ReplicatedLinear returns (output, bias); bias is None. + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + return fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + 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), k = 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, ) return self.indexer_op(hidden_states, q_quant, k, weights)