diff --git a/tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py b/tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py new file mode 100644 index 0000000000..2ac22eae02 --- /dev/null +++ b/tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.offloader as offloader +from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLAAttention, +) + + +class _WeightOnlyLinear(torch.nn.Module): + def __init__(self, weight: torch.Tensor) -> None: + super().__init__() + self.weight = torch.nn.Parameter(weight, requires_grad=False) + + +class _Compressor(torch.nn.Module): + def __init__(self, weight: torch.Tensor) -> None: + super().__init__() + self.fused_wkv_wgate = _WeightOnlyLinear(weight) + + +class _IndexerWeightsProjection(torch.nn.Module): + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]: + return hidden_states.sum(dim=-1, keepdim=True), None + + +class _Indexer(torch.nn.Module): + def __init__(self, compressor_weight: torch.Tensor) -> None: + super().__init__() + self.compressor = _Compressor(compressor_weight) + self.weights_proj = _IndexerWeightsProjection() + + +def _make_attention( + main_weight: torch.Tensor, + indexer_weight: torch.Tensor, +) -> DeepseekV4ROCMAiterMLAAttention: + attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention) + torch.nn.Module.__init__(attention) + attention.compressor = _Compressor(main_weight) + attention.indexer = _Indexer(indexer_weight) + attention.register_buffer("_fused_compressor_weight", None, persistent=False) + attention._fused_compressor_split_sizes = None + return attention + + +def test_prepare_compressor_gemm_fusion_aliases_original_weights(): + main_value = torch.arange(12, dtype=torch.float32).reshape(4, 3) + indexer_value = torch.arange(6, dtype=torch.float32).reshape(2, 3) + 100 + attention = _make_attention(main_value, indexer_value) + main_weight = attention.compressor.fused_wkv_wgate.weight + indexer_weight = attention.indexer.compressor.fused_wkv_wgate.weight + + assert attention.prepare_compressor_gemm_fusion() is True + + fused_weight = attention._fused_compressor_weight + assert fused_weight is not None + assert attention._fused_compressor_split_sizes == (4, 2) + torch.testing.assert_close(fused_weight[:4], main_value) + torch.testing.assert_close(fused_weight[4:], indexer_value) + + fused_storage = fused_weight.untyped_storage().data_ptr() + assert main_weight.untyped_storage().data_ptr() == fused_storage + assert indexer_weight.untyped_storage().data_ptr() == fused_storage + assert main_weight.storage_offset() == 0 + assert indexer_weight.storage_offset() == main_weight.numel() + + state_dict = attention.state_dict() + assert "_fused_compressor_weight" not in state_dict + torch.testing.assert_close( + state_dict["compressor.fused_wkv_wgate.weight"], main_value + ) + torch.testing.assert_close( + state_dict["indexer.compressor.fused_wkv_wgate.weight"], + indexer_value, + ) + + fused_data_ptr = fused_weight.data_ptr() + assert attention.prepare_compressor_gemm_fusion() is False + assert attention._fused_compressor_weight.data_ptr() == fused_data_ptr + + +def test_prepare_compressor_gemm_fusion_skips_weight_offloading(monkeypatch): + attention = _make_attention(torch.ones(4, 3), torch.ones(2, 3)) + main_weight = attention.compressor.fused_wkv_wgate.weight + indexer_weight = attention.indexer.compressor.fused_wkv_wgate.weight + + monkeypatch.setattr(offloader, "get_offloader", object) + + assert attention.prepare_compressor_gemm_fusion() is False + assert attention._fused_compressor_weight is None + assert main_weight.untyped_storage().data_ptr() != ( + indexer_weight.untyped_storage().data_ptr() + ) + + +def test_fused_input_projection_uses_one_mm_and_returns_expected_tuple( + monkeypatch, +): + main_weight = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float32) + indexer_weight = torch.tensor([[0.0, 0.0, 1.0]], dtype=torch.float32) + attention = _make_attention(main_weight, indexer_weight) + assert attention.prepare_compressor_gemm_fusion() is True + + hidden_states = torch.tensor( + [[2.0, 3.0, 5.0], [7.0, 11.0, 13.0]], dtype=torch.float32 + ) + qr_kv = torch.full((2, 4), 17.0) + monkeypatch.setattr( + DeepseekV4ROCMAiterMLAAttention, + "_fused_wqa_wkv_gemm", + lambda self, hidden: qr_kv, + ) + + mm_calls: list[tuple[torch.Tensor, torch.Tensor, torch.dtype | None]] = [] + + def cpu_mm_with_out_dtype( + left: torch.Tensor, + right: torch.Tensor, + *, + out_dtype: torch.dtype | None = None, + ) -> torch.Tensor: + mm_calls.append((left, right, out_dtype)) + return torch.matmul(left, right).to(out_dtype) + + monkeypatch.setattr(torch, "mm", cpu_mm_with_out_dtype) + + qr_out, kv_score, indexer_kv_score, indexer_weights = ( + attention._run_parallel_input_projections(hidden_states) + ) + + assert qr_out is qr_kv + torch.testing.assert_close(kv_score, hidden_states[:, :2]) + torch.testing.assert_close(indexer_kv_score, hidden_states[:, 2:]) + torch.testing.assert_close(indexer_weights, hidden_states.sum(dim=-1, keepdim=True)) + assert len(mm_calls) == 1 + assert mm_calls[0][0] is hidden_states + assert mm_calls[0][2] is torch.float32 + assert mm_calls[0][1].shape == (3, 3) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 2a82915f6e..7a67c8f420 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -1044,11 +1044,18 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): def process_weights_after_loading(self) -> None: # After per-layer quant finalize, so we preshuffle the final fp8 weights. + fused_compressor_layers = 0 for module in self.modules(): if isinstance(module, DeepseekV4ROCMAiterMLAAttention): + fused_compressor_layers += module.prepare_compressor_gemm_fusion() module.prepare_attn_preshuffle() elif isinstance(module, DeepseekV4MLP): module.prepare_gateup_preshuffle() + if fused_compressor_layers: + logger.info( + "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers", + fused_compressor_layers, + ) def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 5e7609bd11..dc6620e6c8 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -11,6 +11,7 @@ from vllm.distributed import ( tensor_model_parallel_all_reduce, ) from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache from vllm.models.deepseek_v4.sparse_mla import ( @@ -36,6 +37,8 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( ) from vllm.v1.worker.workspace import current_workspace_manager +logger = init_logger(__name__) + def _trust_dsv4_extra_cache_nan_free( kv_cache_dtype: str, @@ -486,6 +489,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None + self._fused_compressor_weight: torch.Tensor | None + self.register_buffer("_fused_compressor_weight", None, persistent=False) + self._fused_compressor_split_sizes: tuple[int, int] | None = None @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: @@ -524,6 +530,46 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) self._wo_b_scale = _prep(self.wo_b) + def prepare_compressor_gemm_fusion(self) -> bool: + if self._fused_compressor_weight is not None: + return False + + from vllm.model_executor.offloader import NoopOffloader, get_offloader + + if not isinstance(get_offloader(), NoopOffloader): + logger.warning_once( + "DeepSeek V4 compressor GEMM fusion is incompatible with " + "weight offloading and will remain disabled." + ) + return False + + compressor = self.compressor + indexer = self.indexer + if compressor is None or indexer is None: + return False + + main_weight = compressor.fused_wkv_wgate.weight + indexer_weight = indexer.compressor.fused_wkv_wgate.weight + if main_weight.ndim != 2 or indexer_weight.ndim != 2: + raise ValueError("DeepSeek V4 compressor weights must be matrices") + if main_weight.shape[1] != indexer_weight.shape[1]: + raise ValueError("DeepSeek V4 compressor weights must share K") + if main_weight.dtype != indexer_weight.dtype: + raise ValueError("DeepSeek V4 compressor weights must share dtype") + if main_weight.device != indexer_weight.device: + raise ValueError("DeepSeek V4 compressor weights must share device") + + main_size = main_weight.shape[0] + indexer_size = indexer_weight.shape[0] + fused_weight = torch.cat((main_weight, indexer_weight), dim=0) + with torch.no_grad(): + main_weight.set_(fused_weight[:main_size]) + indexer_weight.set_(fused_weight[main_size:]) + + self._fused_compressor_weight = fused_weight + self._fused_compressor_split_sizes = (main_size, indexer_size) + return True + def _bpre_attn_gemm( self, weight: torch.Tensor, @@ -548,6 +594,33 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): ) return super()._fused_wqa_wkv_gemm(hidden_states) + def _run_parallel_input_projections( + self, hidden_states: torch.Tensor + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + fused_weight = self._fused_compressor_weight + split_sizes = self._fused_compressor_split_sizes + if fused_weight is None or split_sizes is None: + return super()._run_parallel_input_projections(hidden_states) + + indexer = self.indexer + if indexer is None: + raise RuntimeError("Fused compressor weight requires a C4 indexer") + + qr_kv = self._fused_wqa_wkv_gemm(hidden_states) + fused_scores = torch.mm( + hidden_states, + fused_weight.T, + out_dtype=torch.float32, + ) + kv_score, indexer_kv_score = fused_scores.split(split_sizes, dim=-1) + indexer_weights, _ = indexer.weights_proj(hidden_states) + return qr_kv, kv_score, indexer_kv_score, indexer_weights + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. z = rocm_inv_rope_einsum(