[ROCm][Bugfix] Fix AITER preshuffled FP8 block-scale kernel (#56433)

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

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
         )

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()

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
 

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):

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: