vllm.model_executor.kernels.linear.scaled_mm.aiter ¶

class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
    @classmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        if not current_platform.is_rocm():
            return False, "Requires ROCm."

        if compute_capability is not None and compute_capability < 90:
            return False, "requires compute capability 90 and above."

        try:
            import aiter  # noqa: F401 # deliberately attempt to import aiter
        except Exception:
            return False, "requires `aiter` to be installed."

        if not rocm_aiter_ops.is_linear_enabled():
            return (
                False,
                "requires setting `VLLM_ROCM_USE_AITER=1` "
                "and `VLLM_ROCM_USE_AITER_LINEAR=1`. "
                "`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
            )
        return True, None

    @classmethod
    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
        if not c.input_symmetric:
            return False, "supports symmetric quantization only."
        return True, None

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        `AiterInt8ScaledMMLinearKernel` implements a fused version of
            `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
        where scale_a * a and scale_b * b are implemented using numpy-style
        broadcasting.
        Currently only support per-tensor-per-tensor GEMM
        and per-token-per-channel GEMM through AITER
        w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` also does not support
        ATIER block scaled GEMM and mix-precision GEMM.
        """
        w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)

        # ops.scaled_int8_quant supports both dynamic and static quant:
        # * dynamic, i_s is None and x_s computed from x.
        # * static, i_s is scalar and x_s is i_s.
        symmetric = azp_adj is None
        assert symmetric, (
            "AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
        )
        x_q, x_s, x_zp = ops.scaled_int8_quant(x, i_s, i_zp, symmetric=symmetric)

        assert x_zp is None, (
            "AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
        )
        out_dtype = x.dtype

        assert w_q.shape[0] % 16 == 0 and w_q.shape[1] % 16 == 0
        assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
        assert bias is None or bias.shape[0] == w_q.shape[1] and bias.dtype == out_dtype

        m = x_q.shape[0]  # a
        n = w_q.shape[1]  # b

        per_tensor_scale_a = x_s.numel() == 1
        per_tensor_scale_b = w_s.numel() == 1
        per_token_scale_a = x_s.numel() == m
        per_channel_scale_b = w_s.numel() == n

        # @TODO:
        # Maybe broadcast the per-tensor-scale into per-channel-scale
        # if one of the scale is a per-channel-scale.
        # For now, it only supports:
        # - per-tensor-per-tensor a8w8 scaled GEMM, and
        # - per-token-per-channel a8w8 scaled GEMM
        assert (per_tensor_scale_a and per_tensor_scale_b) or (
            per_token_scale_a and per_channel_scale_b
        ), (
            "Currently only support per-tensor-per-tensor GEMM "
            " and per-token-per-channel GEMM through AITER"
            " w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` "
            "does not support AITER block scaled GEMM."
        )

        # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
        # a to be [M, K]
        # b to be [N, K]
        # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
        return rocm_aiter_ops.w8a8_gemm(x_q, w_q.t(), x_s, w_s, bias, out_dtype)

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
    ) -> tuple[bool, str | None]:
        return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(
            compute_capability
        )

    @classmethod
    def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
        can_implement_base, reason = super().can_implement(config)
        if not can_implement_base:
            return can_implement_base, reason

        act_quant_desc = config.activation_quant_key.scale
        if act_quant_desc.group_shape != GroupShape(1, 128):
            return (
                False,
                (
                    "Supports only dynamic per token group activation "
                    "quantization with group_shape=(1,128)."
                ),
            )

        # bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8
        # 2D weight with N and K divisible by 128. Require a tuned config as the
        # fallback can trigger faults or numerical errors.
        if not rocm_aiter_ops.is_linear_fp8_enabled():
            return (
                False,
                (
                    "requires setting `VLLM_ROCM_USE_AITER=1` "
                    "and `VLLM_ROCM_USE_AITER_LINEAR=1`."
                ),
            )

        n, k = config.weight_shape
        if not (n % 128 == 0 and k % 128 == 0):
            return (
                False,
                (
                    f"requires N and K dimensions divisible by 128, received "
                    f"N={n} and K={k}."
                ),
            )

        if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k):
            return (
                False,
                (
                    f"requires a tuned aiter blockscale bpreshuffle config for "
                    f"N={n} and K={k}."
                ),
            )

        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
        else:
            ws, attr = params.weight_scale, params.WEIGHT_SCALE
        if ws is not None and ws.dtype == torch.float8_e8m0fnu:
            replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous())

        weight = params.weight
        # runtime safety net
        assert (
            weight.dim() == 2
            and weight.dtype == current_platform.fp8_dtype()
            and weight.shape[0] % 128 == 0
            and weight.shape[1] % 128 == 0
        ), (
            "AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight "
            "with N and K divisible by 128."
        )

        shuffled_weight = rocm_aiter_ops.shuffle_weight(
            weight.contiguous(), layout=(16, 16)
        )
        replace_parameter(
            layer,
            params.WEIGHT,
            torch.nn.Parameter(shuffled_weight.data, requires_grad=False),
        )

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        params = self._get_layer_params(layer)
        Bs = (
            params.weight_scale
            if params.weight_scale_inv is None
            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=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])

    def apply_block_scaled_mm(
        self,
        A: torch.Tensor,
        B: torch.Tensor,
        As: torch.Tensor,
        Bs: torch.Tensor,
    ) -> torch.Tensor:
        """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
        )