From b4da4d17ae0c1fb44f7358585f04fdba7e4e1c91 Mon Sep 17 00:00:00 2001 From: Matt <156021403+mawong-amd@users.noreply.github.com> Date: Fri, 11 Sep 2026 04:44:57 -0700 Subject: [PATCH] [ROCm][Bugfix] Fix AITER preshuffled FP8 block-scale kernel (#56433) Signed-off-by: Matthew Wong --- .../kernels/linear/scaled_mm/aiter.py | 56 +++++++++++++-- .../layers/attention/mla_attention.py | 7 +- vllm/model_executor/layers/linear.py | 6 +- vllm/models/deepseek_v4/amd/model.py | 17 +++-- vllm/models/deepseek_v4/amd/rocm.py | 70 ++++++++++++++----- 5 files changed, 120 insertions(+), 36 deletions(-) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index 45f6d418b9..c9046e89f9 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import ClassVar + import torch from vllm import _custom_ops as ops @@ -462,6 +464,12 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight.""" + # gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major. + wants_transposed_act_scale: ClassVar[bool] = True + + # process_weights_after_loading shuffles layer.weight to layout (16, 16). + preshuffles_weight: ClassVar[bool] = True + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -519,9 +527,25 @@ class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): return True, None + @staticmethod + def _reads_weight_directly(layer: torch.nn.Module) -> bool: + """True when something other than apply_weights consumes layer.weight. + + Such a weight must stay in the plain layout: ``is_bmm`` marks a stack + of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj. + Both are stamped after construction, so can_implement cannot see them. + """ + return bool( + getattr(layer, "is_bmm", False) + or getattr(layer, "skip_weight_relayout", False) + ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: super().process_weights_after_loading(layer) + if self._reads_weight_directly(layer): + return + params = FP8BlockParams.from_layer(layer) if params.weight_scale_inv is not None: ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV @@ -565,11 +589,25 @@ class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): else params.weight_scale_inv ) + # Left unshuffled above; kv_b_proj still reaches here from MLA's + # prefill-context path. + plain = self._reads_weight_directly(layer) + x_2d = x.view(-1, x.shape[-1]) - A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=True) - output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( - A, params.weight, As, Bs, output_dtype=self.config.out_dtype - ) + A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain) + if plain: + output = rocm_aiter_ops.gemm_a8w8_blockscale( + A, + params.weight, + As, + Bs, + list(self.weight_group_shape), + output_dtype=self.config.out_dtype, + ) + else: + output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + A, params.weight, As, Bs, output_dtype=self.config.out_dtype + ) if bias is not None: output = output + bias return output.view(*x.shape[:-1], params.weight.shape[0]) @@ -581,7 +619,11 @@ class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): As: torch.Tensor, Bs: torch.Tensor, ) -> torch.Tensor: - raise NotImplementedError( - "AiterPreshuffledFp8BlockScaledMMKernel overrides apply_weights and " - "does not use apply_block_scaled_mm." + """Block-scaled GEMM for callers that pre-quantize their activations. + + ``As`` must be column-major; a row-major one will not raise, it just + returns wrong numbers for M > 1. + """ + return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + A, B, As, Bs, output_dtype=self.config.out_dtype ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 34df61fc82..25eef7462e 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -564,10 +564,9 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) - # AMX reads kv_b_proj's weight directly and never calls it live; the - # reference CPU MLA backend calls it but isn't perf-critical. Skip - # the packed-kernel dispatch either way. - kv_b_proj._cpu_skip_gemm_dispatch = True + # MLA reads this weight directly to build W_UK/W_UV, so a backend must + # not relayout it at load time. + kv_b_proj.skip_weight_relayout = True self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 2e77fac06c..212dd95375 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -204,9 +204,9 @@ class UnquantizedLinearMethod(LinearMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if current_platform.is_cpu(): - # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not - # perf-critical, so skip packing and use a plain fallback. - if getattr(layer, "_cpu_skip_gemm_dispatch", False): + # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical, + # so skip packing and use a plain fallback. + if getattr(layer, "skip_weight_relayout", False): layer.cpu_linear = torch.nn.functional.linear return diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 903635ffe7..e9ff6cd80f 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -71,7 +71,10 @@ from vllm.model_executor.models.utils import ( make_layers, maybe_prefix, ) -from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention +from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLAAttention, + weight_already_preshuffled, +) from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950 from vllm.sequence import IntermediateTensors @@ -147,11 +150,13 @@ class DeepseekV4MLP(nn.Module): return if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - replace_parameter( - self.gate_up_proj, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Skip if the linear's kernel already shuffled it. + if not weight_already_preshuffled(self.gate_up_proj): + replace_parameter( + self.gate_up_proj, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) self._gateup_scale = ws def forward(self, x): diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 3b1df97dc7..a798053a49 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -61,6 +61,16 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: return indptr +def weight_already_preshuffled(linear: torch.nn.Module) -> bool: + """True when the linear's kernel already B-preshuffled ``weight``. + + The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped + for those, since shuffle_weight is a permutation rather than an involution. + """ + kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) + return bool(getattr(kernel, "preshuffles_weight", False)) + + def apply_pre_quantized_block_scaled_mm( linear: torch.nn.Module, x_fp8: torch.Tensor, @@ -551,12 +561,13 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): return None if ws.dtype == torch.float8_e8m0fnu: ws = _upcast_e8m0_to_fp32(ws).contiguous() - # Shuffle the weight in place (single weight, no unshuffled copy). - replace_parameter( - linear, - "weight", - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), - ) + # Skip if the linear's kernel already shuffled it. + if not weight_already_preshuffled(linear): + replace_parameter( + linear, + "weight", + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)), + ) return ws self._wqa_wkv_scale = _prep(self.fused_wqa_wkv) @@ -654,16 +665,26 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): return qr_kv, kv_score, indexer_kv_score, indexer_weights @functools.cached_property - def _wq_b_uses_aiter_block_scaled(self) -> bool: - """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel. + def _wq_b_act_scale_transpose(self) -> bool | None: + """Activation-scale byte order every wq_b consumer agrees on, else None. + + None means "do not take the fused norm+quant path". Otherwise the value + is the ``transpose_scale`` the producer must pass so that the scale it + writes matches what the selected GEMM reads. Cached: the linear kernels and the aiter env gates are fixed once the model is built, so this is evaluated at the first forward only. - The fused norm+quant path is only valid if the quant and GEMM it - replaces are exactly the aiter ones; otherwise fall back to the - shared path. + Two conditions, both necessary: + + * The fused norm+quant path is only valid if the quant and GEMM it + replaces are exactly the aiter ones; otherwise fall back to the + shared path. + * One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have + different (N, K) and so can resolve to kernels wanting opposite byte + orders. A single producer cannot serve both, so refuse the fast path + when they disagree rather than guessing. """ from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.kernels.linear.scaled_mm import ( @@ -671,16 +692,31 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): ) if not rocm_aiter_ops.is_linear_fp8_enabled(): - return False + return None linears = [self.wq_b] if self.indexer is not None: linears.append(self.indexer.wq_b) + layouts: set[bool] = set() for linear in linears: kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None) if not isinstance(kernel, Fp8BlockScaledMMLinearKernel): - return False - return True + return None + layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False))) + if len(layouts) != 1: + logger.warning_once( + "DeepSeek-V4 wq_b consumers disagree on activation-scale layout; " + "disabling the fused q/kv norm+quant path.", + scope="global", + ) + return None + transpose_scale = layouts.pop() + logger.debug_once( + "DeepSeek-V4 wq_b: emitting %s-major activation scales", + "column" if transpose_scale else "row", + scope="global", + ) + return transpose_scale def _split_qkv_and_norm( self, qr_kv: torch.Tensor @@ -697,11 +733,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): when the aiter linear path is not active. """ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + transpose_scale = self._wq_b_act_scale_transpose if not ( qr.dim() == 2 and qr.shape[0] > 0 and self.q_lora_rank % 128 == 0 - and self._wq_b_uses_aiter_block_scaled + and transpose_scale is not None ): return super()._split_qkv_and_norm(qr_kv) @@ -715,7 +752,8 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): kv_weight=self.kv_norm.weight.data, kv_epsilon=self.eps, group_size=128, - transpose_scale=False, + # Emit the byte order the wq_b GEMMs read. + transpose_scale=transpose_scale, ) def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: