From 8a039951e309f3673b1c67fa8b5fe79d2c20dc59 Mon Sep 17 00:00:00 2001 From: Liuyinfeng01 Date: Wed, 2 Sep 2026 05:28:58 +0000 Subject: [PATCH] [ROCm][DSV4][Perf] Use FP8 WO_A output projection Signed-off-by: Liuyinfeng01 --- tests/models/test_deepseek_v4_rocm_wo_a.py | 44 ++++++ vllm/models/deepseek_v4/amd/rocm.py | 152 +++++++++++++++++++-- 2 files changed, 185 insertions(+), 11 deletions(-) create mode 100644 tests/models/test_deepseek_v4_rocm_wo_a.py diff --git a/tests/models/test_deepseek_v4_rocm_wo_a.py b/tests/models/test_deepseek_v4_rocm_wo_a.py new file mode 100644 index 000000000000..69c6ed95307c --- /dev/null +++ b/tests/models/test_deepseek_v4_rocm_wo_a.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.models.deepseek_v4.amd.rocm import _wo_a_block_scale_to_e8m0 + + +def test_wo_a_block_scale_to_e8m0_from_float(): + scale = torch.tensor([[0.5, 1.0, 2.0, 4.0]], dtype=torch.float32) + + encoded = _wo_a_block_scale_to_e8m0(scale) + + assert encoded is not None + torch.testing.assert_close( + encoded, + torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8), + ) + assert encoded.is_contiguous() + + +def test_wo_a_block_scale_to_e8m0_preserves_encoded_scales(): + raw = torch.tensor([[125, 127, 131]], dtype=torch.uint8) + encoded = raw.view(torch.float8_e8m0fnu) + + converted = _wo_a_block_scale_to_e8m0(encoded) + + assert converted is not None + torch.testing.assert_close(converted, raw) + + +@pytest.mark.parametrize( + "scale", + [ + torch.tensor([[0.0, 1.0]]), + torch.tensor([[-1.0, 1.0]]), + torch.tensor([[0.75, 1.0]]), + torch.tensor([[float("inf"), 1.0]]), + torch.ones(1, dtype=torch.int32), + ], +) +def test_wo_a_block_scale_to_e8m0_rejects_invalid_scales(scale: torch.Tensor): + assert _wo_a_block_scale_to_e8m0(scale) is None diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3c8d23c4b4fc..0e4e5d511957 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -7,6 +7,7 @@ import torch +from vllm import envs from vllm.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, @@ -41,6 +42,49 @@ logger = init_logger(__name__) +def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None: + """Normalize checkpoint WO_A scales to raw OCP MX E8M0 bytes. + + E8M0 is an unsigned exponent-only scale format with bias 127. A finite + encoded byte ``b`` in ``[0, 254]`` represents ``2 ** (b - 127)``; + ``0xFF`` is reserved for NaN. This is the vendor-neutral OCP encoding used + by AMD AITER/OPUS, not an NVIDIA-specific convention. Reference: OCP + Microscaling Formats (MX) Specification, section 5.4.1: + https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + + Loaders may preserve the encoded byte as ``float8_e8m0fnu``/``uint8``, + or decode it to a floating-point power of two. Floating-point inputs are + accepted only when the original E8M0 byte can be recovered losslessly. + This function does not quantize or round arbitrary scales. + """ + if scale.dtype == torch.float8_e8m0fnu: + # Reinterpret the native E8M0 storage. ``to(uint8)`` would perform a + # numeric conversion instead of preserving the encoded exponent byte. + return scale.view(torch.uint8).contiguous() + if scale.dtype == torch.uint8: + # The checkpoint loader already exposed the E8M0 wire representation. + return scale.contiguous() + if not scale.dtype.is_floating_point: + return None + + scale_f32 = scale.detach().float() + if not bool(torch.isfinite(scale_f32).all()) or bool((scale_f32 <= 0).any()): + return None + + # With no mantissa, E8M0 can represent only exact powers of two. Rebuild + # the value before encoding so this adapter never silently quantizes a + # general floating-point checkpoint scale. + exponent = torch.round(torch.log2(scale_f32)) + if not torch.equal(torch.exp2(exponent), scale_f32): + return None + + encoded = exponent.to(torch.int32) + 127 + # 0xFF is NaN in OCP E8M0, not a finite exponent. + if int(encoded.min()) < 0 or int(encoded.max()) > 254: + return None + return encoded.to(torch.uint8).contiguous() + + def _trust_dsv4_extra_cache_nan_free( kv_cache_dtype: str, has_kv_transfer: bool, @@ -520,6 +564,10 @@ def __init__(self, *args, **kwargs): # 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._wo_a_fp8_weight: torch.Tensor | None = None + self._wo_a_e8m0_scale: torch.Tensor | None = None + self._wo_a_cos_cache: torch.Tensor | None = None + self._wo_a_sin_cache: 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 @@ -560,6 +608,61 @@ def _prep(linear) -> torch.Tensor | None: self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) self._wo_b_scale = _prep(self.wo_b) + if _ON_GFX950 and envs.VLLM_ROCM_USE_AITER_FP8BMM: + self._prepare_fp8_wo_a() + + def _prepare_fp8_wo_a(self) -> None: + try: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale as mxscale_op, + ) + from aiter.ops.inverse_rope_group_quant import ( + inverse_rope_group_quant as inverse_quant_op, + ) + except ImportError: + logger.warning_once( + "The DeepSeek V4 FP8 WO_A path requires AITER >= 0.1.20; " + "falling back to BF16 WO_A." + ) + return + del mxscale_op, inverse_quant_op + + weight = getattr(self.wo_a, "weight", None) + scale = getattr(self.wo_a, "weight_scale_inv", None) + if ( + weight is None + or scale is None + or weight.dim() != 2 + or scale.dim() != 2 + or weight.dtype not in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + ): + return + + groups = self.n_local_groups + out_per_group = self.o_lora_rank + out_features, in_features = weight.shape + if ( + out_features != groups * out_per_group + or out_per_group % 128 != 0 + or in_features % 128 != 0 + or scale.shape != (out_features // 128, in_features // 128) + ): + return + + e8m0_scale = _wo_a_block_scale_to_e8m0(scale) + if e8m0_scale is None: + return + + self._wo_a_fp8_weight = weight.view(groups, out_per_group, in_features) + self._wo_a_e8m0_scale = e8m0_scale.view( + groups, out_per_group // 128, in_features // 128 + ) + cache = getattr(self.rotary_emb, "cos_sin_cache_bf16", None) + if cache is None: + cache = self.rotary_emb.cos_sin_cache.to(dtype=torch.bfloat16) + cos_cache, sin_cache = cache.chunk(2, dim=-1) + self._wo_a_cos_cache = cos_cache.contiguous() + self._wo_a_sin_cache = sin_cache.contiguous() def prepare_compressor_gemm_fusion(self) -> bool: if self._fused_compressor_weight is not None: @@ -718,17 +821,44 @@ def _split_qkv_and_norm( ) 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( - self.rotary_emb, - o, - positions, - self.rope_head_dim, - self.n_local_groups, - self.o_lora_rank, - self.wo_a, - ) - zf = z.flatten(1) + if self._wo_a_fp8_weight is not None: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale, + ) + from aiter.ops.inverse_rope_group_quant import ( + inverse_rope_group_quant, + ) + + assert self._wo_a_cos_cache is not None + assert self._wo_a_sin_cache is not None + o_fp8, o_scale = inverse_rope_group_quant( + o.view(o.shape[0], self.n_local_heads, self.head_dim), + positions.to(torch.int64), + self._wo_a_cos_cache, + self._wo_a_sin_cache, + num_groups=self.n_local_groups, + quant_group_size=128, + ) + assert self._wo_a_e8m0_scale is not None + zf = batched_gemm_a8w8_mxscale( + o_fp8, + self._wo_a_fp8_weight, + o_scale, + self._wo_a_e8m0_scale, + dtype=o.dtype, + ).flatten(1) + else: + # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b. + z = rocm_inv_rope_einsum( + self.rotary_emb, + o, + positions, + self.rope_head_dim, + self.n_local_groups, + self.o_lora_rank, + self.wo_a, + ) + zf = z.flatten(1) if self._wo_b_scale is not None and zf.dim() == 2: return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True) return self.wo_b(zf)