vLLM/tests/model_executor/kernels/test_b12x_linear.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import importlib
import types
from dataclasses import dataclass

import pytest
import torch

from vllm.model_executor.kernels.linear import (
    _LINEAR_BACKEND_KERNEL_MAP,
    _POSSIBLE_FP8_BLOCK_KERNELS,
    _POSSIBLE_FP8_KERNELS,
    _POSSIBLE_MXFP4_KERNELS,
    _POSSIBLE_MXFP8_KERNELS,
    _POSSIBLE_NVFP4_KERNELS,
    B12xFp8BlockScaledMMKernel,
    B12xMxFp4LinearKernel,
    B12xMxfp8LinearKernel,
    B12xNvFp4LinearKernel,
    B12xNvFp4W4A16LinearKernel,
    B12xTensorFP8ScaledMMLinearKernel,
    FP8ScaledMMLinearLayerConfig,
    Mxfp8LinearLayerConfig,
    init_fp8_linear_kernel,
    init_mxfp4_linear_kernel,
    init_mxfp8_linear_kernel,
    init_nvfp4_linear_kernel,
)
from vllm.model_executor.kernels.linear.nvfp4.marlin import (
    MarlinNvFp4LinearKernel,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    kFp8Dynamic128Sym,
    kFp8Static128BlockSym,
    kFp8StaticTensorSym,
    kMxfp4Dynamic,
)
from vllm.platforms import PlatformEnum


@pytest.mark.parametrize(
    ("kernel_cls", "kernels", "before", "after", "initializer", "kwargs"),
    [
        (
            B12xMxFp4LinearKernel,
            _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA],
            "HummingMxFp4LinearKernel",
            "EmulationMxfp4LinearKernel",
            init_mxfp4_linear_kernel,
            {"activation_quant_key": kMxfp4Dynamic},
        ),
        (
            B12xNvFp4LinearKernel,
            _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA],
            "FbgemmNvFp4LinearKernel",
            "EmulationNvFp4LinearKernel",
            init_nvfp4_linear_kernel,
            {},
        ),
        (
            B12xMxfp8LinearKernel,
            _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA],
            "MarlinMxfp8LinearKernel",
            "EmulationMxfp8LinearKernel",
            init_mxfp8_linear_kernel,
            {"weight_shape": (2048, 2048)},
        ),
        (
            B12xTensorFP8ScaledMMLinearKernel,
            _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA],
            "CutlassFP8ScaledMMLinearKernel",
            "PerTensorTorchFP8ScaledMMLinearKernel",
            init_fp8_linear_kernel,
            {
                "activation_quant_key": kFp8StaticTensorSym,
                "weight_quant_key": kFp8StaticTensorSym,
                "input_dtype": torch.bfloat16,
                "out_dtype": torch.bfloat16,
                "weight_shape": (2048, 2048),
            },
        ),
        (
            B12xFp8BlockScaledMMKernel,
            _POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA],
            "CutlassFp8BlockScaledMMKernel",
            "MarlinFP8ScaledMMLinearKernel",
            init_fp8_linear_kernel,
            {
                "activation_quant_key": kFp8Dynamic128Sym,
                "weight_quant_key": kFp8Static128BlockSym,
                "input_dtype": torch.bfloat16,
                "out_dtype": torch.bfloat16,
                "weight_shape": (2048, 2048),
            },
        ),
    ],
)
def test_b12x_backend_registration_priority_and_selection(
    monkeypatch,
    default_vllm_config,
    kernel_cls,
    kernels,
    before: str,
    after: str,
    initializer,
    kwargs: dict,
) -> None:
    import vllm.model_executor.kernels.linear as linear_mod

    assert kernel_cls in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
    names = [kernel.__name__ for kernel in kernels]
    assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after)

    monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
    monkeypatch.setattr(
        linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"
    )
    monkeypatch.setattr(
        kernel_cls,
        "is_supported",
        classmethod(lambda cls, compute_capability=None: (True, None)),
    )
    monkeypatch.setattr(
        kernel_cls,
        "can_implement",
        classmethod(lambda cls, config: (True, None)),
    )

    assert isinstance(initializer(**kwargs), kernel_cls)


def test_b12x_module_lookup_is_dynamo_safe(monkeypatch) -> None:
    import vllm.utils.b12x as b12x_utils

    module = types.ModuleType("b12x.gemm.blockscaled")
    module.run = lambda x: x + 1  # type: ignore[attr-defined]
    monkeypatch.setitem(
        b12x_utils._B12X_SUBMODULES,
        "b12x.gemm.blockscaled",
        module,
    )

    @torch.compile(backend="eager", fullgraph=True)
    def forward(x: torch.Tensor) -> torch.Tensor:
        blockscaled = b12x_utils.get_b12x_blockscaled()
        assert blockscaled is not None
        return blockscaled.run(x)  # type: ignore[attr-defined]

    x = torch.ones(1)
    torch.testing.assert_close(forward(x), x + 1)


def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
    config = FP8ScaledMMLinearLayerConfig(
        activation_quant_key=kFp8StaticTensorSym,
        weight_quant_key=kFp8StaticTensorSym,
        weight_shape=(64, 128),
        input_dtype=torch.bfloat16,
        out_dtype=torch.bfloat16,
    )

    can_implement, reason = B12xTensorFP8ScaledMMLinearKernel.can_implement(config)

    assert can_implement
    assert reason is None


def test_b12x_block_fp8_checks_runtime_support(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod

    platform = types.SimpleNamespace(
        is_cuda=lambda: True,
        is_device_capability_family=lambda family: family == 120,
    )
    monkeypatch.setattr(b12x_mod, "current_platform", platform)

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_blockscaled",
        lambda: types.SimpleNamespace(is_supported=lambda: False),
    )

    supported, reason = B12xFp8BlockScaledMMKernel.is_supported()

    assert not supported
    assert reason == "B12X regular block-FP8 GEMM is not supported"


def test_b12x_block_fp8_requires_matching_supported_dtypes() -> None:
    def config(input_dtype: torch.dtype, out_dtype: torch.dtype):
        return FP8ScaledMMLinearLayerConfig(
            activation_quant_key=kFp8Dynamic128Sym,
            weight_quant_key=kFp8Static128BlockSym,
            weight_shape=(256, 128),
            input_dtype=input_dtype,
            out_dtype=out_dtype,
        )

    can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement(
        config(torch.float32, torch.float32)
    )
    assert not can_implement
    assert reason == "Supports only bf16/fp16 input dtype"

    can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement(
        config(torch.bfloat16, torch.float16)
    )
    assert not can_implement
    assert reason == "Input and output dtype must match"

    can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement(
        config(torch.float16, torch.float16)
    )
    assert can_implement
    assert reason is None


def test_b12x_block_fp8_requires_aligned_features() -> None:
    def can_implement(weight_shape: tuple[int, int]):
        config = FP8ScaledMMLinearLayerConfig(
            activation_quant_key=kFp8Dynamic128Sym,
            weight_quant_key=kFp8Static128BlockSym,
            weight_shape=weight_shape,
            input_dtype=torch.bfloat16,
            out_dtype=torch.bfloat16,
        )
        return B12xFp8BlockScaledMMKernel.can_implement(config)

    assert can_implement((256, 192)) == (
        False,
        "Input features must be a positive multiple of 128",
    )
    assert can_implement((192, 256)) == (
        False,
        "Output features must be a positive multiple of 128",
    )


def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout(
    monkeypatch,
) -> None:
    import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod

    calls = []
    packed = types.SimpleNamespace(out_features=64)

    def pack(weight: torch.Tensor, output_scale: torch.Tensor):
        calls.append((weight, output_scale))
        return packed

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_tensor_fp8",
        lambda: types.SimpleNamespace(pack_weight=pack),
    )
    layer = torch.nn.Module()
    layer.prefix = "model.layers.0.self_attn.qkv_proj"
    original_weight = (
        torch.randn((128, 64), dtype=torch.float32).clamp(-4, 4).to(torch.float8_e4m3fn)
    )
    layer.weight = torch.nn.Parameter(original_weight, requires_grad=False)
    layer.weight_scale = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False)
    layer.input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False)
    weight_loader = object()
    scale_loader = object()
    layer.weight.weight_loader = weight_loader
    layer.weight_scale.weight_loader = scale_loader
    kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel)
    kernel.config = types.SimpleNamespace(weight_shape=(64, 128))
    kernel.layer_param_names = (
        "weight",
        "weight_scale",
        "input_scale",
        "input_scale_ub",
    )

    kernel.process_weights_after_loading(layer)

    assert layer.b12x_tensor_fp8_packed_weight is packed
    assert layer.b12x_warmup_provider is kernel
    assert len(calls) == 1
    weight, output_scale = calls[0]
    torch.testing.assert_close(weight, original_weight.T.contiguous())
    torch.testing.assert_close(output_scale, torch.tensor([0.125]))
    assert layer.weight.numel() == 0
    assert layer.weight_scale.numel() == 0
    assert layer.weight.weight_loader is weight_loader
    assert layer.weight_scale.weight_loader is scale_loader
    torch.testing.assert_close(layer.input_scale, torch.tensor(0.5))


def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight(
    monkeypatch,
) -> None:
    import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod

    calls = []

    def mm(
        source: torch.Tensor,
        packed_weight,
        *,
        bias: torch.Tensor | None = None,
        out_dtype: torch.dtype,
        expected_m: int,
        stream: object = None,
    ) -> torch.Tensor:
        del stream
        calls.append((source, packed_weight, bias, out_dtype, expected_m))
        return torch.full(
            (source.shape[0], packed_weight.out_features),
            3.0,
            dtype=out_dtype,
        )

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_tensor_fp8",
        lambda: types.SimpleNamespace(mm=mm),
    )
    monkeypatch.setattr(
        b12x_mod,
        "current_stream",
        lambda: types.SimpleNamespace(cuda_stream=object()),
    )
    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)

    layer = torch.nn.Module()
    packed = types.SimpleNamespace(out_features=48)
    layer.b12x_tensor_fp8_packed_weight = packed
    layer.weight = torch.nn.Parameter(
        torch.empty((128, 48), dtype=torch.float8_e4m3fn),
        requires_grad=False,
    )
    layer.weight_scale = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False)
    layer.input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False)
    x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
    x_q = torch.empty((6, 128), dtype=torch.float8_e4m3fn)
    bias = torch.empty((48,), dtype=torch.bfloat16)
    kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel)
    kernel.config = types.SimpleNamespace(out_dtype=torch.bfloat16)
    kernel.layer_param_names = (
        "weight",
        "weight_scale",
        "input_scale",
        "input_scale_ub",
    )
    kernel.quant_fp8 = lambda source, scale, scale_ub: (x_q, scale)

    output = kernel.apply_weights(layer, x, bias)

    assert output.shape == (2, 3, 48)
    assert output.dtype == torch.bfloat16
    assert len(calls) == 1
    source, called_packed, called_bias, out_dtype, expected_m = calls[0]
    assert source.data_ptr() == x_q.data_ptr()
    assert called_packed is packed
    assert called_bias is bias
    assert out_dtype == torch.bfloat16
    assert expected_m == 6


def test_b12x_mxfp8_can_implement_supported_config() -> None:
    can_implement, reason = B12xMxfp8LinearKernel.can_implement(
        Mxfp8LinearLayerConfig(weight_shape=(256, 512))
    )

    assert can_implement
    assert reason is None


def test_b12x_mxfp8_support_check_reports_missing_import(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod

    monkeypatch.setattr(b12x_mod.current_platform, "is_cuda", lambda: True)
    monkeypatch.setattr(
        b12x_mod.current_platform,
        "is_device_capability_family",
        lambda family: family == 120,
    )
    monkeypatch.setattr(b12x_mod, "_import_b12x_mxfp8", lambda: None)

    is_supported, reason = B12xMxfp8LinearKernel.is_supported()

    assert not is_supported
    assert reason == "Install the B12X backend with `pip install vllm[b12x]`"


def test_b12x_mxfp8_support_respects_runtime_probe(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod

    monkeypatch.setattr(b12x_mod.current_platform, "is_cuda", lambda: True)
    monkeypatch.setattr(
        b12x_mod.current_platform,
        "is_device_capability_family",
        lambda family: family == 120,
    )
    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_mxfp8",
        lambda: types.SimpleNamespace(is_supported=lambda: False),
    )

    is_supported, reason = B12xMxfp8LinearKernel.is_supported()

    assert not is_supported
    assert reason == "b12x.gemm.mxfp8_linear is not supported"


def test_b12x_mxfp8_process_weights_packs_modelopt_layout(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod

    calls = []
    packed = types.SimpleNamespace(out_features=48)

    def pack(weight: torch.Tensor, weight_scale: torch.Tensor):
        calls.append((weight, weight_scale))
        return packed

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_mxfp8",
        lambda: types.SimpleNamespace(pack_weight=pack),
    )

    layer = torch.nn.Module()
    layer.prefix = "model.layers.0.self_attn.qkv_proj"
    layer.weight = torch.nn.Parameter(
        torch.empty((48, 128), dtype=torch.float8_e4m3fn),
        requires_grad=False,
    )
    layer.weight_scale = torch.nn.Parameter(
        torch.empty((64, 8), dtype=torch.uint8),
        requires_grad=False,
    )
    weight_loader = object()
    scale_loader = object()
    layer.weight.weight_loader = weight_loader
    layer.weight_scale.weight_loader = scale_loader
    kernel = object.__new__(B12xMxfp8LinearKernel)

    kernel.process_weights_after_loading(layer)

    assert layer.b12x_mxfp8_packed_weight is packed
    assert layer.b12x_warmup_provider is kernel
    assert len(calls) == 1
    weight, weight_scale = calls[0]
    assert weight.shape == (48, 128)
    assert weight_scale.shape == (48, 4)
    assert weight.dtype == torch.float8_e4m3fn
    assert weight_scale.dtype == torch.uint8
    assert layer.weight.numel() == 0
    assert layer.weight_scale.numel() == 0
    assert layer.weight.weight_loader is weight_loader
    assert layer.weight_scale.weight_loader is scale_loader


def test_b12x_mxfp8_reload_reuses_packed_tensor_addresses(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod

    @dataclass(frozen=True)
    class PackedWeight:
        values: torch.Tensor
        scales: torch.Tensor
        out_features: int

    def pack(weight: torch.Tensor, weight_scale: torch.Tensor) -> PackedWeight:
        return PackedWeight(
            values=weight.clone(),
            scales=weight_scale.clone(),
            out_features=int(weight.shape[0]),
        )

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_mxfp8",
        lambda: types.SimpleNamespace(pack_weight=pack),
    )
    layer = torch.nn.Module()
    layer.prefix = "model.layers.0.mlp.down_proj"
    layer.weight = torch.nn.Parameter(
        torch.zeros((48, 128), dtype=torch.float8_e4m3fn),
        requires_grad=False,
    )
    layer.weight_scale = torch.nn.Parameter(
        torch.zeros((48, 4), dtype=torch.uint8),
        requires_grad=False,
    )
    kernel = object.__new__(B12xMxfp8LinearKernel)

    kernel.process_weights_after_loading(layer)
    packed = layer.b12x_mxfp8_packed_weight
    values_ptr = packed.values.data_ptr()
    scales_ptr = packed.scales.data_ptr()

    layer.weight = torch.nn.Parameter(
        torch.ones((48, 128), dtype=torch.float8_e4m3fn),
        requires_grad=False,
    )
    layer.weight_scale = torch.nn.Parameter(
        torch.full((48, 4), 3, dtype=torch.uint8),
        requires_grad=False,
    )
    kernel.process_weights_after_loading(layer)

    assert layer.b12x_mxfp8_packed_weight is packed
    assert packed.values.data_ptr() == values_ptr
    assert packed.scales.data_ptr() == scales_ptr
    torch.testing.assert_close(
        packed.values,
        torch.ones((48, 128), dtype=torch.float8_e4m3fn),
    )
    torch.testing.assert_close(
        packed.scales,
        torch.full((48, 4), 3, dtype=torch.uint8),
    )
    assert layer.weight.numel() == 0
    assert layer.weight_scale.numel() == 0


@pytest.fixture
def _mock_b12x_cuda_fp8_platform(monkeypatch: pytest.MonkeyPatch) -> None:
    import vllm.model_executor.layers.quantization.utils.fp8_utils as fp8_utils

    monkeypatch.setattr(
        fp8_utils,
        "current_platform",
        types.SimpleNamespace(
            is_fp8_fnuz=lambda: False,
            is_rocm=lambda: False,
            fp8_dtype=lambda: torch.float8_e4m3fn,
            is_xpu=lambda: False,
            is_cuda_alike=lambda: True,
        ),
    )


@pytest.mark.usefixtures("_mock_b12x_cuda_fp8_platform")
def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None:
    layer = torch.nn.Module()
    layer.weight = torch.nn.Parameter(
        torch.empty((128, 128), dtype=torch.float8_e4m3fn),
        requires_grad=False,
    )
    layer.weight_scale_inv = torch.nn.Parameter(
        torch.empty((1, 1), dtype=torch.float32),
        requires_grad=False,
    )
    layer.weight_block_size = [128, 128]
    weight_loader = object()
    scale_loader = object()
    layer.weight.weight_loader = weight_loader
    layer.weight_scale_inv.weight_loader = scale_loader
    kernel = object.__new__(B12xFp8BlockScaledMMKernel)

    kernel.process_weights_after_loading(layer)

    assert layer.b12x_warmup_provider is kernel
    assert layer.weight.shape == (128, 128)
    assert layer.weight.dtype == torch.float8_e4m3fn
    assert layer.weight_scale_inv.shape == (1, 1)
    assert layer.weight_scale_inv.dtype == torch.float32
    assert layer.weight.weight_loader is weight_loader
    assert layer.weight_scale_inv.weight_loader is scale_loader


@pytest.mark.parametrize("scale_dtype", [torch.float8_e8m0fnu, torch.uint8])
@pytest.mark.usefixtures("_mock_b12x_cuda_fp8_platform")
def test_b12x_block_fp8_upcasts_e8m0_weight_scales(scale_dtype) -> None:
    layer = torch.nn.Module()
    layer.weight = torch.nn.Parameter(
        torch.empty((128, 128), dtype=torch.float8_e4m3fn),
        requires_grad=False,
    )
    scale_bytes = torch.tensor([[125]], dtype=torch.uint8)
    layer.weight_scale_inv = torch.nn.Parameter(
        scale_bytes.view(scale_dtype),
        requires_grad=False,
    )
    layer.weight_block_size = [128, 128]
    kernel = object.__new__(B12xFp8BlockScaledMMKernel)

    kernel.process_weights_after_loading(layer)

    assert layer.weight_scale_inv.dtype == torch.float32
    torch.testing.assert_close(
        layer.weight_scale_inv,
        torch.tensor([[0.25]], dtype=torch.float32),
    )


def test_b12x_mxfp8_apply_uses_packed_weight(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod

    calls = []

    def mxfp8_linear(
        source: torch.Tensor,
        packed_weight,
        *,
        bias: torch.Tensor | None = None,
        expected_m: int | None = None,
        stream: object = None,
    ) -> torch.Tensor:
        del stream
        calls.append((source, packed_weight, bias, expected_m))
        return source.new_full((source.shape[0], packed_weight.out_features), 3.0)

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_mxfp8",
        lambda: types.SimpleNamespace(mm=mxfp8_linear),
    )

    layer = torch.nn.Module()
    packed = types.SimpleNamespace(out_features=48)
    layer.b12x_mxfp8_packed_weight = packed
    x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
    bias = torch.empty((48,), dtype=torch.bfloat16)
    kernel = object.__new__(B12xMxfp8LinearKernel)

    output = kernel.apply_weights(layer, x, bias)

    assert output.shape == (2, 3, 48)
    assert output.dtype == x.dtype
    assert len(calls) == 1
    source, called_packed, called_bias, expected_m = calls[0]
    assert source.shape == (6, 128)
    assert called_packed is packed
    assert called_bias is bias
    assert expected_m == 6


def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod

    calls = []

    def mm_block_fp8(*args, **kwargs):
        calls.append((args, kwargs))
        return torch.full(
            (args[0].shape[0], args[2].shape[0]),
            13.0,
            dtype=kwargs["out_dtype"],
        )

    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_blockscaled",
        lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8),
    )

    a = torch.empty((6, 128), dtype=torch.float8_e4m3fn)
    weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn)
    a_scale = torch.empty((6, 1), dtype=torch.float32)
    weight_scale = torch.empty((2, 1), dtype=torch.float32)
    kernel = object.__new__(B12xFp8BlockScaledMMKernel)
    kernel.config = types.SimpleNamespace(out_dtype=torch.bfloat16)

    output = kernel.apply_block_scaled_mm(a, weight, a_scale, weight_scale)

    assert output.shape == (6, 256)
    assert output.dtype == torch.bfloat16
    assert len(calls) == 1
    assert calls[0] == (
        (a, a_scale, weight, weight_scale),
        {"out_dtype": torch.bfloat16},
    )
    torch.testing.assert_close(output, torch.full_like(output, 13.0))


def test_b12x_mxfp4_requires_dynamic_activations() -> None:
    config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic)
    can_implement, reason = B12xMxFp4LinearKernel.can_implement(config)

    assert can_implement
    assert reason is None

    config.activation_quant_key = None
    can_implement, reason = B12xMxFp4LinearKernel.can_implement(config)

    assert not can_implement
    assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations"


@pytest.mark.parametrize(
    ("kernel_cls", "module_name", "scale_dtype"),
    [
        (
            B12xMxFp4LinearKernel,
            "vllm.model_executor.kernels.linear.mxfp4.b12x",
            torch.uint8,
        ),
        (
            B12xNvFp4LinearKernel,
            "vllm.model_executor.kernels.linear.nvfp4.b12x",
            torch.float8_e4m3fn,
        ),
    ],
)
def test_b12x_fp4_processes_scale_and_preserves_loader(
    monkeypatch,
    kernel_cls,
    module_name: str,
    scale_dtype: torch.dtype,
) -> None:
    scale = torch.empty((48, 8), dtype=scale_dtype)
    swizzled_scale = torch.empty((128, 8), dtype=scale_dtype)
    intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale)
    monkeypatch.setattr(
        importlib.import_module(module_name),
        "_import_b12x_intrinsics",
        lambda: intrinsics,
    )
    layer = torch.nn.Module()
    layer.prefix = "model.layers.0.mlp.shared_expert.down_proj"
    layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
    weight_loader = object()
    layer.weight_scale.weight_loader = weight_loader
    kernel = object.__new__(kernel_cls)

    kernel.process_weights_after_loading(layer)

    assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr()
    assert layer.weight_scale.weight_loader is weight_loader
    assert layer.b12x_warmup_provider is kernel


def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod
    import vllm.utils.flashinfer as flashinfer_utils

    calls: list[tuple] = []
    x_packed = torch.empty((6, 64), dtype=torch.uint8)
    x_scale_storage = torch.empty((128, 4), dtype=torch.uint8)

    def mm_mxfp4(*args, **kwargs):
        calls.append((args, kwargs))
        return torch.full((6, 48), 3.0, dtype=torch.bfloat16)

    monkeypatch.setattr(
        flashinfer_utils,
        "flashinfer_mxfp4_quantize",
        lambda *args, **kwargs: (x_packed, x_scale_storage),
    )
    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_blockscaled",
        lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4),
    )

    layer = torch.nn.Module()
    layer.output_size_per_partition = 48
    layer.weight = torch.empty((48, 64), dtype=torch.uint8)
    layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8)
    x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
    bias = torch.ones(48, dtype=torch.bfloat16)
    kernel = object.__new__(B12xMxFp4LinearKernel)

    output = kernel.apply_weights(layer, x, bias)

    assert output.shape == (2, 3, 48)
    torch.testing.assert_close(output, torch.full_like(output, 4.0))
    assert len(calls) == 1
    args, kwargs = calls[0]
    assert args == (
        x_packed,
        x_scale_storage,
        layer.weight,
        layer.weight_scale,
    )
    assert kwargs == {"out_dtype": torch.bfloat16}


def test_b12x_nvfp4_can_implement_supported_config() -> None:
    can_implement, reason = B12xNvFp4LinearKernel.can_implement(None)

    assert can_implement
    assert reason is None


def test_b12x_backend_selects_w4a16(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear as linear_mod

    monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
    monkeypatch.setattr(
        linear_mod, "_get_linear_backend", lambda *, quantization: "b12x"
    )
    monkeypatch.setitem(
        linear_mod._POSSIBLE_NVFP4_KERNELS,
        PlatformEnum.CUDA,
        [B12xNvFp4LinearKernel, B12xNvFp4W4A16LinearKernel, MarlinNvFp4LinearKernel],
    )
    monkeypatch.setattr(
        B12xNvFp4W4A16LinearKernel,
        "is_supported",
        classmethod(lambda cls, compute_capability=None: (True, None)),
    )

    kernel = init_nvfp4_linear_kernel(use_a16=True)

    assert isinstance(kernel, B12xNvFp4W4A16LinearKernel)


def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
    import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod

    calls: list[tuple] = []
    quant_calls: list[tuple] = []
    x_packed = torch.empty((6, 64), dtype=torch.uint8)
    x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn)

    def quant(*args, **kwargs):
        quant_calls.append((args, kwargs))
        return x_packed, x_scale_storage

    def mm_nvfp4(*args, **kwargs):
        calls.append((args, kwargs))
        return torch.full((6, 48), 3.0, dtype=torch.bfloat16)

    monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", quant)
    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_blockscaled",
        lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4),
    )

    layer = torch.nn.Module()
    layer.output_size_per_partition = 48
    layer.weight = torch.empty((48, 64), dtype=torch.uint8)
    layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn)
    layer.input_global_scale_inv = torch.tensor(2.0)
    layer.alpha = torch.tensor(0.25)
    x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2]
    bias = torch.ones(48, dtype=torch.bfloat16)
    kernel = object.__new__(B12xNvFp4LinearKernel)

    output = kernel.apply_weights(layer, x, bias)

    assert output.shape == (2, 3, 48)
    torch.testing.assert_close(output, torch.full_like(output, 4.0))
    assert len(quant_calls) == 1
    quant_args, quant_kwargs = quant_calls[0]
    assert quant_args[0].shape == (6, 128)
    assert quant_args[0].data_ptr() == x.data_ptr()
    assert quant_args[1] is layer.input_global_scale_inv
    assert not quant_args[0].is_contiguous()
    assert quant_kwargs == {"is_sf_swizzled_layout": True}
    assert len(calls) == 1
    args, kwargs = calls[0]
    assert args == (
        x_packed,
        x_scale_storage,
        layer.weight,
        layer.weight_scale,
        layer.alpha,
    )
    assert kwargs == {"out_dtype": torch.bfloat16}


def test_b12x_nvfp4_w4a16_uses_weight_only_scale_and_a16_dispatch(monkeypatch):
    import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod

    calls = []
    packed_weight = object()

    def pack_weight(weight, scale, *, recipe, global_scale):
        assert weight is layer.weight
        assert scale is layer.weight_scale
        assert recipe == "nvfp4"
        assert global_scale is layer.weight_global_scale
        return packed_weight

    def mm(source, weight, *, required_mode):
        assert weight is packed_weight
        assert required_mode == "a16"
        assert source.dtype == torch.bfloat16
        assert source.is_contiguous()
        calls.append(source)
        return torch.full((source.shape[0], 48), 3.0, dtype=source.dtype)

    def reject_activation_quantization(*args, **kwargs):
        pytest.fail("W4A16 must not quantize activations")

    monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", reject_activation_quantization)
    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_blockscaled",
        lambda: types.SimpleNamespace(pack_weight=pack_weight, mm=mm),
    )
    monkeypatch.setattr(
        b12x_mod,
        "_import_b12x_intrinsics",
        lambda: types.SimpleNamespace(swizzle_block_scale=lambda scale: scale),
    )
    layer = torch.nn.Module()
    layer.weight = torch.nn.Parameter(
        torch.empty((48, 64), dtype=torch.uint8), requires_grad=False
    )
    layer.weight_scale = torch.nn.Parameter(
        torch.empty((48, 8), dtype=torch.float8_e4m3fn), requires_grad=False
    )
    layer.weight_global_scale = torch.tensor([0.125])
    kernel = object.__new__(B12xNvFp4W4A16LinearKernel)
    kernel.process_weights_after_loading(layer)
    assert layer.b12x_warmup_provider is kernel
    kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16).compile()
    assert [source.shape[0] for source in calls] == [1, 8]

    x = torch.randn((2, 3, 256), dtype=torch.bfloat16)[..., ::2]
    output = kernel.apply_weights(layer, x, torch.ones(48, dtype=torch.bfloat16))
    torch.testing.assert_close(calls[-1], x.reshape(6, 128))
    torch.testing.assert_close(output, torch.full((2, 3, 48), 4.0, dtype=x.dtype))
    with pytest.raises(ValueError, match="requires BF16 input"):
        kernel.apply_weights(layer, x.float())


@pytest.mark.parametrize("m,n,k", [(2, 136, 96), (17, 256, 256)])
@torch.inference_mode()
def test_b12x_nvfp4_w4a16_preserves_bf16_activations_under_graph_replay(
    monkeypatch, default_vllm_config, m, n, k
) -> None:
    import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod

    supported, reason = B12xNvFp4W4A16LinearKernel.is_supported()
    if not supported:
        pytest.skip(reason)

    def reject_activation_quantization(*args, **kwargs):
        pytest.fail("W4A16 must not quantize activations")

    monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", reject_activation_quantization)
    torch.manual_seed(17)
    codes = torch.randint(0, 16, (n, k), device="cuda")
    scales = (2.0 ** torch.randint(-2, 2, (n, k // 16), device="cuda")).to(
        torch.float8_e4m3fn
    )
    layer = torch.nn.Module()
    layer.weight = torch.nn.Parameter(
        (codes[:, ::2] | codes[:, 1::2] << 4).to(torch.uint8), requires_grad=False
    )
    layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
    layer.weight_global_scale = torch.nn.Parameter(
        torch.tensor([0.125], device="cuda"), requires_grad=False
    )
    lut = torch.tensor(
        [0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6],
        device="cuda",
    )
    decoded = lut[codes] * scales.float().repeat_interleave(16, dim=1) * 0.125
    x_storage = torch.randn(1, m, k * 2, device="cuda", dtype=torch.bfloat16)
    x = x_storage[..., ::2]
    bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
    kernel = B12xNvFp4W4A16LinearKernel(None)
    kernel.process_weights_after_loading(layer)
    kernel.get_b12x_warmup_unit(layer, (1, 2, 4, 32), x.dtype).compile()
    run = torch.compile(
        lambda inputs: kernel.apply_weights(layer, inputs, bias),
        fullgraph=True,
        dynamic=True,
    )
    run(x)
    torch.accelerator.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        actual = run(x)
    for _ in range(2):
        x_storage.normal_()
        graph.replay()
        expected = (x.double() @ decoded.double().T).bfloat16() + bias
        torch.testing.assert_close(actual, expected, atol=0.005, rtol=0.01)


def _block_weights(n, k, codec):
    block_size, block_bytes = (
        (32, 34) if codec == "q8_0" else (256, 66 if codec == "iq2_xxs" else 74)
    )
    raw = torch.randint(0, 256, (n, k // block_size, block_bytes), dtype=torch.uint8)
    raw[..., :2] = torch.full(
        (n, k // block_size, 1), 1 / 128, dtype=torch.float16
    ).view(torch.uint8)
    return raw


@pytest.mark.parametrize("codec", ["iq2_xs", "iq2_xxs", "q8_0"])
@pytest.mark.parametrize("profile_first", [False, True])
@torch.inference_mode()
def test_b12x_block_quant_dense_compiled_graph(
    default_vllm_config, codec, profile_first
):
    supported, reason = B12xNvFp4W4A16LinearKernel.is_supported()
    if not supported:
        pytest.skip(reason)
    from b12x.testing.iq2_xs_reference import dequantize_blocks

    from vllm.model_executor.layers.quantization.modelopt_block_quant import (
        ModelOptBlockQuantLinearMethod,
    )

    default_vllm_config.scheduler_config.max_num_batched_tokens = 32
    default_vllm_config.scheduler_config.max_num_scheduled_tokens = 32
    default_vllm_config.compilation_config.cudagraph_capture_sizes = [1, 4]
    method = ModelOptBlockQuantLinearMethod(codec)
    raw = _block_weights(128, 256, codec)
    decoded = dequantize_blocks(raw).bfloat16().cuda()
    layer = torch.nn.Module()
    layer.weight = torch.nn.Parameter(raw.cuda(), requires_grad=False)
    method.process_weights_after_loading(layer)
    assert layer.b12x_block_plan.prepared is None
    if profile_first:
        method.apply(layer, torch.zeros(32, 256, device="cuda", dtype=torch.bfloat16))
    layer.b12x_warmup_provider.get_b12x_warmup_unit(
        layer, (1, 4, 32), torch.bfloat16
    ).compile()
    assert layer.b12x_block_plan.prepared is not None
    run = torch.compile(lambda x: method.apply(layer, x), fullgraph=True, dynamic=True)
    for rows in (1, 4, 17):
        x = torch.randn(rows, 256, device="cuda", dtype=torch.bfloat16)
        run(x)
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            output = run(x)
        for _ in range(2):
            x.normal_()
            allocated = torch.accelerator.memory_allocated()
            graph.replay()
            torch.accelerator.synchronize()
            assert torch.accelerator.memory_allocated() == allocated
            expected = (x.float() @ decoded.float().T).bfloat16()
            torch.testing.assert_close(output, expected, rtol=0.015, atol=0.03125)
        graph.reset()


@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
@pytest.mark.parametrize("profile_first", [False, True])
@torch.inference_mode()
def test_b12x_q8_embedding_graph(default_vllm_config, id_dtype, profile_first):
    supported, reason = B12xNvFp4W4A16LinearKernel.is_supported()
    if not supported:
        pytest.skip(reason)
    from b12x.testing.q8_0_reference import dequantize_blocks

    from vllm.model_executor.layers.quantization.modelopt_block_quant import (
        ModelOptBlockQuantLinearMethod,
    )

    default_vllm_config.scheduler_config.max_num_batched_tokens = 32
    default_vllm_config.scheduler_config.max_num_scheduled_tokens = 32
    default_vllm_config.compilation_config.cudagraph_capture_sizes = [1, 4]
    method = ModelOptBlockQuantLinearMethod("q8_0")
    method.is_embedding = True
    raw = _block_weights(128, 256, "q8_0")
    decoded = dequantize_blocks(raw).bfloat16().cuda()
    layer = torch.nn.Module()
    layer.weight = torch.nn.Parameter(raw.cuda(), requires_grad=False)
    method.process_weights_after_loading(layer)
    assert all(plan.prepared is None for plan in layer.b12x_embedding_plans.values())
    if profile_first:
        method.embedding(layer, torch.zeros(32, device="cuda", dtype=id_dtype))
    layer.b12x_warmup_provider.get_b12x_warmup_unit(
        layer, (1, 4, 32), torch.bfloat16
    ).compile()
    assert all(
        plan.prepared is not None for plan in layer.b12x_embedding_plans.values()
    )
    run = torch.compile(
        lambda x: method.embedding(layer, x), fullgraph=True, dynamic=True
    )
    for count in (7, 9, 70):
        ids = torch.arange(count, device="cuda", dtype=id_dtype)
        torch._dynamo.mark_dynamic(ids, 0)
        torch.testing.assert_close(run(ids), decoded[ids.long()], rtol=0, atol=0)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        output = run(ids)
    ids.add_(20)
    allocated = torch.accelerator.memory_allocated()
    graph.replay()
    torch.accelerator.synchronize()
    assert torch.accelerator.memory_allocated() == allocated
    torch.testing.assert_close(output, decoded[ids.long()], rtol=0, atol=0)
    graph.reset()