[ROCm][DSV4][Perf] Use FP8 WO_A output projection (#54894)
Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
yinfengLiu committed
Sep 20, 2026 at 17:13 UTC
0e110f696db450f5645ce1daf9945241212d6ef3
2 files changed
+217
-11
tests/models/test_deepseek_v4_rocm_wo_a.py
new
+44
@@ -0,0 +1,44 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0
8
+
9
+
10
+def test_wo_a_block_scale_to_e8m0_from_float():
11
+ scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32)
12
+
13
+ encoded = _wo_a_block_scale_to_e8m0(scale)
14
+
15
+ assert encoded is not None
16
+ torch.testing.assert_close(
17
+ encoded,
18
+ torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8),
19
+ )
20
+ assert encoded.is_contiguous()
21
+
22
+
23
+def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales():
24
+ raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8)
25
+ encoded = raw.view(torch.float8_e8m0fnu)
26
+
27
+ converted = _wo_a_block_scale_to_e8m0(encoded)
28
+
29
+ assert converted is not None
30
+ torch.testing.assert_close(converted, raw)
31
+
32
+
33
+@pytest.mark.parametrize(
34
+ "scale",
35
+ [
36
+ torch.tensor([[0.0, 1.0]]),
37
+ torch.tensor([[-1.0, 1.0]]),
38
+ torch.tensor([[0.75, 1.0]]),
39
+ torch.tensor([[float("inf"), 1.0]]),
40
+ torch.ones(1, dtype=torch.int32),
41
+ ],
42
+)
43
+def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor):
44
+ assert _wo_a_block_scale_to_e8m0(scale) is None
vllm/models/deepseek_v4/amd/rocm.py
+173
-11
@@ -7,6 +7,7 @@ from typing import cast
7
8
import torch
9
10
+from vllm import envs
11
from vllm.distributed import (
12
get_tensor_model_parallel_world_size,
13
tensor_model_parallel_all_reduce,
@@ -42,6 +43,49 @@ from vllm.v1.worker.workspace import current_workspace_manager
43
logger = init_logger(__name__)
44
45
46
+def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None:
47
+ """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes.
48
+
49
+ E8M0 is an unsigned exponent-only scale format with bias 127. A finite
50
+ encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``;
51
+ ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used
52
+ by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP
53
+ Microscaling Formats (MX) Specification, section 5.4.1:
54
+ https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
55
+
56
+ Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``,
57
+ or decode it to a floating-point power of two. Floating-point inputs are
58
+ accepted only when the original E8M0 byte can be recovered losslessly.
59
+ This function does not quantize or round arbitrary scales.
60
+ """
61
+ if scale.dtype == torch.float8_e8m0fnu:
62
+ # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a
63
+ # numeric conversion instead of preserving the encoded exponent byte.
64
+ return scale.view(torch.uint8).contiguous()
65
+ if scale.dtype == torch.uint8:
66
+ # The checkpoint loader already exposed the E8M0 wire representation.
67
+ return scale.contiguous()
68
+ if not scale.dtype.is_floating_point:
69
+ return None
70
+
71
+ scale_f32 = scale.detach().float()
72
+ if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()):
73
+ return None
74
+
75
+ # With no mantissa, E8M0 can represent only exact powers of two. Rebuild
76
+ # the value before encoding so this adapter never silently quantizes a
77
+ # general floating-point checkpoint scale.
78
+ exponent = torch.round(torch.log2(scale_f32))
79
+ if not torch.equal(torch.exp2(exponent), scale_f32):
80
+ return None
81
+
82
+ encoded = exponent.to(torch.int32) + 127
83
+ # 0xFF is NaN in OCP E8M0, not a finite exponent.
84
+ if int(encoded.min()) < 0 or int(encoded.max()) > 254:
85
+ return None
86
+ return encoded.to(torch.uint8).contiguous()
87
+
88
+
89
def _trust_dsv4_extra_cache_nan_free(
90
kv_cache_dtype: str,
91
has_kv_transfer: bool,
@@ -555,6 +599,10 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
599
# Block scale for the preshuffled weight; None = not preshuffled.
600
self._wqa_wkv_scale: torch.Tensor | None = None
601
self._wo_b_scale: torch.Tensor | None = None
602
+ self._wo_a_fp8_weight: torch.Tensor | None = None
603
+ self._wo_a_e8m0_scale: torch.Tensor | None = None
604
+ self._wo_a_cos_cache: torch.Tensor | None = None
605
+ self._wo_a_sin_cache: torch.Tensor | None = None
606
self._fused_compressor_weight: torch.Tensor | None
607
self.register_buffer("_fused_compressor_weight", None, persistent=False)
608
self._fused_compressor_split_sizes: tuple[int, int] | None = None
@@ -815,6 +863,93 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
863
864
self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
865
self._wo_b_scale = _prep(self.wo_b)
866
+ if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM:
867
+ self._prepare_fp8_wo_a()
868
+
869
+ def _prepare_fp8_wo_a(self) -> None:
870
+ try:
871
+ from aiter.ops.batched_gemm_op_a8w8 import (
872
+ batched_gemm_a8w8_mxscale as mxscale_op,
873
+ )
874
+ from aiter.ops.inverse_rope_group_quant import (
875
+ inverse_rope_group_quant as inverse_quant_op,
876
+ )
877
+ except ImportError:
878
+ logger.warning_once(
879
+ "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; "
880
+ "falling back to BF16 WO_A."
881
+ )
882
+ return
883
+ del mxscale_op, inverse_quant_op
884
+
885
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
886
+ get_fp8_block_weight_scale,
887
+ is_fp8,
888
+ )
889
+
890
+ weight = getattr(self.wo_a, "weight", None)
891
+ scale = get_fp8_block_weight_scale(self.wo_a)
892
+ if scale is None:
893
+ # ModelOpt MXFP8 stores the multiplicative E8M0 scale without the
894
+ # historical ``_inv`` suffix.
895
+ scale = getattr(self.wo_a, "weight_scale", None)
896
+ if weight is None or scale is None:
897
+ logger.warning_once(
898
+ "DeepSeek V4 FP8 WO_A needs a block-scaled FP8 wo_a weight; "
899
+ "the layer exposes no weight/weight scale. Falling back to "
900
+ "BF16 WO_A."
901
+ )
902
+ return
903
+ if weight.dim() != 2 or scale.dim() != 2 or not is_fp8(weight.dtype):
904
+ logger.warning_once(
905
+ "DeepSeek V4 FP8 WO_A needs a 2-D FP8 wo_a weight with a 2-D "
906
+ "block scale, got weight %s%s and scale %s. Falling back to "
907
+ "BF16 WO_A.",
908
+ weight.dtype,
909
+ tuple(weight.shape),
910
+ tuple(scale.shape),
911
+ )
912
+ return
913
+
914
+ groups = self.n_local_groups
915
+ out_per_group = self.o_lora_rank
916
+ out_features, in_features = weight.shape
917
+ if (
918
+ out_features != groups * out_per_group
919
+ or out_per_group % 128 != 0
920
+ or in_features % 128 != 0
921
+ or scale.shape != (out_features // 128, in_features // 128)
922
+ ):
923
+ logger.warning_once(
924
+ "DeepSeek V4 FP8 WO_A needs group-128 blocks for %d groups of "
925
+ "%d outputs, got weight %s and scale %s. Falling back to BF16 "
926
+ "WO_A.",
927
+ groups,
928
+ out_per_group,
929
+ tuple(weight.shape),
930
+ tuple(scale.shape),
931
+ )
932
+ return
933
+
934
+ e8m0_scale = _wo_a_block_scale_to_e8m0(scale)
935
+ if e8m0_scale is None:
936
+ logger.warning_once(
937
+ "DeepSeek V4 FP8 WO_A could not losslessly encode the %s wo_a "
938
+ "block scale as OCP E8M0. Falling back to BF16 WO_A.",
939
+ scale.dtype,
940
+ )
941
+ return
942
+
943
+ self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features)
944
+ self._wo_a_e8m0_scale = e8m0_scale.view(
945
+ groups, out_per_group // 128, in_features // 128
946
+ )
947
+ cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None)
948
+ if cache is None:
949
+ cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16)
950
+ cos_cache, sin_cache = cache.chunk(2, dim=-1)
951
+ self._wo_a_cos_cache = cos_cache.contiguous()
952
+ self._wo_a_sin_cache = sin_cache.contiguous()
953
954
def prepare_compressor_gemm_fusion(self) -> bool:
955
if self._fused_compressor_weight is not None:
@@ -973,17 +1108,44 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
1108
)
1109
1110
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
976
- # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
977
- z = rocm_inv_rope_einsum(
978
- self.rotary_emb,
979
- o,
980
- positions,
981
- self.rope_head_dim,
982
- self.n_local_groups,
983
- self.o_lora_rank,
984
- self.wo_a,
985
- )
986
- zf = z.flatten(1)
1111
+ if self._wo_a_fp8_weight is not None:
1112
+ from aiter.ops.batched_gemm_op_a8w8 import (
1113
+ batched_gemm_a8w8_mxscale,
1114
+ )
1115
+ from aiter.ops.inverse_rope_group_quant import (
1116
+ inverse_rope_group_quant,
1117
+ )
1118
+
1119
+ assert self._wo_a_cos_cache is not None
1120
+ assert self._wo_a_sin_cache is not None
1121
+ o_fp8, o_scale = inverse_rope_group_quant(
1122
+ o.view(o.shape[0], self.n_local_heads, self.head_dim),
1123
+ positions.to(torch.int64),
1124
+ self._wo_a_cos_cache,
1125
+ self._wo_a_sin_cache,
1126
+ num_groups=self.n_local_groups,
1127
+ quant_group_size=128,
1128
+ )
1129
+ assert self._wo_a_e8m0_scale is not None
1130
+ zf = batched_gemm_a8w8_mxscale(
1131
+ o_fp8,
1132
+ self._wo_a_fp8_weight,
1133
+ o_scale,
1134
+ self._wo_a_e8m0_scale,
1135
+ dtype=o.dtype,
1136
+ ).flatten(1)
1137
+ else:
1138
+ # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
1139
+ z = rocm_inv_rope_einsum(
1140
+ self.rotary_emb,
1141
+ o,
1142
+ positions,
1143
+ self.rope_head_dim,
1144
+ self.n_local_groups,
1145
+ self.o_lora_rank,
1146
+ self.wo_a,
1147
+ )
1148
+ zf = z.flatten(1)
1149
if self._wo_b_scale is not None and zf.dim() == 2:
1150
return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
1151
return self.wo_b(zf)