@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b4da4d17ae

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

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

Matt committed Sep 11, 2026 at 04:44 UTC b4da4d17ae0c1fb44f7358585f04fdba7e4e1c91

5 files changed +120 -36

vllm/model_executor/kernels/linear/scaled_mm/aiter.py

+49 -7

@@ -2,6 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

5 +from typing import ClassVar

6 +

7 import torch

9 from vllm import _custom_ops as ops

@@ -462,6 +464,12 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

464 class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

465 """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight."""

466

467 + # gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major.

468 + wants_transposed_act_scale: ClassVar[bool] = True

469 +

470 + # process_weights_after_loading shuffles layer.weight to layout (16, 16).

471 + preshuffles_weight: ClassVar[bool] = True

472 +

473 @classmethod

474 def is_supported(

475 cls, compute_capability: int | None = None

@@ -519,9 +527,25 @@ class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

527

528 return True, None

529

530 + @staticmethod

531 + def _reads_weight_directly(layer: torch.nn.Module) -> bool:

532 + """True when something other than apply_weights consumes layer.weight.

533 +

534 + Such a weight must stay in the plain layout: ``is_bmm`` marks a stack

535 + of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj.

536 + Both are stamped after construction, so can_implement cannot see them.

537 + """

538 + return bool(

539 + getattr(layer, "is_bmm", False)

540 + or getattr(layer, "skip_weight_relayout", False)

541 + )

542 +

543 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

544 super().process_weights_after_loading(layer)

545

546 + if self._reads_weight_directly(layer):

547 + return

548 +

549 params = FP8BlockParams.from_layer(layer)

550 if params.weight_scale_inv is not None:

551 ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV

@@ -565,11 +589,25 @@ class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

589 else params.weight_scale_inv

590 )

591

592 + # Left unshuffled above; kv_b_proj still reaches here from MLA's

593 + # prefill-context path.

594 + plain = self._reads_weight_directly(layer)

595 +

596 x_2d = x.view(-1, x.shape[-1])

569 - A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=True)

570 - output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(

571 - A, params.weight, As, Bs, output_dtype=self.config.out_dtype

572 - )

597 + A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain)

598 + if plain:

599 + output = rocm_aiter_ops.gemm_a8w8_blockscale(

600 + A,

601 + params.weight,

602 + As,

603 + Bs,

604 + list(self.weight_group_shape),

605 + output_dtype=self.config.out_dtype,

606 + )

607 + else:

608 + output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(

609 + A, params.weight, As, Bs, output_dtype=self.config.out_dtype

610 + )

611 if bias is not None:

612 output = output + bias

613 return output.view(*x.shape[:-1], params.weight.shape[0])

@@ -581,7 +619,11 @@ class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

619 As: torch.Tensor,

620 Bs: torch.Tensor,

621 ) -> torch.Tensor:

584 - raise NotImplementedError(

585 - "AiterPreshuffledFp8BlockScaledMMKernel overrides apply_weights and "

586 - "does not use apply_block_scaled_mm."

622 + """Block-scaled GEMM for callers that pre-quantize their activations.

623 +

624 + ``As`` must be column-major; a row-major one will not raise, it just

625 + returns wrong numbers for M > 1.

626 + """

627 + return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(

628 + A, B, As, Bs, output_dtype=self.config.out_dtype

629 )

vllm/model_executor/layers/attention/mla_attention.py

+3 -4

@@ -564,10 +564,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):

564 )

565 self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)

566 self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)

567 - # AMX reads kv_b_proj's weight directly and never calls it live; the

568 - # reference CPU MLA backend calls it but isn't perf-critical. Skip

569 - # the packed-kernel dispatch either way.

570 - kv_b_proj._cpu_skip_gemm_dispatch = True

567 + # MLA reads this weight directly to build W_UK/W_UV, so a backend must

568 + # not relayout it at load time.

569 + kv_b_proj.skip_weight_relayout = True

570 self.use_direct_call = not current_platform.opaque_attention_op()

571

572 vllm_config = get_current_vllm_config()

vllm/model_executor/layers/linear.py

+3 -3

@@ -204,9 +204,9 @@ class UnquantizedLinearMethod(LinearMethodBase):

204

205 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

206 if current_platform.is_cpu():

207 - # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not

208 - # perf-critical, so skip packing and use a plain fallback.

209 - if getattr(layer, "_cpu_skip_gemm_dispatch", False):

207 + # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical,

208 + # so skip packing and use a plain fallback.

209 + if getattr(layer, "skip_weight_relayout", False):

210 layer.cpu_linear = torch.nn.functional.linear

211 return

212

vllm/models/deepseek_v4/amd/model.py

+11 -6

@@ -71,7 +71,10 @@ from vllm.model_executor.models.utils import (

71 make_layers,

72 maybe_prefix,

73 )

74 -from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention

74 +from vllm.models.deepseek_v4.amd.rocm import (

75 + DeepseekV4ROCMAiterMLAAttention,

76 + weight_already_preshuffled,

77 +)

78 from vllm.platforms import current_platform

79 from vllm.platforms.rocm import on_gfx950

80 from vllm.sequence import IntermediateTensors

@@ -147,11 +150,13 @@ class DeepseekV4MLP(nn.Module):

150 return

151 if ws.dtype == torch.float8_e8m0fnu:

152 ws = _upcast_e8m0_to_fp32(ws).contiguous()

150 - replace_parameter(

151 - self.gate_up_proj,

152 - "weight",

153 - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),

154 - )

153 + # Skip if the linear's kernel already shuffled it.

154 + if not weight_already_preshuffled(self.gate_up_proj):

155 + replace_parameter(

156 + self.gate_up_proj,

157 + "weight",

158 + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),

159 + )

160 self._gateup_scale = ws

161

162 def forward(self, x):

vllm/models/deepseek_v4/amd/rocm.py

+54 -16

@@ -61,6 +61,16 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:

61 return indptr

64 +def weight_already_preshuffled(linear: torch.nn.Module) -> bool:

65 + """True when the linear's kernel already B-preshuffled ``weight``.

66 +

67 + The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped

68 + for those, since shuffle_weight is a permutation rather than an involution.

69 + """

70 + kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)

71 + return bool(getattr(kernel, "preshuffles_weight", False))

72 +

73 +

74 def apply_pre_quantized_block_scaled_mm(

75 linear: torch.nn.Module,

76 x_fp8: torch.Tensor,

@@ -551,12 +561,13 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

561 return None

562 if ws.dtype == torch.float8_e8m0fnu:

563 ws = _upcast_e8m0_to_fp32(ws).contiguous()

554 - # Shuffle the weight in place (single weight, no unshuffled copy).

555 - replace_parameter(

556 - linear,

557 - "weight",

558 - rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),

559 - )

564 + # Skip if the linear's kernel already shuffled it.

565 + if not weight_already_preshuffled(linear):

566 + replace_parameter(

567 + linear,

568 + "weight",

569 + rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),

570 + )

571 return ws

572

573 self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)

@@ -654,16 +665,26 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

665 return qr_kv, kv_score, indexer_kv_score, indexer_weights

666

667 @functools.cached_property

657 - def _wq_b_uses_aiter_block_scaled(self) -> bool:

658 - """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.

668 + def _wq_b_act_scale_transpose(self) -> bool | None:

669 + """Activation-scale byte order every wq_b consumer agrees on, else None.

670 +

671 + None means "do not take the fused norm+quant path". Otherwise the value

672 + is the ``transpose_scale`` the producer must pass so that the scale it

673 + writes matches what the selected GEMM reads.

674

675 Cached: the linear kernels and the aiter env gates are fixed once

676 the model is built, so this is evaluated at the first forward

677 only.

678

664 - The fused norm+quant path is only valid if the quant and GEMM it

665 - replaces are exactly the aiter ones; otherwise fall back to the

666 - shared path.

679 + Two conditions, both necessary:

680 +

681 + * The fused norm+quant path is only valid if the quant and GEMM it

682 + replaces are exactly the aiter ones; otherwise fall back to the

683 + shared path.

684 + * One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have

685 + different (N, K) and so can resolve to kernels wanting opposite byte

686 + orders. A single producer cannot serve both, so refuse the fast path

687 + when they disagree rather than guessing.

688 """

689 from vllm._aiter_ops import rocm_aiter_ops

690 from vllm.model_executor.kernels.linear.scaled_mm import (

@@ -671,16 +692,31 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

692 )

693

694 if not rocm_aiter_ops.is_linear_fp8_enabled():

674 - return False

695 + return None

696

697 linears = [self.wq_b]

698 if self.indexer is not None:

699 linears.append(self.indexer.wq_b)

700 + layouts: set[bool] = set()

701 for linear in linears:

702 kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)

703 if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):

682 - return False

683 - return True

704 + return None

705 + layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False)))

706 + if len(layouts) != 1:

707 + logger.warning_once(

708 + "DeepSeek-V4 wq_b consumers disagree on activation-scale layout; "

709 + "disabling the fused q/kv norm+quant path.",

710 + scope="global",

711 + )

712 + return None

713 + transpose_scale = layouts.pop()

714 + logger.debug_once(

715 + "DeepSeek-V4 wq_b: emitting %s-major activation scales",

716 + "column" if transpose_scale else "row",

717 + scope="global",

718 + )

719 + return transpose_scale

720

721 def _split_qkv_and_norm(

722 self, qr_kv: torch.Tensor

@@ -697,11 +733,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

733 when the aiter linear path is not active.

734 """

735 qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)

736 + transpose_scale = self._wq_b_act_scale_transpose

737 if not (

738 qr.dim() == 2

739 and qr.shape[0] > 0

740 and self.q_lora_rank % 128 == 0

704 - and self._wq_b_uses_aiter_block_scaled

741 + and transpose_scale is not None

742 ):

743 return super()._split_qkv_and_norm(qr_kv)

744

@@ -715,7 +752,8 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):

752 kv_weight=self.kv_norm.weight.data,

753 kv_epsilon=self.eps,

754 group_size=128,

718 - transpose_scale=False,

755 + # Emit the byte order the wq_b GEMMs read.

756 + transpose_scale=transpose_scale,

757 )

758

759 def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: