[ROCm][DSV4][Perf] Fuse DeepSeek V4 C4 compressor GEMMs (#53838)
Signed-off-by: fai <fangzhouai@gmail.com>
Fangzhou Ai committed
Aug 26, 2026 at 08:36 UTC
657f9b9ce241a5452a36fc37d8e2e1e936722bfc
3 files changed
+221
tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
new
+141
@@ -0,0 +1,141 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+import vllm.model_executor.offloader as offloader
7
+from vllm.models.deepseek_v4.amd.rocm import (
8
+ DeepseekV4ROCMAiterMLAAttention,
9
+)
10
+
11
+
12
+class _WeightOnlyLinear(torch.nn.Module):
13
+ def __init__(self, weight: torch.Tensor) -> None:
14
+ super().__init__()
15
+ self.weight = torch.nn.Parameter(weight, requires_grad=False)
16
+
17
+
18
+class _Compressor(torch.nn.Module):
19
+ def __init__(self, weight: torch.Tensor) -> None:
20
+ super().__init__()
21
+ self.fused_wkv_wgate = _WeightOnlyLinear(weight)
22
+
23
+
24
+class _IndexerWeightsProjection(torch.nn.Module):
25
+ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]:
26
+ return hidden_states.sum(dim=-1, keepdim=True), None
27
+
28
+
29
+class _Indexer(torch.nn.Module):
30
+ def __init__(self, compressor_weight: torch.Tensor) -> None:
31
+ super().__init__()
32
+ self.compressor = _Compressor(compressor_weight)
33
+ self.weights_proj = _IndexerWeightsProjection()
34
+
35
+
36
+def _make_attention(
37
+ main_weight: torch.Tensor,
38
+ indexer_weight: torch.Tensor,
39
+) -> DeepseekV4ROCMAiterMLAAttention:
40
+ attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention)
41
+ torch.nn.Module.__init__(attention)
42
+ attention.compressor = _Compressor(main_weight)
43
+ attention.indexer = _Indexer(indexer_weight)
44
+ attention.register_buffer("_fused_compressor_weight", None, persistent=False)
45
+ attention._fused_compressor_split_sizes = None
46
+ return attention
47
+
48
+
49
+def test_prepare_compressor_gemm_fusion_aliases_original_weights():
50
+ main_value = torch.arange(12, dtype=torch.float32).reshape(4, 3)
51
+ indexer_value = torch.arange(6, dtype=torch.float32).reshape(2, 3) + 100
52
+ attention = _make_attention(main_value, indexer_value)
53
+ main_weight = attention.compressor.fused_wkv_wgate.weight
54
+ indexer_weight = attention.indexer.compressor.fused_wkv_wgate.weight
55
+
56
+ assert attention.prepare_compressor_gemm_fusion() is True
57
+
58
+ fused_weight = attention._fused_compressor_weight
59
+ assert fused_weight is not None
60
+ assert attention._fused_compressor_split_sizes == (4, 2)
61
+ torch.testing.assert_close(fused_weight[:4], main_value)
62
+ torch.testing.assert_close(fused_weight[4:], indexer_value)
63
+
64
+ fused_storage = fused_weight.untyped_storage().data_ptr()
65
+ assert main_weight.untyped_storage().data_ptr() == fused_storage
66
+ assert indexer_weight.untyped_storage().data_ptr() == fused_storage
67
+ assert main_weight.storage_offset() == 0
68
+ assert indexer_weight.storage_offset() == main_weight.numel()
69
+
70
+ state_dict = attention.state_dict()
71
+ assert "_fused_compressor_weight" not in state_dict
72
+ torch.testing.assert_close(
73
+ state_dict["compressor.fused_wkv_wgate.weight"], main_value
74
+ )
75
+ torch.testing.assert_close(
76
+ state_dict["indexer.compressor.fused_wkv_wgate.weight"],
77
+ indexer_value,
78
+ )
79
+
80
+ fused_data_ptr = fused_weight.data_ptr()
81
+ assert attention.prepare_compressor_gemm_fusion() is False
82
+ assert attention._fused_compressor_weight.data_ptr() == fused_data_ptr
83
+
84
+
85
+def test_prepare_compressor_gemm_fusion_skips_weight_offloading(monkeypatch):
86
+ attention = _make_attention(torch.ones(4, 3), torch.ones(2, 3))
87
+ main_weight = attention.compressor.fused_wkv_wgate.weight
88
+ indexer_weight = attention.indexer.compressor.fused_wkv_wgate.weight
89
+
90
+ monkeypatch.setattr(offloader, "get_offloader", object)
91
+
92
+ assert attention.prepare_compressor_gemm_fusion() is False
93
+ assert attention._fused_compressor_weight is None
94
+ assert main_weight.untyped_storage().data_ptr() != (
95
+ indexer_weight.untyped_storage().data_ptr()
96
+ )
97
+
98
+
99
+def test_fused_input_projection_uses_one_mm_and_returns_expected_tuple(
100
+ monkeypatch,
101
+):
102
+ main_weight = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=torch.float32)
103
+ indexer_weight = torch.tensor([[0.0, 0.0, 1.0]], dtype=torch.float32)
104
+ attention = _make_attention(main_weight, indexer_weight)
105
+ assert attention.prepare_compressor_gemm_fusion() is True
106
+
107
+ hidden_states = torch.tensor(
108
+ [[2.0, 3.0, 5.0], [7.0, 11.0, 13.0]], dtype=torch.float32
109
+ )
110
+ qr_kv = torch.full((2, 4), 17.0)
111
+ monkeypatch.setattr(
112
+ DeepseekV4ROCMAiterMLAAttention,
113
+ "_fused_wqa_wkv_gemm",
114
+ lambda self, hidden: qr_kv,
115
+ )
116
+
117
+ mm_calls: list[tuple[torch.Tensor, torch.Tensor, torch.dtype | None]] = []
118
+
119
+ def cpu_mm_with_out_dtype(
120
+ left: torch.Tensor,
121
+ right: torch.Tensor,
122
+ *,
123
+ out_dtype: torch.dtype | None = None,
124
+ ) -> torch.Tensor:
125
+ mm_calls.append((left, right, out_dtype))
126
+ return torch.matmul(left, right).to(out_dtype)
127
+
128
+ monkeypatch.setattr(torch, "mm", cpu_mm_with_out_dtype)
129
+
130
+ qr_out, kv_score, indexer_kv_score, indexer_weights = (
131
+ attention._run_parallel_input_projections(hidden_states)
132
+ )
133
+
134
+ assert qr_out is qr_kv
135
+ torch.testing.assert_close(kv_score, hidden_states[:, :2])
136
+ torch.testing.assert_close(indexer_kv_score, hidden_states[:, 2:])
137
+ torch.testing.assert_close(indexer_weights, hidden_states.sum(dim=-1, keepdim=True))
138
+ assert len(mm_calls) == 1
139
+ assert mm_calls[0][0] is hidden_states
140
+ assert mm_calls[0][2] is torch.float32
141
+ assert mm_calls[0][1].shape == (3, 3)
vllm/models/deepseek_v4/amd/model.py
+7
@@ -1044,11 +1044,18 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
1044
1045
def process_weights_after_loading(self) -> None:
1046
# After per-layer quant finalize, so we preshuffle the final fp8 weights.
1047
+ fused_compressor_layers = 0
1048
for module in self.modules():
1049
if isinstance(module, DeepseekV4ROCMAiterMLAAttention):
1050
+ fused_compressor_layers += module.prepare_compressor_gemm_fusion()
1051
module.prepare_attn_preshuffle()
1052
elif isinstance(module, DeepseekV4MLP):
1053
module.prepare_gateup_preshuffle()
1054
+ if fused_compressor_layers:
1055
+ logger.info(
1056
+ "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers",
1057
+ fused_compressor_layers,
1058
+ )
1059
1060
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1061
return self.model.get_expert_mapping()
vllm/models/deepseek_v4/amd/rocm.py
+73
@@ -11,6 +11,7 @@ from vllm.distributed import (
11
tensor_model_parallel_all_reduce,
12
)
13
from vllm.forward_context import get_forward_context
14
+from vllm.logger import init_logger
15
from vllm.models.deepseek_v4.attention import DeepseekV4Attention
16
from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache
17
from vllm.models.deepseek_v4.sparse_mla import (
@@ -36,6 +37,8 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
37
)
38
from vllm.v1.worker.workspace import current_workspace_manager
39
40
+logger = init_logger(__name__)
41
+
42
43
def _trust_dsv4_extra_cache_nan_free(
44
kv_cache_dtype: str,
@@ -486,6 +489,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
489
# Block scale for the preshuffled weight; None = not preshuffled.
490
self._wqa_wkv_scale: torch.Tensor | None = None
491
self._wo_b_scale: torch.Tensor | None = None
492
+ self._fused_compressor_weight: torch.Tensor | None
493
+ self.register_buffer("_fused_compressor_weight", None, persistent=False)
494
+ self._fused_compressor_split_sizes: tuple[int, int] | None = None
495
496
@classmethod
497
def get_padded_num_q_heads(cls, num_heads: int) -> int:
@@ -524,6 +530,46 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
530
self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
531
self._wo_b_scale = _prep(self.wo_b)
532
533
+ def prepare_compressor_gemm_fusion(self) -> bool:
534
+ if self._fused_compressor_weight is not None:
535
+ return False
536
+
537
+ from vllm.model_executor.offloader import NoopOffloader, get_offloader
538
+
539
+ if not isinstance(get_offloader(), NoopOffloader):
540
+ logger.warning_once(
541
+ "DeepSeek V4 compressor GEMM fusion is incompatible with "
542
+ "weight offloading and will remain disabled."
543
+ )
544
+ return False
545
+
546
+ compressor = self.compressor
547
+ indexer = self.indexer
548
+ if compressor is None or indexer is None:
549
+ return False
550
+
551
+ main_weight = compressor.fused_wkv_wgate.weight
552
+ indexer_weight = indexer.compressor.fused_wkv_wgate.weight
553
+ if main_weight.ndim != 2 or indexer_weight.ndim != 2:
554
+ raise ValueError("DeepSeek V4 compressor weights must be matrices")
555
+ if main_weight.shape[1] != indexer_weight.shape[1]:
556
+ raise ValueError("DeepSeek V4 compressor weights must share K")
557
+ if main_weight.dtype != indexer_weight.dtype:
558
+ raise ValueError("DeepSeek V4 compressor weights must share dtype")
559
+ if main_weight.device != indexer_weight.device:
560
+ raise ValueError("DeepSeek V4 compressor weights must share device")
561
+
562
+ main_size = main_weight.shape[0]
563
+ indexer_size = indexer_weight.shape[0]
564
+ fused_weight = torch.cat((main_weight, indexer_weight), dim=0)
565
+ with torch.no_grad():
566
+ main_weight.set_(fused_weight[:main_size])
567
+ indexer_weight.set_(fused_weight[main_size:])
568
+
569
+ self._fused_compressor_weight = fused_weight
570
+ self._fused_compressor_split_sizes = (main_size, indexer_size)
571
+ return True
572
+
573
def _bpre_attn_gemm(
574
self,
575
weight: torch.Tensor,
@@ -548,6 +594,33 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
594
)
595
return super()._fused_wqa_wkv_gemm(hidden_states)
596
597
+ def _run_parallel_input_projections(
598
+ self, hidden_states: torch.Tensor
599
+ ) -> tuple[
600
+ torch.Tensor,
601
+ torch.Tensor | None,
602
+ torch.Tensor | None,
603
+ torch.Tensor | None,
604
+ ]:
605
+ fused_weight = self._fused_compressor_weight
606
+ split_sizes = self._fused_compressor_split_sizes
607
+ if fused_weight is None or split_sizes is None:
608
+ return super()._run_parallel_input_projections(hidden_states)
609
+
610
+ indexer = self.indexer
611
+ if indexer is None:
612
+ raise RuntimeError("Fused compressor weight requires a C4 indexer")
613
+
614
+ qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
615
+ fused_scores = torch.mm(
616
+ hidden_states,
617
+ fused_weight.T,
618
+ out_dtype=torch.float32,
619
+ )
620
+ kv_score, indexer_kv_score = fused_scores.split(split_sizes, dim=-1)
621
+ indexer_weights, _ = indexer.weights_proj(hidden_states)
622
+ return qr_kv, kv_score, indexer_kv_score, indexer_weights
623
+
624
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
625
# ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
626
z = rocm_inv_rope_einsum(