diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 44797874a4..a268ca1974 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -612,6 +612,7 @@ class Indexer(nn.Module): cache_config: CacheConfig | None, topk_indices_buffer: torch.Tensor | None, prefix: str = "", + is_inplace_rope: bool = False, ): super().__init__() self.vllm_config = vllm_config @@ -673,15 +674,21 @@ class Indexer(nn.Module): self.topk_indices_buffer, ) + self.is_inplace_rope = is_inplace_rope + def forward( self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb ) -> torch.Tensor: q, _ = self.wq_b(qr) q = q.view(-1, self.n_head, self.head_dim) - if current_platform.is_rocm(): + if current_platform.is_rocm() and self.is_inplace_rope: # This path should works on all platform, will remove extra # branches in the future + # This fast path relies on rotary_emb mutating q and k inplace. + # On ROCm, this is only valid for kernels used as custom ops. + # In pytorch-native rope for inductor fusion, rotated q/k tensors + # are not mutated inplace but returned as new tensors. # Fused wk + weights_proj: one GEMM, then split kw, _ = self.wk_weights_proj(hidden_states) k = kw[:, : self.head_dim] @@ -1002,6 +1009,7 @@ class DeepseekV2MLAAttention(nn.Module): cache_config, topk_indices_buffer, f"{prefix}.indexer", + is_inplace_rope=self.indexer_rope_emb.enabled(), ) # Enable IndexCache for DeepSeek models to reduce redundant top-k diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 80731296fc..7ddcc49344 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -107,6 +107,7 @@ def indexer_k_quant_and_cache_triton( kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype) kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32) head_tile_size = head_tile_size // kv_cache.element_size() + layout = "NORMAL" if block_size == 1 else "SHUFFLE" grid = (num_tokens,) _indexer_k_quant_and_cache_kernel[grid]( k, @@ -118,7 +119,7 @@ def indexer_k_quant_and_cache_triton( block_size, num_tokens, head_dim, - "SHUFFLE", + layout, block_tile_size, head_tile_size, IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz, @@ -229,6 +230,7 @@ def cp_gather_indexer_k_quant_cache_triton( k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32) grid = (num_tokens,) k_fp8_scale = k_fp8_scale.view(torch.float32) + layout = "NORMAL" if block_size == 1 else "SHUFFLE" _cp_gather_indexer_quant_cache_kernel[grid]( k_cache_value, k_cache_scale, @@ -241,7 +243,7 @@ def cp_gather_indexer_k_quant_cache_triton( block_table_stride, k_cache_value.stride(0), k_cache_scale.stride(0), - "SHUFFLE", + layout, head_dim, block_tile_size, head_tile_size, @@ -433,7 +435,7 @@ def rocm_fp8_paged_mqa_logits( block_tables, max_model_len, ChunkK=256, - Preshuffle=block_size == 64, + Preshuffle=block_size > 1, KVBlockSize=block_size, WavePerEU=2, )