iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][DSV4][Perf] Fuse DeepSeek V4 C4 compressor GEMMs (#53838)

Browse Source 
Signed-off-by: fai <fangzhouai@gmail.com>
This commit is contained in:
Fangzhou Ai
2026-08-26 23:36:14 +08:00
committed by GitHub
parent 080a66a69c
commit 657f9b9ce2
3 changed files with 221 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
+141
View File
@@ -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)
vllm/models/deepseek_v4/amd/model.py
+7
View File
@@ -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()
vllm/models/deepseek_v4/amd/rocm.py
+73
View File
@@ -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(