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..ff5ad6f4088d 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,27 @@ logger = init_logger(__name__) +def _wo_a_block_scale_to_e8m0(scale: torch.Tensor) -> torch.Tensor | None: + """Return raw E8M0 exponent bytes for a WO_A block-scale tensor.""" + if scale.dtype == torch.float8_e8m0fnu: + return scale.view(torch.uint8).contiguous() + if scale.dtype == torch.uint8: + 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 + exponent = torch.round(torch.log2(scale_f32)) + if not torch.equal(torch.exp2(exponent), scale_f32): + return None + biased = exponent.to(torch.int32) + 127 + if int(biased.min()) < 0 or int(biased.max()) > 255: + return None + return biased.to(torch.uint8).contiguous() + + def _trust_dsv4_extra_cache_nan_free( kv_cache_dtype: str, has_kv_transfer: bool, @@ -520,6 +542,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 +586,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 +799,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)