From 1d8a5c251ffb8d3d5cfeeee115ab3ce966f4d5d1 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Fri, 7 Aug 2026 09:35:40 -0700 Subject: [PATCH 1/5] feat(linear): add B12X FP8 backends Add native B12X kernels for per-tensor, block-scaled, and MXFP8 linear layers on SM120/SM121. Support explicit backend selection and an opt-in FP8 override for automatic dispatch, with weight packing and compile-safe custom ops. Assisted-by: OpenAI Codex Signed-off-by: Luke Alonso --- setup.py | 1 + tests/kernels/quantization/test_block_fp8.py | 55 + .../kernels/test_b12x_mxfp8_linear.py | 1105 +++++++++++++++++ vllm/config/kernel.py | 2 + vllm/envs.py | 5 + vllm/model_executor/kernels/b12x_utils.py | 48 + .../model_executor/kernels/linear/__init__.py | 20 + .../kernels/linear/mxfp8/b12x.py | 332 +++++ .../kernels/linear/scaled_mm/b12x.py | 290 +++++ .../kernels/linear/scaled_mm/b12x_tensor.py | 378 ++++++ 10 files changed, 2236 insertions(+) create mode 100644 tests/model_executor/kernels/test_b12x_mxfp8_linear.py create mode 100644 vllm/model_executor/kernels/b12x_utils.py create mode 100644 vllm/model_executor/kernels/linear/mxfp8/b12x.py create mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x.py create mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py diff --git a/setup.py b/setup.py index 1ddb3ea5d568..3d8406251305 100644 --- a/setup.py +++ b/setup.py @@ -1291,6 +1291,7 @@ def add_vllm_package_data(filename: str) -> None: # only; also needs system GStreamer + libv4l (see docs). "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"], "flashinfer": [], # Kept for backwards compatibility + "b12x": ["b12x>=1.2.2"], # Optional deps for Helion kernel development # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 9b41765de6bf..52f87c403e75 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -14,6 +14,10 @@ ) from tests.kernels.utils import fp8_ulp_distance from vllm.config import VllmConfig +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( + B12xFp8BlockScaledMMKernel, + _run_b12x_fp8_block_scaled_mm, +) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, @@ -353,3 +357,54 @@ def test_w8a8_block_fp8_flashinfer_matmul(M, N, K, block_size, out_dtype, seed): torch.abs(out.to(torch.bfloat16) - ref_out.to(torch.bfloat16)) ) / torch.mean(torch.abs(ref_out.to(torch.bfloat16))) assert rel_diff < 0.001 + + +@pytest.mark.parametrize( + "M,N,K", + [(1, 128, 256), (8, 256, 512), (129, 256, 256), (2, 4096, 4096)], +) +@torch.inference_mode() +def test_w8a8_block_fp8_b12x_matmul(M, N, K): + supported, reason = B12xFp8BlockScaledMMKernel.is_supported() + if not supported: + pytest.skip(reason) + + torch.manual_seed(M) + fp8_max = torch.finfo(torch.float8_e4m3fn).max + A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max + B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max + A_fp8, As = per_token_group_quant_fp8(A_bf16, 128, use_ue8m0=False) + B_fp8, Bs = per_block_cast_to_fp8( + B_bf16, + block_size=[128, 128], + use_ue8m0=False, + ) + As = As.float() + Bs = Bs.float() + + ref_out = native_w8a8_block_matmul( + A_fp8, + B_fp8, + As, + Bs, + [128, 128], + torch.bfloat16, + ) + out = _run_b12x_fp8_block_scaled_mm( + A_fp8, + B_fp8, + As, + Bs, + torch.bfloat16, + ) + + rel_diff = torch.mean(torch.abs(out.float() - ref_out.float())) / torch.mean( + torch.abs(ref_out.float()) + ) + cosine = torch.nn.functional.cosine_similarity( + out.float().flatten(), + ref_out.float().flatten(), + dim=0, + ) + assert rel_diff < 0.002 + assert cosine >= 0.9999 diff --git a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py b/tests/model_executor/kernels/test_b12x_mxfp8_linear.py new file mode 100644 index 000000000000..b4c2ab19e14f --- /dev/null +++ b/tests/model_executor/kernels/test_b12x_mxfp8_linear.py @@ -0,0 +1,1105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import types +from dataclasses import dataclass + +import pytest +import torch + +from vllm.config import VllmConfig +from vllm.forward_context import set_forward_context +from vllm.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + _POSSIBLE_FP8_KERNELS, + _POSSIBLE_MXFP8_KERNELS, + init_fp8_linear_kernel, + init_mxfp8_linear_kernel, +) +from vllm.model_executor.kernels.linear.mxfp8.b12x import ( + B12xMxfp8LinearKernel, + _b12x_mxfp8_expected_m, + _b12x_mxfp8_linear, + _b12x_mxfp8_warmup_token_counts, + warmup_b12x_mxfp8_linear, +) +from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( + B12xFp8BlockScaledMMKernel, + _b12x_block_fp8_warmup_token_counts, + _run_b12x_fp8_block_scaled_mm, + warmup_b12x_block_fp8_linear, +) +from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( + B12xTensorFP8ScaledMMLinearKernel, + _b12x_tensor_fp8_linear, + warmup_b12x_tensor_fp8_linear, +) +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( + FP8ScaledMMLinearLayerConfig, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, +) +from vllm.platforms import PlatformEnum + + +def test_b12x_backend_maps_mxfp8_kernel() -> None: + assert B12xMxfp8LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + assert B12xMxfp8LinearKernel in _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA] + + +def test_b12x_backend_maps_tensor_fp8_kernel() -> None: + assert B12xTensorFP8ScaledMMLinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + assert B12xTensorFP8ScaledMMLinearKernel in _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA] + + +def test_b12x_fp8_env_selects_per_tensor_fp8_with_auto_backend( + monkeypatch, + default_vllm_config, +) -> 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: "auto") + monkeypatch.setattr(linear_mod.envs, "VLLM_USE_B12X_FP8_GEMM", True) + monkeypatch.setattr( + B12xTensorFP8ScaledMMLinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + monkeypatch.setattr( + B12xTensorFP8ScaledMMLinearKernel, + "can_implement", + classmethod(lambda cls, config: (True, None)), + ) + + kernel = init_fp8_linear_kernel( + activation_quant_key=kFp8StaticTensorSym, + weight_quant_key=kFp8StaticTensorSym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + weight_shape=(2048, 2048), + ) + + assert isinstance(kernel, B12xTensorFP8ScaledMMLinearKernel) + + +def test_b12x_tensor_fp8_env_enables_auto_backend(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "auto") + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP8_GEMM", False) + 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 not can_implement + assert reason == "b12x tensor FP8 GEMM is not enabled" + + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP8_GEMM", True) + can_implement, reason = B12xTensorFP8ScaledMMLinearKernel.can_implement(config) + + assert can_implement + assert reason is None + + +def test_b12x_fp8_env_selects_block_fp8_with_auto_backend( + monkeypatch, + default_vllm_config, +) -> 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: "auto") + monkeypatch.setattr(linear_mod.envs, "VLLM_USE_B12X_FP8_GEMM", True) + monkeypatch.setattr( + B12xFp8BlockScaledMMKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_fp8_linear_kernel( + activation_quant_key=kFp8Dynamic128Sym, + weight_quant_key=kFp8Static128BlockSym, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + weight_shape=(2048, 2048), + ) + + assert isinstance(kernel, B12xFp8BlockScaledMMKernel) + + +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(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "b12x") + + 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(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "b12x") + + 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_block_fp8_warmup_covers_serving_sizes() -> None: + assert _b12x_block_fp8_warmup_token_counts( + max_tokens=2048, + cudagraph_capture_sizes=[1, 2, 8, 128], + ) == (1, 2, 8, 128, 2048) + + +def test_warmup_b12x_block_fp8_dedupes_weight_signatures(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + + def run(a, weight, a_scale, weight_scale, out_dtype): + calls.append((a.shape, weight, a_scale.shape, weight_scale, out_dtype)) + return torch.empty((a.shape[0], weight.shape[0]), dtype=out_dtype) + + 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(), + ) + monkeypatch.setattr(b12x_mod, "_run_b12x_fp8_block_scaled_mm", run) + + def layer(in_features: int, out_features: int): + return types.SimpleNamespace( + b12x_block_fp8_linear=True, + weight=torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn), + weight_scale_inv=torch.empty( + (out_features // 128, in_features // 128), dtype=torch.float32 + ), + ) + + layer_a = layer(128, 256) + layer_b = layer(256, 128) + modules = [ + layer_a, + layer_a, + layer_b, + types.SimpleNamespace(), + ] + model = types.SimpleNamespace(modules=lambda: iter(modules)) + + warmed = warmup_b12x_block_fp8_linear( + model, + max_tokens=32, + cudagraph_capture_sizes=[2, 8], + output_dtype=torch.bfloat16, + ) + + assert warmed == 8 + assert [call[0][0] for call in calls] == [1, 2, 8, 32] * 2 + assert [call[2][0] for call in calls] == [1, 2, 8, 32] * 2 + assert calls[0][1] is layer_a.weight + assert calls[4][1] is layer_b.weight + assert calls[0][3] is layer_a.weight_scale_inv + assert calls[4][3] is layer_b.weight_scale_inv + assert all(call[4] == torch.bfloat16 for call in calls) + + +def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout( + monkeypatch, +) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor 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), + ) + vllm_config = VllmConfig() + monkeypatch.setattr( + b12x_mod, "get_current_vllm_config_or_none", lambda: vllm_config + ) + + 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 vllm_config.compilation_config.static_forward_context[layer.prefix] is layer + 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_warmup_b12x_tensor_fp8_dedupes_weight_signatures(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + + calls = [] + + def prewarm(packed_weight, token_counts, *, out_dtype, stream): + del stream + calls.append((packed_weight, tuple(token_counts), out_dtype)) + return len(tuple(token_counts)) + + 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, "_b12x_tensor_fp8_enabled", lambda: True) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_tensor_fp8", + lambda: types.SimpleNamespace(prewarm=prewarm), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: types.SimpleNamespace(cuda_stream=object()), + ) + + def packed(in_features: int, padded_in_features: int, out_features: int): + return types.SimpleNamespace( + in_features=in_features, + padded_in_features=padded_in_features, + out_features=out_features, + values=torch.empty(1), + ) + + packed_a = packed(128, 128, 256) + packed_b = packed(160, 256, 512) + modules = [ + types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a), + types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a), + types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_b), + types.SimpleNamespace(), + ] + model = types.SimpleNamespace(modules=lambda: iter(modules)) + + warmed = warmup_b12x_tensor_fp8_linear( + model, + max_tokens=2048, + cudagraph_capture_sizes=[1, 2], + ) + + assert warmed == 6 + assert calls == [ + (packed_a, (1, 2, 2048), torch.bfloat16), + (packed_b, (1, 2, 2048), torch.bfloat16), + ] + + +def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight( + monkeypatch, +) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor 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, + ) -> 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_tensor_fp8_custom_op_body_uses_forward_context(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + + calls = [] + + def mm( + source: torch.Tensor, + packed_weight, + *, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + expected_m: int, + stream: object, + ) -> torch.Tensor: + del stream + calls.append((source, packed_weight, bias, out_dtype, expected_m)) + return torch.full( + (source.shape[0], packed_weight.out_features), + 5.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()), + ) + + layer = torch.nn.Module() + layer.prefix = "model.layers.2.self_attn.qkv_proj" + packed = types.SimpleNamespace(out_features=32) + layer.b12x_tensor_fp8_packed_weight = packed + x_q = torch.empty((2, 3, 128), dtype=torch.float8_e4m3fn) + bias = torch.empty((32,), dtype=torch.bfloat16) + vllm_config = VllmConfig() + vllm_config.compilation_config.static_forward_context[layer.prefix] = layer + + with set_forward_context({}, vllm_config): + output = _b12x_tensor_fp8_linear( + x_q, + bias, + layer.prefix, + 32, + torch.bfloat16, + ) + + assert output.shape == (2, 3, 32) + assert len(calls) == 1 + source, called_packed, called_bias, out_dtype, expected_m = calls[0] + assert source.shape == (6, 128) + assert called_packed is packed + assert called_bias is bias + assert out_dtype == torch.bfloat16 + assert expected_m == 6 + torch.testing.assert_close(output, torch.full_like(output, 5.0)) + + +def test_b12x_mxfp8_explicit_backend_selects_kernel(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: "b12x") + monkeypatch.setattr( + B12xMxfp8LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + monkeypatch.setattr( + B12xMxfp8LinearKernel, + "can_implement", + classmethod(lambda cls, c: (True, None)), + ) + + kernel = init_mxfp8_linear_kernel() + + assert isinstance(kernel, B12xMxfp8LinearKernel) + + +def test_b12x_mxfp8_can_implement_requires_opt_in(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "auto") + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP8_GEMM", False) + + can_implement, reason = B12xMxfp8LinearKernel.can_implement( + Mxfp8LinearLayerConfig() + ) + + assert not can_implement + assert reason == "b12x MXFP8 GEMM is not enabled" + + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP8_GEMM", True) + can_implement, reason = B12xMxfp8LinearKernel.can_implement( + Mxfp8LinearLayerConfig() + ) + + assert can_implement + assert reason is None + + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP8_GEMM", False) + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "b12x") + can_implement, reason = B12xMxfp8LinearKernel.can_implement( + Mxfp8LinearLayerConfig() + ) + + assert can_implement + assert reason is None + + +def test_b12x_mxfp8_expected_m_uses_live_m() -> None: + assert _b12x_mxfp8_expected_m(0) == 1 + assert _b12x_mxfp8_expected_m(1) == 1 + assert _b12x_mxfp8_expected_m(2) == 2 + assert _b12x_mxfp8_expected_m(8) == 8 + assert _b12x_mxfp8_expected_m(9) == 9 + assert _b12x_mxfp8_expected_m(128) == 128 + assert _b12x_mxfp8_expected_m(129) == 129 + assert _b12x_mxfp8_expected_m(2048) == 2048 + + +def test_b12x_mxfp8_warmup_token_counts_cover_serving_regimes() -> None: + assert _b12x_mxfp8_warmup_token_counts( + max_tokens=2048, + cudagraph_capture_sizes=[1, 2, 4, 8], + ) == (1, 2, 4, 8, 2048) + + +def test_warmup_b12x_mxfp8_linear_dedupes_weight_signatures( + monkeypatch, +) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + calls = [] + + def mm( + 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.shape, packed_weight, bias, expected_m)) + return source.new_empty((source.shape[0], packed_weight.out_features)) + + 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, "_b12x_mxfp8_enabled", lambda: True) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: types.SimpleNamespace(mm=mm), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: types.SimpleNamespace(cuda_stream=object()), + ) + + def packed(in_features: int, padded_in_features: int, out_features: int): + return types.SimpleNamespace( + in_features=in_features, + padded_in_features=padded_in_features, + out_features=out_features, + weight=types.SimpleNamespace(values=torch.empty(1)), + ) + + packed_a = packed(128, 128, 256) + packed_b = packed(128, 128, 512) + modules = [ + types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a), + types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a), + types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_b), + types.SimpleNamespace(), + ] + model = types.SimpleNamespace(modules=lambda: iter(modules)) + + warmed = warmup_b12x_mxfp8_linear( + model, + max_tokens=2048, + cudagraph_capture_sizes=[1, 2], + ) + + assert warmed == 6 + assert [call[0] for call in calls] == [ + torch.Size([1, 128]), + torch.Size([2, 128]), + torch.Size([2048, 128]), + torch.Size([1, 128]), + torch.Size([2, 128]), + torch.Size([2048, 128]), + ] + assert [call[3] for call in calls] == [1, 2, 2048, 1, 2, 2048] + assert calls[0][1] is packed_a + assert calls[3][1] is packed_b + + +def test_b12x_mxfp8_disabled_support_check_skips_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, "_current_linear_backend", lambda: "auto") + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP8_GEMM", False) + + def fail_import(): + raise AssertionError("B12X MXFP8 import should require opt-in") + + monkeypatch.setattr(b12x_mod, "_import_b12x_mxfp8", fail_import) + + is_supported, reason = B12xMxfp8LinearKernel.is_supported() + + assert not is_supported + assert reason == "b12x MXFP8 GEMM is not enabled" + + +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, "_current_linear_backend", lambda: "b12x") + 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) + vllm_config = VllmConfig() + + monkeypatch.setattr( + b12x_mod, "get_current_vllm_config_or_none", lambda: vllm_config + ) + + kernel.process_weights_after_loading(layer) + + assert layer.b12x_mxfp8_packed_weight is packed + assert vllm_config.compilation_config.static_forward_context[layer.prefix] is layer + 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), + ) + monkeypatch.setattr( + b12x_mod, + "get_current_vllm_config_or_none", + lambda: VllmConfig(), + ) + + 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 + + +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_block_fp8_linear + 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]) +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_mxfp8_compile_path_uses_forward_context_custom_op( + monkeypatch, +) -> None: + calls = [] + + def op( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: str, + out_features: int, + ) -> torch.Tensor: + calls.append((x, bias, layer_name, out_features)) + return x.new_full((*x.shape[:-1], out_features), 7.0) + + monkeypatch.setattr(torch.ops.vllm, "b12x_mxfp8_linear", op) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + layer = torch.nn.Module() + layer.prefix = "model.layers.1.mlp.gate_up_proj" + packed = types.SimpleNamespace(out_features=32) + layer.b12x_mxfp8_packed_weight = packed + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + bias = torch.empty((32,), dtype=torch.bfloat16) + kernel = object.__new__(B12xMxfp8LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 32) + assert output.dtype == x.dtype + assert len(calls) == 1 + source, called_bias, layer_name, out_features = calls[0] + assert source is x + assert called_bias is bias + assert getattr(layer_name, "value", layer_name) == layer.prefix + assert out_features == 32 + torch.testing.assert_close(output, torch.full_like(output, 7.0)) + + +def test_b12x_mxfp8_custom_op_body_uses_forward_context(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), 11.0) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: types.SimpleNamespace(mm=mxfp8_linear), + ) + + layer = torch.nn.Module() + layer.prefix = "model.layers.2.mlp.down_proj" + packed = types.SimpleNamespace(out_features=16) + layer.b12x_mxfp8_packed_weight = packed + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + bias = torch.empty((16,), dtype=torch.bfloat16) + vllm_config = VllmConfig() + vllm_config.compilation_config.static_forward_context[layer.prefix] = layer + + with set_forward_context({}, vllm_config): + output = _b12x_mxfp8_linear(x, bias, layer.prefix, 16) + + assert output.shape == (2, 3, 16) + 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 + torch.testing.assert_close(output, torch.full_like(output, 11.0)) + + +def test_b12x_block_fp8_apply_uses_opaque_prequantized_op(monkeypatch) -> None: + calls = [] + + def op( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, + ) -> torch.Tensor: + calls.append((a, weight, a_scale, weight_scale, out_dtype)) + return torch.full((a.shape[0], weight.shape[0]), 13.0, dtype=out_dtype) + + monkeypatch.setattr(torch.ops.vllm, "b12x_fp8_block_scaled_mm", op) + + 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, weight, a_scale, weight_scale, torch.bfloat16) + torch.testing.assert_close(output, torch.full_like(output, 13.0)) + + +def test_b12x_block_fp8_op_uses_regular_compact_scale_api(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + + def mm(lhs, rhs, **kwargs): + calls.append((lhs, rhs, kwargs)) + return torch.full((6, 256, 1), 17.0, dtype=torch.bfloat16) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm=mm), + ) + stream = object() + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: types.SimpleNamespace(cuda_stream=stream), + ) + + 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) + output = _run_b12x_fp8_block_scaled_mm( + a, + weight, + a_scale, + weight_scale, + torch.bfloat16, + ) + + assert output.shape == (6, 256) + lhs, rhs, kwargs = calls[0] + assert lhs[0].shape == (6, 128, 1) + assert lhs[1] is a_scale + assert rhs[0].shape == (256, 128, 1) + assert rhs[1] is weight_scale + assert kwargs == { + "ab_dtype": "float8_e4m3fn", + "sf_dtype": "float32", + "c_dtype": "bfloat16", + "sf_vec_size": 128, + "block_fp8": True, + "expected_m": 6, + "stream": stream, + } + torch.testing.assert_close(output, torch.full_like(output, 17.0)) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 59195f307b80..61bbfe33e8eb 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -141,6 +141,7 @@ def with_default( LinearBackend = Literal[ "auto", + "b12x", "cutlass", "flashinfer_cutlass", "flashinfer_cutedsl", @@ -192,6 +193,7 @@ class KernelConfig: """Backend for MoE expert computation kernels. Available options: - "auto": Automatically select the best backend based on model and hardware + - "b12x": Use B12X kernels for SM12x FP8 and FP4 linear layers - "triton": Use Triton-based fused MoE kernels - "batched_triton": Use batched Triton experts (moe_mmk) on the batched activation format ([E_local, max_num_tokens, K]) diff --git a/vllm/envs.py b/vllm/envs.py index 1460fa1a8946..29c35b339a8d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -182,6 +182,7 @@ VLLM_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None VLLM_HUMMING_USE_F16_ACCUM: bool = False VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None + VLLM_USE_B12X_FP8_GEMM: bool = False VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False VLLM_TPU_USING_PATHWAYS: bool = False @@ -1568,6 +1569,10 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) ), + # Prefer B12X dense GEMMs for supported FP8 linear layers. + "VLLM_USE_B12X_FP8_GEMM": lambda: bool( + int(os.getenv("VLLM_USE_B12X_FP8_GEMM", "0")) + ), # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool( int(os.getenv("VLLM_USE_FLASHINFER_MOE_INT4", "0")) diff --git a/vllm/model_executor/kernels/b12x_utils.py b/vllm/model_executor/kernels/b12x_utils.py new file mode 100644 index 000000000000..47116d8c9333 --- /dev/null +++ b/vllm/model_executor/kernels/b12x_utils.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import fields, is_dataclass +from typing import Any + +import torch + + +def _same_packed_layout(current: Any, replacement: Any) -> bool: + if type(current) is not type(replacement): + return False + if isinstance(current, torch.Tensor): + return ( + current.shape == replacement.shape + and current.stride() == replacement.stride() + and current.dtype == replacement.dtype + and current.device == replacement.device + ) + if is_dataclass(current): + return all( + _same_packed_layout( + getattr(current, field.name), + getattr(replacement, field.name), + ) + for field in fields(current) + ) + return bool(current == replacement) + + +def _copy_packed_tensors(current: Any, replacement: Any) -> None: + if isinstance(current, torch.Tensor): + current.copy_(replacement) + elif is_dataclass(current): + for field in fields(current): + _copy_packed_tensors( + getattr(current, field.name), + getattr(replacement, field.name), + ) + + +@torch.no_grad() +def reuse_packed_weight_storage(current: Any, replacement: Any) -> Any: + """Reuse packed tensor addresses when a compatible weight is reloaded.""" + if current is None or not _same_packed_layout(current, replacement): + return replacement + _copy_packed_tensors(current, replacement) + return current diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 90e983653eb2..0b582989f616 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -103,6 +103,9 @@ Mxfp8LinearKernel, Mxfp8LinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mxfp8.b12x import ( + B12xMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp8.emulation import ( EmulationMxfp8LinearKernel, ) @@ -163,6 +166,12 @@ AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( + B12xFp8BlockScaledMMKernel, +) +from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( + B12xTensorFP8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, CPUInt8ScaledMMLinearKernel, @@ -228,6 +237,11 @@ def _get_linear_backend() -> str: # set are considered candidates. If none can implement the layer config, # an error is raised to respect the user's explicit intent. _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { + "b12x": { + B12xFp8BlockScaledMMKernel, + B12xMxfp8LinearKernel, + B12xTensorFP8ScaledMMLinearKernel, + }, "cutlass": { CutlassInt8ScaledMMLinearKernel, CutlassFP8ScaledMMLinearKernel, @@ -374,6 +388,7 @@ def _resolve_backend_kernels( # in priority/performance order (when available) _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { PlatformEnum.CUDA: [ + B12xTensorFP8ScaledMMLinearKernel, MarlinFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, CutlassFP8ScaledMMLinearKernel, @@ -408,6 +423,7 @@ def _resolve_backend_kernels( PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]] ] = { PlatformEnum.CUDA: [ + B12xFp8BlockScaledMMKernel, FlashInferFp8DeepGEMMDynamicBlockScaledKernel, DeepGemmFp8BlockScaledMMKernel, CutlassFp8BlockScaledMMKernel, @@ -478,6 +494,7 @@ def _resolve_backend_kernels( # in priority/performance order (when available) _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = { PlatformEnum.CUDA: [ + B12xMxfp8LinearKernel, FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, MarlinMxfp8LinearKernel, @@ -1174,6 +1191,7 @@ def register_linear_kernel( "init_mxfp8_linear_kernel", "Mxfp8LinearKernel", "Mxfp8LinearLayerConfig", + "B12xMxfp8LinearKernel", "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", @@ -1202,4 +1220,6 @@ def register_linear_kernel( "_KernelT", "DeepGemmFp8BlockScaledMMKernel", "FlashInferFp8DeepGEMMDynamicBlockScaledKernel", + "B12xFp8BlockScaledMMKernel", + "B12xTensorFP8ScaledMMLinearKernel", ] diff --git a/vllm/model_executor/kernels/linear/mxfp8/b12x.py b/vllm/model_executor/kernels/linear/mxfp8/b12x.py new file mode 100644 index 000000000000..ad82bc0d28eb --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/b12x.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config_or_none +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.b12x_utils import ( + reuse_packed_weight_storage, +) +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + _USE_LAYERNAME, + LayerName, + _encode_layer_name, + current_stream, + direct_register_custom_op, +) + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | LayerName +else: + _layer_name_type = LayerName if _USE_LAYERNAME else str + +_B12X_MXFP8: Any | None = None +_B12X_MXFP8_MISSING = False + + +def _import_b12x_mxfp8() -> Any | None: + global _B12X_MXFP8, _B12X_MXFP8_MISSING + if _B12X_MXFP8 is not None: + return _B12X_MXFP8 + if _B12X_MXFP8_MISSING: + return None + try: + _B12X_MXFP8 = importlib.import_module("b12x.gemm.mxfp8_linear") + except ImportError: + _B12X_MXFP8_MISSING = True + return None + return _B12X_MXFP8 + + +def _current_linear_backend() -> str: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return "auto" + return str(getattr(vllm_config.kernel_config, "linear_backend", "auto")).lower() + + +def _b12x_mxfp8_enabled() -> bool: + return _current_linear_backend() == "b12x" or envs.VLLM_USE_B12X_FP8_GEMM + + +def _b12x_mxfp8_expected_m(tokens: int) -> int: + return max(1, int(tokens)) + + +def _b12x_mxfp8_warmup_token_counts( + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), +) -> tuple[int, ...]: + counts = {1} + counts.update(int(size) for size in cudagraph_capture_sizes if int(size) > 0) + if int(max_tokens) > 0: + counts.add(int(max_tokens)) + return tuple(sorted(counts)) + + +@torch.compiler.assume_constant_result +def _resolve_layer_name(layer_name: str | LayerName) -> str: + from torch._library.fake_class_registry import FakeScriptObject + + if isinstance(layer_name, LayerName): + return layer_name.value + elif isinstance(layer_name, FakeScriptObject): + return layer_name.real_obj.value + return layer_name + + +def _register_b12x_mxfp8_linear_layer(layer: torch.nn.Module) -> None: + prefix = getattr(layer, "prefix", "") + if not prefix: + return + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return + static_forward_context = vllm_config.compilation_config.static_forward_context + existing = static_forward_context.get(prefix) + if existing is not None and existing is not layer: + raise ValueError(f"Duplicate B12X MXFP8 linear layer name: {prefix}") + static_forward_context[prefix] = layer + + +def _apply_b12x_mxfp8_packed_linear( + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) + if packed_weight is None: + raise RuntimeError( + "b12x MXFP8 packed weights are missing; " + "process_weights_after_loading did not run for this layer" + ) + + input_2d = x.reshape(-1, x.shape[-1]).contiguous() + output_shape = [*x.shape[:-1], int(packed_weight.out_features)] + + mxfp8 = _import_b12x_mxfp8() + if mxfp8 is None: + raise ImportError("b12x.gemm.mxfp8_linear is not importable") + output = mxfp8.mm( + input_2d, + packed_weight, + bias=bias, + expected_m=_b12x_mxfp8_expected_m(int(input_2d.shape[0])), + stream=current_stream().cuda_stream, + ) + return output.view(*output_shape) + + +def warmup_b12x_mxfp8_linear( + model: torch.nn.Module, + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), + output_dtype: torch.dtype = torch.bfloat16, +) -> int: + if not _b12x_mxfp8_enabled(): + return 0 + if not current_platform.is_cuda(): + return 0 + if not current_platform.is_device_capability_family(120): + return 0 + if output_dtype not in (torch.bfloat16, torch.float16): + output_dtype = torch.bfloat16 + + mxfp8 = _import_b12x_mxfp8() + if mxfp8 is None: + return 0 + + token_counts = _b12x_mxfp8_warmup_token_counts( + max_tokens=max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + ) + seen_signatures: set[tuple[int, int, int, torch.dtype]] = set() + warmed = 0 + last_device: torch.device | None = None + + with torch.inference_mode(): + for layer in model.modules(): + packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) + if packed_weight is None: + continue + signature = ( + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ) + if signature in seen_signatures: + continue + seen_signatures.add(signature) + + device = torch.device(packed_weight.weight.values.device) + last_device = device + for tokens in token_counts: + source = torch.zeros( + (tokens, int(packed_weight.in_features)), + dtype=output_dtype, + device=device, + ) + mxfp8.mm( + source, + packed_weight, + expected_m=_b12x_mxfp8_expected_m(tokens), + stream=current_stream().cuda_stream, + ) + warmed += 1 + + if warmed > 0 and last_device is not None and last_device.type == "cuda": + torch.accelerator.synchronize(last_device) + + return warmed + + +def _b12x_mxfp8_linear( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, +) -> torch.Tensor: + del out_features + layer = get_forward_context().no_compile_layers[_resolve_layer_name(layer_name)] + return _apply_b12x_mxfp8_packed_linear(layer, x, bias) + + +def _b12x_mxfp8_linear_fake( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, +) -> torch.Tensor: + del bias, layer_name + return x.new_empty((*x.shape[:-1], out_features)) + + +direct_register_custom_op( + op_name="b12x_mxfp8_linear", + op_func=_b12x_mxfp8_linear, + fake_impl=_b12x_mxfp8_linear_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class B12xMxfp8LinearKernel(Mxfp8LinearKernel): + """ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "b12x MXFP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x MXFP8 kernels require a Blackwell 12x device" + if not _b12x_mxfp8_enabled(): + return False, "b12x MXFP8 GEMM is not enabled" + mxfp8 = _import_b12x_mxfp8() + if mxfp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not mxfp8.is_supported(): + return False, "b12x.gemm.mxfp8_linear is not supported" + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + del c + if not _b12x_mxfp8_enabled(): + return False, "b12x MXFP8 GEMM is not enabled" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight.data + if weight.dtype != MXFP8_VALUE_DTYPE: + raise ValueError( + f"b12x MXFP8 requires {MXFP8_VALUE_DTYPE}, got {weight.dtype}" + ) + if weight.ndim != 2: + raise ValueError(f"b12x MXFP8 weight must be 2D, got {weight.ndim}D") + if not hasattr(layer, "weight_scale"): + raise ValueError("b12x MXFP8 linear requires weight_scale") + + out_features, in_features = map(int, weight.shape) + if in_features % MXFP8_BLOCK_SIZE != 0: + raise ValueError( + "b12x MXFP8 requires input features divisible by " + f"{MXFP8_BLOCK_SIZE}, got {in_features}" + ) + weight_scale = layer.weight_scale.data + if weight_scale.dtype != MXFP8_SCALE_DTYPE: + raise ValueError( + f"b12x MXFP8 requires {MXFP8_SCALE_DTYPE} weight_scale, " + f"got {weight_scale.dtype}" + ) + if weight_scale.ndim != 2: + raise ValueError( + f"b12x MXFP8 weight_scale must be 2D, got {weight_scale.ndim}D" + ) + + mxfp8 = _import_b12x_mxfp8() + if mxfp8 is None: + raise ImportError("b12x.gemm.mxfp8_linear is not importable") + scale_k = in_features // MXFP8_BLOCK_SIZE + packed_weight = mxfp8.pack_weight( + weight[:out_features, :in_features].detach(), + weight_scale[:out_features, :scale_k].detach(), + ) + layer.b12x_mxfp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_mxfp8_packed_weight", None), + packed_weight, + ) + _register_b12x_mxfp8_linear_layer(layer) + replace_parameter(layer, "weight", weight.new_empty((0,))) + replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,))) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if torch.compiler.is_compiling(): + prefix = getattr(layer, "prefix", "") + if not prefix: + raise RuntimeError( + "B12X MXFP8 linear requires a layer prefix under torch.compile" + ) + packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) + if packed_weight is None: + raise RuntimeError( + "b12x MXFP8 packed weights are missing; " + "process_weights_after_loading did not run for this layer" + ) + return torch.ops.vllm.b12x_mxfp8_linear( + x, + bias, + _encode_layer_name(prefix), + int(packed_weight.out_features), + ) + + return _apply_b12x_mxfp8_packed_linear(layer, x, bias) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py new file mode 100644 index 000000000000..a507f6aa422f --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +from collections.abc import Iterable +from typing import Any + +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config_or_none +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.torch_utils import current_stream, direct_register_custom_op + +from .BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) + +_B12X_BLOCKSCALED: Any | None = None +_B12X_BLOCKSCALED_MISSING = False + + +def _import_b12x_blockscaled() -> Any | None: + global _B12X_BLOCKSCALED, _B12X_BLOCKSCALED_MISSING + if _B12X_BLOCKSCALED is not None: + return _B12X_BLOCKSCALED + if _B12X_BLOCKSCALED_MISSING: + return None + try: + _B12X_BLOCKSCALED = importlib.import_module("b12x.gemm.blockscaled") + except ImportError: + _B12X_BLOCKSCALED_MISSING = True + return None + return _B12X_BLOCKSCALED + + +def _current_linear_backend() -> str: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return "auto" + return str(getattr(vllm_config.kernel_config, "linear_backend", "auto")).lower() + + +def _b12x_block_fp8_warmup_token_counts( + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), +) -> tuple[int, ...]: + counts = {1} + counts.update(int(size) for size in cudagraph_capture_sizes if int(size) > 0) + if int(max_tokens) > 0: + counts.add(int(max_tokens)) + return tuple(sorted(counts)) + + +def _run_b12x_fp8_block_scaled_mm( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + if blockscaled is None: + raise ImportError("b12x regular block-FP8 GEMM is not importable") + + m, k = map(int, a.shape) + n = int(weight.shape[0]) + output = blockscaled.mm( + (a.reshape(m, k, 1), a_scale), + (weight.reshape(n, k, 1), weight_scale), + ab_dtype="float8_e4m3fn", + sf_dtype="float32", + c_dtype=str(out_dtype).removeprefix("torch."), + sf_vec_size=128, + block_fp8=True, + expected_m=m, + stream=current_stream().cuda_stream, + ) + return output[:, :, 0] + + +def _b12x_fp8_block_scaled_mm( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + return _run_b12x_fp8_block_scaled_mm( + a, + weight, + a_scale, + weight_scale, + out_dtype, + ) + + +def _b12x_fp8_block_scaled_mm_fake( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + del a_scale, weight_scale + return torch.empty( + (a.shape[0], weight.shape[0]), + dtype=out_dtype, + device=a.device, + ) + + +direct_register_custom_op( + op_name="b12x_fp8_block_scaled_mm", + op_func=_b12x_fp8_block_scaled_mm, + fake_impl=_b12x_fp8_block_scaled_mm_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +def warmup_b12x_block_fp8_linear( + model: torch.nn.Module, + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), + output_dtype: torch.dtype = torch.bfloat16, +) -> int: + if not current_platform.is_cuda(): + return 0 + if not current_platform.is_device_capability_family(120): + return 0 + if output_dtype not in (torch.bfloat16, torch.float16): + output_dtype = torch.bfloat16 + + blockscaled = _import_b12x_blockscaled() + if blockscaled is None: + return 0 + token_counts = _b12x_block_fp8_warmup_token_counts( + max_tokens=max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + ) + seen_signatures: set[tuple[Any, ...]] = set() + warmed = 0 + + with torch.inference_mode(): + for layer in model.modules(): + if not getattr(layer, "b12x_block_fp8_linear", False): + continue + weight = layer.weight + weight_scale = getattr(layer, "weight_scale_inv", None) + if weight_scale is None: + weight_scale = layer.weight_scale + n, k = map(int, weight.shape) + signature = ( + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ) + if signature in seen_signatures: + continue + seen_signatures.add(signature) + for tokens in token_counts: + a = torch.empty( + (tokens, k), + dtype=weight.dtype, + device=weight.device, + ) + a_scale = torch.empty( + (tokens, k // 128), + dtype=torch.float32, + device=weight.device, + ) + _run_b12x_fp8_block_scaled_mm( + a, + weight, + a_scale, + weight_scale, + output_dtype, + ) + warmed += 1 + return warmed + + +class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """K128 block-FP8 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X FP8 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "B12X regular block-FP8 GEMM is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + can_implement_base, reason = super().can_implement(config) + if not can_implement_base: + return can_implement_base, reason + + if _current_linear_backend() != "b12x" and not envs.VLLM_USE_B12X_FP8_GEMM: + return False, "B12X FP8 GEMM is not enabled" + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "Supports only bf16/fp16 input dtype" + if config.input_dtype != config.out_dtype: + return False, "Input and output dtype must match" + + act_group_shape = config.activation_quant_key.scale.group_shape + if act_group_shape != GroupShape(1, 128): + return ( + False, + "Supports only dynamic per-token group activation quantization " + "with group_shape=(1,128)", + ) + weight_group_shape = config.weight_quant_key.scale.group_shape + if weight_group_shape != GroupShape(128, 128): + return False, "Supports only 128x128 block-scaled FP8 weights" + + out_features, in_features = config.weight_shape + if in_features <= 0 or in_features % 128 != 0: + return False, "Input features must be a positive multiple of 128" + if out_features <= 0 or out_features % 128 != 0: + return False, "Output features must be a positive multiple of 128" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + params = self._get_layer_params(layer) + if params.weight_scale_inv is not None: + weight_scale = params.weight_scale_inv + scale_attr = params.WEIGHT_SCALE_INV + else: + weight_scale = params.weight_scale + scale_attr = params.WEIGHT_SCALE + if weight_scale is not None and weight_scale.dtype in ( + torch.float8_e8m0fnu, + torch.uint8, + ): + # TODO: Remove once B12X supports 128x128 UE8M0 block scales. + replace_parameter( + layer, + scale_attr, + _upcast_e8m0_to_fp32(weight_scale).contiguous(), + ) + layer.b12x_block_fp8_linear = True + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.vllm.b12x_fp8_block_scaled_mm( + A, + B, + As, + Bs, + self.config.out_dtype, + ) + + +__all__ = [ + "B12xFp8BlockScaledMMKernel", + "warmup_b12x_block_fp8_linear", +] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py new file mode 100644 index 000000000000..df75cec3f9d1 --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config_or_none +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.b12x_utils import reuse_packed_weight_storage +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + _USE_LAYERNAME, + LayerName, + _encode_layer_name, + current_stream, + direct_register_custom_op, +) + +from .ScaledMMLinearKernel import ( + FP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) + +_B12X_TENSOR_FP8: Any | None = None +_B12X_TENSOR_FP8_MISSING = False + +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | LayerName +else: + _layer_name_type = LayerName if _USE_LAYERNAME else str + + +def _import_b12x_tensor_fp8() -> Any | None: + global _B12X_TENSOR_FP8, _B12X_TENSOR_FP8_MISSING + if _B12X_TENSOR_FP8 is not None: + return _B12X_TENSOR_FP8 + if _B12X_TENSOR_FP8_MISSING: + return None + try: + _B12X_TENSOR_FP8 = importlib.import_module("b12x.gemm.tensor_fp8_linear") + except ImportError: + _B12X_TENSOR_FP8_MISSING = True + return None + return _B12X_TENSOR_FP8 + + +def _current_linear_backend() -> str: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return "auto" + return str(getattr(vllm_config.kernel_config, "linear_backend", "auto")).lower() + + +def _b12x_tensor_fp8_enabled() -> bool: + return _current_linear_backend() == "b12x" or envs.VLLM_USE_B12X_FP8_GEMM + + +@torch.compiler.assume_constant_result +def _resolve_layer_name(layer_name: str | LayerName) -> str: + from torch._library.fake_class_registry import FakeScriptObject + + if isinstance(layer_name, LayerName): + return layer_name.value + elif isinstance(layer_name, FakeScriptObject): + return layer_name.real_obj.value + return layer_name + + +def _register_b12x_tensor_fp8_linear_layer(layer: torch.nn.Module) -> None: + prefix = getattr(layer, "prefix", "") + if not prefix: + return + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return + static_forward_context = vllm_config.compilation_config.static_forward_context + existing = static_forward_context.get(prefix) + if existing is not None and existing is not layer: + raise ValueError(f"Duplicate B12X tensor FP8 linear layer name: {prefix}") + static_forward_context[prefix] = layer + + +def _apply_b12x_tensor_fp8_packed_linear( + layer: torch.nn.Module, + x_q: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + packed_weight = getattr(layer, "b12x_tensor_fp8_packed_weight", None) + if packed_weight is None: + raise RuntimeError( + "b12x tensor FP8 packed weights are missing; " + "process_weights_after_loading did not run for this layer" + ) + + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + raise ImportError("b12x.gemm.tensor_fp8_linear is not importable") + + input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous() + output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)] + output = tensor_fp8.mm( + input_2d, + packed_weight, + bias=bias, + out_dtype=out_dtype, + expected_m=max(1, int(input_2d.shape[0])), + stream=current_stream().cuda_stream, + ) + return output.view(*output_shape) + + +def _b12x_tensor_fp8_linear( + x_q: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + del out_features + layer = get_forward_context().no_compile_layers[_resolve_layer_name(layer_name)] + return _apply_b12x_tensor_fp8_packed_linear(layer, x_q, bias, out_dtype) + + +def _b12x_tensor_fp8_linear_fake( + x_q: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + del bias, layer_name + return torch.empty( + (*x_q.shape[:-1], out_features), + dtype=out_dtype, + device=x_q.device, + ) + + +direct_register_custom_op( + op_name="b12x_tensor_fp8_linear", + op_func=_b12x_tensor_fp8_linear, + fake_impl=_b12x_tensor_fp8_linear_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +def _warmup_token_counts( + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), +) -> tuple[int, ...]: + counts = {1} + counts.update(int(size) for size in cudagraph_capture_sizes if int(size) > 0) + if int(max_tokens) > 0: + counts.add(int(max_tokens)) + return tuple(sorted(counts)) + + +def warmup_b12x_tensor_fp8_linear( + model: torch.nn.Module, + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), + output_dtype: torch.dtype = torch.bfloat16, +) -> int: + if not _b12x_tensor_fp8_enabled(): + return 0 + if not current_platform.is_cuda(): + return 0 + if not current_platform.is_device_capability_family(120): + return 0 + + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + return 0 + if output_dtype not in (torch.bfloat16, torch.float16): + output_dtype = torch.bfloat16 + + token_counts = _warmup_token_counts( + max_tokens=max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + ) + seen_signatures: set[tuple[int, int, int, torch.dtype]] = set() + warmed = 0 + last_device: torch.device | None = None + + with torch.inference_mode(): + for layer in model.modules(): + packed_weight = getattr( + layer, + "b12x_tensor_fp8_packed_weight", + None, + ) + if packed_weight is None: + continue + signature = ( + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ) + if signature in seen_signatures: + continue + seen_signatures.add(signature) + last_device = torch.device(packed_weight.values.device) + warmed += int( + tensor_fp8.prewarm( + packed_weight, + token_counts, + out_dtype=output_dtype, + stream=current_stream().cuda_stream, + ) + ) + + if warmed > 0 and last_device is not None and last_device.type == "cuda": + torch.accelerator.synchronize(last_device) + + return warmed + + +class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """Static per-tensor FP8 linear through the B12X SM12x dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "b12x tensor FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x tensor FP8 kernels require a Blackwell 12x device" + if not _b12x_tensor_fp8_enabled(): + return False, "b12x tensor FP8 GEMM is not enabled" + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not tensor_fp8.is_supported(): + return False, "b12x.gemm.tensor_fp8_linear is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + if not _b12x_tensor_fp8_enabled(): + return False, "b12x tensor FP8 GEMM is not enabled" + activation_scale = config.activation_quant_key.scale + weight_scale = config.weight_quant_key.scale + if ( + not activation_scale.static + or not activation_scale.group_shape.is_per_tensor() + ): + return False, "requires static per-tensor activation scales" + if not weight_scale.static or not weight_scale.group_shape.is_per_tensor(): + return False, "requires static per-tensor weight scales" + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 input dtype" + if config.out_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 output dtype" + out_features, in_features = config.weight_shape + if out_features <= 0 or in_features <= 0 or in_features % 32 != 0: + return False, "weight dimensions must be positive with K divisible by 32" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight, weight_scale, input_scale, _ = self._get_layer_params(layer) + if weight.dtype != torch.float8_e4m3fn: + raise ValueError( + f"b12x tensor FP8 requires float8_e4m3fn weight, got {weight.dtype}" + ) + if weight_scale.numel() != 1 or input_scale is None or input_scale.numel() != 1: + raise ValueError( + "b12x tensor FP8 requires scalar weight and activation scales" + ) + + out_features, in_features = map(int, self.config.weight_shape) + if tuple(weight.shape) != (in_features, out_features): + raise ValueError( + "b12x tensor FP8 expects the processed weight in [K,N] layout, " + f"got {tuple(weight.shape)} for N={out_features}, K={in_features}" + ) + + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + raise ImportError("b12x.gemm.tensor_fp8_linear is not importable") + output_scale = ( + input_scale.detach().to(torch.float32).reshape(1) + * weight_scale.detach().to(torch.float32).reshape(1) + ).contiguous() + packed_weight = tensor_fp8.pack_weight( + weight.detach().T.contiguous(), + output_scale, + ) + layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_tensor_fp8_packed_weight", None), + packed_weight, + ) + _register_b12x_tensor_fp8_linear_layer(layer) + weight_name, weight_scale_name, _, _ = self.layer_param_names + replace_parameter(layer, weight_name, weight.new_empty((0,))) + replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,))) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if not isinstance(x, torch.Tensor): + raise TypeError("b12x tensor FP8 linear requires a Tensor input") + _, _, input_scale, input_scale_ub = self._get_layer_params(layer) + input_2d = x.reshape(-1, x.shape[-1]) + x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub) + out_dtype = self.config.out_dtype + if torch.compiler.is_compiling(): + prefix = getattr(layer, "prefix", "") + if not prefix: + raise RuntimeError( + "B12X tensor FP8 linear requires a layer prefix under torch.compile" + ) + packed_weight = getattr( + layer, + "b12x_tensor_fp8_packed_weight", + None, + ) + if packed_weight is None: + raise RuntimeError( + "b12x tensor FP8 packed weights are missing; " + "process_weights_after_loading did not run for this layer" + ) + output = torch.ops.vllm.b12x_tensor_fp8_linear( + x_q, + bias, + _encode_layer_name(prefix), + int(packed_weight.out_features), + out_dtype, + ) + else: + output = _apply_b12x_tensor_fp8_packed_linear( + layer, + x_q, + bias, + out_dtype, + ) + return output.view(*x.shape[:-1], output.shape[-1]) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + del A, B, out_dtype, As, Bs, bias, output_shape + raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights") + + +__all__ = [ + "B12xTensorFP8ScaledMMLinearKernel", + "warmup_b12x_tensor_fp8_linear", +] From a5b529afd994d68c880479ddfa14432fb2fdeb70 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Fri, 7 Aug 2026 09:41:27 -0700 Subject: [PATCH 2/5] feat(linear): add B12X FP4 backends Add native B12X NVFP4 and MXFP4 dense linear kernels, selected explicitly or through the opt-in FP4 override. Preserve automatic fallback for unsupported W4A16 linears and cover dispatch, packing, and execution. Assisted-by: OpenAI Codex Signed-off-by: Luke Alonso --- .../kernels/test_b12x_mxfp4_linear.py | 186 ++++++++++++++ .../kernels/test_b12x_nvfp4_linear.py | 191 +++++++++++++++ vllm/envs.py | 6 +- .../model_executor/kernels/linear/__init__.py | 12 + .../kernels/linear/mxfp4/b12x.py | 216 +++++++++++++++++ .../kernels/linear/nvfp4/b12x.py | 226 ++++++++++++++++++ 6 files changed, 836 insertions(+), 1 deletion(-) create mode 100644 tests/model_executor/kernels/test_b12x_mxfp4_linear.py create mode 100644 tests/model_executor/kernels/test_b12x_nvfp4_linear.py create mode 100644 vllm/model_executor/kernels/linear/mxfp4/b12x.py create mode 100644 vllm/model_executor/kernels/linear/nvfp4/b12x.py diff --git a/tests/model_executor/kernels/test_b12x_mxfp4_linear.py b/tests/model_executor/kernels/test_b12x_mxfp4_linear.py new file mode 100644 index 000000000000..c18b0b1c8f35 --- /dev/null +++ b/tests/model_executor/kernels/test_b12x_mxfp4_linear.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import types + +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + _POSSIBLE_MXFP4_KERNELS, + init_mxfp4_linear_kernel, +) +from vllm.model_executor.kernels.linear.mxfp4.b12x import ( + B12xMxFp4LinearKernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) +from vllm.platforms import PlatformEnum + + +def test_b12x_backend_maps_mxfp4_kernel() -> None: + assert B12xMxFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + assert B12xMxFp4LinearKernel in _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA] + + +def test_b12x_mxfp4_explicit_backend_selects_native_kernel(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: "b12x") + monkeypatch.setattr( + B12xMxFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + monkeypatch.setattr( + B12xMxFp4LinearKernel, + "can_implement", + classmethod(lambda cls, config: (True, None)), + ) + + kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) + + assert isinstance(kernel, B12xMxFp4LinearKernel) + + +def test_b12x_fp4_env_selects_mxfp4_with_auto_backend(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: "auto") + monkeypatch.setattr(linear_mod.envs, "VLLM_USE_B12X_FP4_GEMM", True) + monkeypatch.setattr( + B12xMxFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) + + assert isinstance(kernel, B12xMxFp4LinearKernel) + + +def test_b12x_mxfp4_env_enables_auto_backend(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "auto") + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP4_GEMM", False) + + config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic) + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) + + assert not can_implement + assert reason == "B12X MXFP4 GEMM is not enabled" + + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP4_GEMM", True) + 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" + + +def test_b12x_mxfp4_processes_scale_and_preserves_loader(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod + + scale = torch.empty((48, 8), dtype=torch.uint8) + swizzled_scale = torch.empty((128, 8), dtype=torch.uint8) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_intrinsics", + lambda: types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale), + ) + + 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 + vllm_config = VllmConfig() + monkeypatch.setattr( + b12x_mod, "get_current_vllm_config_or_none", lambda: vllm_config + ) + kernel = object.__new__(B12xMxFp4LinearKernel) + + 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 vllm_config.compilation_config.static_forward_context[layer.prefix] is layer + + +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) + x_scale = torch.empty((32, 4, 1, 4, 1, 1), dtype=torch.float8_e8m0fnu) + weight_scale = torch.empty((32, 4, 1, 4, 1, 1), dtype=torch.float8_e8m0fnu) + + def as_grouped_scale_view_mx(storage, rows: int, cols: int): + return x_scale if rows == 6 else weight_scale + + def mm(lhs, rhs, **kwargs): + calls.append((lhs, rhs, kwargs)) + return torch.full((6, 48, 1), 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=mm), + ) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_intrinsics", + lambda: types.SimpleNamespace( + as_grouped_scale_view_mx=as_grouped_scale_view_mx + ), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: types.SimpleNamespace(cuda_stream=123), + ) + + 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 + lhs, rhs, kwargs = calls[0] + assert lhs[0].data_ptr() == x_packed.data_ptr() + assert lhs[0].shape == (6, 64, 1) + assert lhs[1] is x_scale + assert rhs[0].data_ptr() == layer.weight.data_ptr() + assert rhs[0].shape == (48, 64, 1) + assert rhs[1] is weight_scale + assert kwargs["ab_dtype"] == "float4_e2m1fn" + assert kwargs["sf_dtype"] == "float8_e8m0fnu" + assert kwargs["sf_vec_size"] == 32 + assert kwargs["expected_m"] == 6 + assert kwargs["stream"] == 123 diff --git a/tests/model_executor/kernels/test_b12x_nvfp4_linear.py b/tests/model_executor/kernels/test_b12x_nvfp4_linear.py new file mode 100644 index 000000000000..d590ff7bed06 --- /dev/null +++ b/tests/model_executor/kernels/test_b12x_nvfp4_linear.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import types + +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + _POSSIBLE_NVFP4_KERNELS, + init_nvfp4_linear_kernel, +) +from vllm.model_executor.kernels.linear.nvfp4.b12x import ( + B12xNvFp4LinearKernel, +) +from vllm.model_executor.kernels.linear.nvfp4.marlin import ( + MarlinNvFp4LinearKernel, +) +from vllm.platforms import PlatformEnum + + +def test_b12x_backend_maps_nvfp4_kernel() -> None: + assert B12xNvFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + assert B12xNvFp4LinearKernel in _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA] + + +def test_b12x_nvfp4_explicit_backend_selects_native_kernel(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: "b12x") + monkeypatch.setattr( + B12xNvFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + monkeypatch.setattr( + B12xNvFp4LinearKernel, + "can_implement", + classmethod(lambda cls, config: (True, None)), + ) + + kernel = init_nvfp4_linear_kernel() + + assert isinstance(kernel, B12xNvFp4LinearKernel) + + +def test_b12x_fp4_env_selects_nvfp4_with_auto_backend(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: "auto") + monkeypatch.setattr(linear_mod.envs, "VLLM_USE_B12X_FP4_GEMM", True) + monkeypatch.setattr( + B12xNvFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_nvfp4_linear_kernel() + + assert isinstance(kernel, B12xNvFp4LinearKernel) + + +def test_b12x_nvfp4_env_enables_auto_backend(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod, "_current_linear_backend", lambda: "auto") + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP4_GEMM", False) + + can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) + + assert not can_implement + assert reason == "B12X NVFP4 GEMM is not enabled" + + monkeypatch.setattr(b12x_mod.envs, "VLLM_USE_B12X_FP4_GEMM", True) + can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) + + assert can_implement + assert reason is None + + +def test_b12x_fp4_env_preserves_w4a16_auto_fallback(monkeypatch) -> None: + import vllm.model_executor.kernels.linear as linear_mod + + monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "auto") + monkeypatch.setattr(linear_mod.envs, "VLLM_USE_B12X_FP4_GEMM", True) + monkeypatch.setattr( + MarlinNvFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_nvfp4_linear_kernel(use_a16=True) + + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_b12x_nvfp4_processes_scale_and_registers_layer(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + scale = torch.empty((48, 8), dtype=torch.float8_e4m3fn) + swizzled_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale) + monkeypatch.setattr(b12x_mod, "_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 + vllm_config = VllmConfig() + monkeypatch.setattr( + b12x_mod, "get_current_vllm_config_or_none", lambda: vllm_config + ) + kernel = object.__new__(B12xNvFp4LinearKernel) + + 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 vllm_config.compilation_config.static_forward_context[layer.prefix] is layer + + +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] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + x_scale = torch.empty((32, 4, 1, 4, 2, 1), dtype=torch.float8_e4m3fn) + weight_scale = torch.empty((32, 4, 1, 4, 2, 1), dtype=torch.float8_e4m3fn) + + def as_grouped_scale_view(storage, rows: int, cols: int): + return x_scale if rows == 6 else weight_scale + + def mm(lhs, rhs, **kwargs): + calls.append((lhs, rhs, kwargs)) + return torch.full((6, 48, 1), 3.0, dtype=torch.bfloat16) + + monkeypatch.setattr( + b12x_mod, + "scaled_fp4_quant", + lambda *args, **kwargs: (x_packed, x_scale_storage), + ) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm=mm), + ) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_intrinsics", + lambda: types.SimpleNamespace(as_grouped_scale_view=as_grouped_scale_view), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: types.SimpleNamespace(cuda_stream=123), + ) + + 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, 128), dtype=torch.bfloat16) + 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(calls) == 1 + lhs, rhs, kwargs = calls[0] + assert lhs[0].data_ptr() == x_packed.data_ptr() + assert lhs[0].shape == (6, 64, 1) + assert lhs[1] is x_scale + assert rhs[0].data_ptr() == layer.weight.data_ptr() + assert rhs[0].shape == (48, 64, 1) + assert rhs[1] is weight_scale + assert kwargs["ab_dtype"] == "float4_e2m1fn" + assert kwargs["sf_dtype"] == "float8_e4m3fn" + assert kwargs["sf_vec_size"] == 16 + assert kwargs["expected_m"] == 6 + assert kwargs["stream"] == 123 diff --git a/vllm/envs.py b/vllm/envs.py index 29c35b339a8d..657c1cd21068 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -183,6 +183,7 @@ VLLM_HUMMING_USE_F16_ACCUM: bool = False VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None VLLM_USE_B12X_FP8_GEMM: bool = False + VLLM_USE_B12X_FP4_GEMM: bool = False VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False VLLM_TPU_USING_PATHWAYS: bool = False @@ -1569,10 +1570,13 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) ), - # Prefer B12X dense GEMMs for supported FP8 linear layers. + # Prefer B12X dense GEMMs for supported FP8 and FP4 linear layers. "VLLM_USE_B12X_FP8_GEMM": lambda: bool( int(os.getenv("VLLM_USE_B12X_FP8_GEMM", "0")) ), + "VLLM_USE_B12X_FP4_GEMM": lambda: bool( + int(os.getenv("VLLM_USE_B12X_FP4_GEMM", "0")) + ), # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool( int(os.getenv("VLLM_USE_FLASHINFER_MOE_INT4", "0")) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 0b582989f616..8935a0968127 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -77,6 +77,9 @@ from vllm.model_executor.kernels.linear.mxfp4.aiter import ( AiterMxfp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.b12x import ( + B12xMxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.emulation import ( EmulationMxfp4LinearKernel, ) @@ -129,6 +132,9 @@ NvFp4LinearKernel, NvFp4LinearLayerConfig, ) +from vllm.model_executor.kernels.linear.nvfp4.b12x import ( + B12xNvFp4LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( CutlassNvFp4LinearKernel, ) @@ -239,7 +245,9 @@ def _get_linear_backend() -> str: _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "b12x": { B12xFp8BlockScaledMMKernel, + B12xMxFp4LinearKernel, B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, B12xTensorFP8ScaledMMLinearKernel, }, "cutlass": { @@ -515,6 +523,7 @@ def _resolve_backend_kernels( _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { PlatformEnum.CUDA: [ + B12xNvFp4LinearKernel, FlashInferCuteDslNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferB12xNvFp4LinearKernel, @@ -542,6 +551,7 @@ def _resolve_backend_kernels( _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { PlatformEnum.CUDA: [ + B12xMxFp4LinearKernel, FlashInferMxFp4LinearKernel, MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel, @@ -1192,6 +1202,8 @@ def register_linear_kernel( "Mxfp8LinearKernel", "Mxfp8LinearLayerConfig", "B12xMxfp8LinearKernel", + "B12xMxFp4LinearKernel", + "B12xNvFp4LinearKernel", "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", diff --git a/vllm/model_executor/kernels/linear/mxfp4/b12x.py b/vllm/model_executor/kernels/linear/mxfp4/b12x.py new file mode 100644 index 000000000000..6ecaffe1cf66 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/b12x.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING, Any + +import torch + +import vllm.envs as envs +from vllm.config import get_current_vllm_config_or_none +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + _USE_LAYERNAME, + LayerName, + _encode_layer_name, + current_stream, + direct_register_custom_op, +) + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | LayerName +else: + _layer_name_type = LayerName if _USE_LAYERNAME else str + +_MXFP4_GROUP_SIZE = 32 +_B12X_BLOCKSCALED: Any | None = None +_B12X_INTRINSICS: Any | None = None + + +def _import_b12x_blockscaled() -> Any | None: + global _B12X_BLOCKSCALED + if _B12X_BLOCKSCALED is None: + try: + _B12X_BLOCKSCALED = importlib.import_module("b12x.gemm.blockscaled") + except ImportError: + return None + return _B12X_BLOCKSCALED + + +def _import_b12x_intrinsics() -> Any | None: + global _B12X_INTRINSICS + if _B12X_INTRINSICS is None: + try: + _B12X_INTRINSICS = importlib.import_module("b12x._lib.intrinsics") + except ImportError: + return None + return _B12X_INTRINSICS + + +def _current_linear_backend() -> str: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return "auto" + return str(getattr(vllm_config.kernel_config, "linear_backend", "auto")).lower() + + +@torch.compiler.assume_constant_result +def _resolve_layer_name(layer_name: str | LayerName) -> str: + from torch._library.fake_class_registry import FakeScriptObject + + if isinstance(layer_name, LayerName): + return layer_name.value + elif isinstance(layer_name, FakeScriptObject): + return layer_name.real_obj.value + return layer_name + + +def _register_b12x_mxfp4_linear_layer(layer: torch.nn.Module) -> None: + prefix = getattr(layer, "prefix", "") + if not prefix: + return + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return + static_forward_context = vllm_config.compilation_config.static_forward_context + existing = static_forward_context.get(prefix) + if existing is not None and existing is not layer: + raise ValueError(f"Duplicate B12X MXFP4 linear layer name: {prefix}") + static_forward_context[prefix] = layer + + +def _apply_b12x_mxfp4_linear( + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + from vllm.utils.flashinfer import flashinfer_mxfp4_quantize + + blockscaled = _import_b12x_blockscaled() + intrinsics = _import_b12x_intrinsics() + if blockscaled is None or intrinsics is None: + raise ImportError("b12x native MXFP4 GEMM is not importable") + + output_size = int(layer.output_size_per_partition) + output_shape = [*x.shape[:-1], output_size] + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + m, k = map(int, x_2d.shape) + x_packed, x_scale_swizzled = flashinfer_mxfp4_quantize(x_2d, backend="cute-dsl") + x_scale = intrinsics.as_grouped_scale_view_mx( + x_scale_swizzled.view(torch.uint8).unsqueeze(0), m, k + ) + weight_scale = intrinsics.as_grouped_scale_view_mx( + layer.weight_scale.view(torch.uint8).unsqueeze(0), output_size, k + ) + output = blockscaled.mm( + (x_packed.unsqueeze(-1), x_scale), + (layer.weight.unsqueeze(-1), weight_scale), + ab_dtype="float4_e2m1fn", + sf_dtype="float8_e8m0fnu", + c_dtype=str(x.dtype).split(".")[-1], + sf_vec_size=_MXFP4_GROUP_SIZE, + expected_m=m, + stream=current_stream().cuda_stream, + )[:, :, 0] + if bias is not None: + output = output + bias + return output.view(*output_shape) + + +def _b12x_mxfp4_linear( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, +) -> torch.Tensor: + del out_features + layer = get_forward_context().no_compile_layers[_resolve_layer_name(layer_name)] + return _apply_b12x_mxfp4_linear(layer, x, bias) + + +def _b12x_mxfp4_linear_fake( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, +) -> torch.Tensor: + del bias, layer_name + return x.new_empty((*x.shape[:-1], out_features)) + + +direct_register_custom_op( + op_name="b12x_mxfp4_linear", + op_func=_b12x_mxfp4_linear, + fake_impl=_b12x_mxfp4_linear_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class B12xMxFp4LinearKernel(MxFp4LinearKernel): + """MXFP4 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X MXFP4 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X MXFP4 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None or _import_b12x_intrinsics() is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "b12x native MXFP4 GEMM is not supported" + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if _current_linear_backend() != "b12x" and not envs.VLLM_USE_B12X_FP4_GEMM: + return False, "B12X MXFP4 GEMM is not enabled" + if config.activation_quant_key != kMxfp4Dynamic: + return False, "B12X MXFP4 GEMM requires dynamic MXFP4 activations" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + intrinsics = _import_b12x_intrinsics() + if intrinsics is None: + raise ImportError("b12x native MXFP4 GEMM is not importable") + replace_parameter( + layer, + "weight_scale", + intrinsics.swizzle_block_scale(layer.weight_scale.data), + ) + _register_b12x_mxfp4_linear_layer(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if torch.compiler.is_compiling(): + prefix = getattr(layer, "prefix", "") + if not prefix: + raise RuntimeError( + "B12X MXFP4 linear requires a layer prefix under torch.compile" + ) + return torch.ops.vllm.b12x_mxfp4_linear( + x, + bias, + _encode_layer_name(prefix), + int(layer.output_size_per_partition), + ) + return _apply_b12x_mxfp4_linear(layer, x, bias) diff --git a/vllm/model_executor/kernels/linear/nvfp4/b12x.py b/vllm/model_executor/kernels/linear/nvfp4/b12x.py new file mode 100644 index 000000000000..e8b576af892d --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/b12x.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING, Any + +import torch + +import vllm.envs as envs +from vllm._custom_ops import scaled_fp4_quant +from vllm.config import get_current_vllm_config_or_none +from vllm.forward_context import get_forward_context +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + _USE_LAYERNAME, + LayerName, + _encode_layer_name, + current_stream, + direct_register_custom_op, +) + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + +if TYPE_CHECKING: + from typing import TypeAlias + + _layer_name_type: TypeAlias = str | LayerName +else: + _layer_name_type = LayerName if _USE_LAYERNAME else str + +_B12X_BLOCKSCALED: Any | None = None +_B12X_INTRINSICS: Any | None = None +_B12X_MISSING = False + + +def _import_b12x_blockscaled() -> Any | None: + global _B12X_BLOCKSCALED, _B12X_MISSING + if _B12X_BLOCKSCALED is not None: + return _B12X_BLOCKSCALED + if _B12X_MISSING: + return None + try: + _B12X_BLOCKSCALED = importlib.import_module("b12x.gemm.blockscaled") + except ImportError: + _B12X_MISSING = True + return None + return _B12X_BLOCKSCALED + + +def _import_b12x_intrinsics() -> Any | None: + global _B12X_INTRINSICS, _B12X_MISSING + if _B12X_INTRINSICS is not None: + return _B12X_INTRINSICS + if _B12X_MISSING: + return None + try: + _B12X_INTRINSICS = importlib.import_module("b12x._lib.intrinsics") + except ImportError: + _B12X_MISSING = True + return None + return _B12X_INTRINSICS + + +def _current_linear_backend() -> str: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return "auto" + return str(getattr(vllm_config.kernel_config, "linear_backend", "auto")).lower() + + +@torch.compiler.assume_constant_result +def _resolve_layer_name(layer_name: str | LayerName) -> str: + from torch._library.fake_class_registry import FakeScriptObject + + if isinstance(layer_name, LayerName): + return layer_name.value + elif isinstance(layer_name, FakeScriptObject): + return layer_name.real_obj.value + return layer_name + + +def _register_b12x_nvfp4_linear_layer(layer: torch.nn.Module) -> None: + prefix = getattr(layer, "prefix", "") + if not prefix: + return + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return + static_forward_context = vllm_config.compilation_config.static_forward_context + existing = static_forward_context.get(prefix) + if existing is not None and existing is not layer: + raise ValueError(f"Duplicate B12X NVFP4 linear layer name: {prefix}") + static_forward_context[prefix] = layer + + +def _apply_b12x_nvfp4_linear( + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + intrinsics = _import_b12x_intrinsics() + if blockscaled is None or intrinsics is None: + raise ImportError("b12x native NVFP4 GEMM is not importable") + + output_size = int(layer.output_size_per_partition) + output_shape = [*x.shape[:-1], output_size] + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + m, k = map(int, x_2d.shape) + x_packed, x_scale_swizzled = scaled_fp4_quant( + x_2d, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="cutlass", + ) + x_scale = intrinsics.as_grouped_scale_view( + x_scale_swizzled.view(torch.uint8).unsqueeze(0), m, k + ) + weight_scale = intrinsics.as_grouped_scale_view( + layer.weight_scale.view(torch.uint8).unsqueeze(0), output_size, k + ) + output = blockscaled.mm( + (x_packed.unsqueeze(-1), x_scale), + (layer.weight.unsqueeze(-1), weight_scale), + ab_dtype="float4_e2m1fn", + sf_dtype="float8_e4m3fn", + c_dtype=str(x.dtype).split(".")[-1], + sf_vec_size=16, + alpha=layer.alpha.view(1), + expected_m=m, + stream=current_stream().cuda_stream, + ) + output = output[:, :, 0] + if bias is not None: + output = output + bias + return output.view(*output_shape) + + +def _b12x_nvfp4_linear( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, +) -> torch.Tensor: + del out_features + layer = get_forward_context().no_compile_layers[_resolve_layer_name(layer_name)] + return _apply_b12x_nvfp4_linear(layer, x, bias) + + +def _b12x_nvfp4_linear_fake( + x: torch.Tensor, + bias: torch.Tensor | None, + layer_name: _layer_name_type, + out_features: int, +) -> torch.Tensor: + del bias, layer_name + return x.new_empty((*x.shape[:-1], out_features)) + + +direct_register_custom_op( + op_name="b12x_nvfp4_linear", + op_func=_b12x_nvfp4_linear, + fake_impl=_b12x_nvfp4_linear_fake, + tags=(torch.Tag.needs_fixed_stride_order,), +) + + +class B12xNvFp4LinearKernel(NvFp4LinearKernel): + """ModelOpt NVFP4 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X NVFP4 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X NVFP4 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None or _import_b12x_intrinsics() is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "b12x native NVFP4 GEMM is not supported" + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + del config + if _current_linear_backend() != "b12x" and not envs.VLLM_USE_B12X_FP4_GEMM: + return False, "B12X NVFP4 GEMM is not enabled" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + intrinsics = _import_b12x_intrinsics() + if intrinsics is None: + raise ImportError("b12x native NVFP4 GEMM is not importable") + replace_parameter( + layer, + "weight_scale", + intrinsics.swizzle_block_scale(layer.weight_scale.data), + ) + _register_b12x_nvfp4_linear_layer(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if torch.compiler.is_compiling(): + prefix = getattr(layer, "prefix", "") + if not prefix: + raise RuntimeError( + "B12X NVFP4 linear requires a layer prefix under torch.compile" + ) + return torch.ops.vllm.b12x_nvfp4_linear( + x, + bias, + _encode_layer_name(prefix), + int(layer.output_size_per_partition), + ) + return _apply_b12x_nvfp4_linear(layer, x, bias) From 9b3afae1fdef4a931a0e1be68a714fce136ba891 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Fri, 7 Aug 2026 11:44:36 -0700 Subject: [PATCH 3/5] feat(moe): add B12X FP4 backend Add a native tensor-parallel B12X backend for NVFP4 and MXFP4 MoE, including native NVFP4 W4A4, W4A16, and supported dynamic W4A8 modes. Integrate backend selection, weight preparation, warmup, and CUDA graph replay without adding EXL3 or NF3 MoE paths. Assisted-by: OpenAI Codex Signed-off-by: Luke Alonso --- tests/kernels/moe/test_b12x_moe.py | 1642 +++++++++++++++++ tests/quantization/test_auto_round.py | 9 +- vllm/config/kernel.py | 6 +- vllm/envs.py | 5 + .../layers/fused_moe/b12x_moe.py | 871 +++++++++ .../model_executor/layers/fused_moe/config.py | 4 + .../layers/fused_moe/oracle/mxfp4.py | 74 +- .../layers/fused_moe/oracle/nvfp4.py | 61 +- .../compressed_tensors_moe_w4a4_mxfp4.py | 11 +- .../compressed_tensors_moe_w4a4_nvfp4.py | 3 + .../quantization/inc/schemes/inc_mxfp4_moe.py | 8 +- .../layers/quantization/modelopt.py | 2 + .../layers/quantization/mxfp4.py | 1 + .../layers/quantization/utils/b12x_moe.py | 151 ++ vllm/model_executor/warmup/kernel_warmup.py | 76 +- 15 files changed, 2902 insertions(+), 22 deletions(-) create mode 100644 tests/kernels/moe/test_b12x_moe.py create mode 100644 vllm/model_executor/layers/fused_moe/b12x_moe.py create mode 100644 vllm/model_executor/layers/quantization/utils/b12x_moe.py diff --git a/tests/kernels/moe/test_b12x_moe.py b/tests/kernels/moe/test_b12x_moe.py new file mode 100644 index 000000000000..89c40c6a32ae --- /dev/null +++ b/tests/kernels/moe/test_b12x_moe.py @@ -0,0 +1,1642 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the native B12X tensor-parallel MoE integration.""" + +from dataclasses import dataclass, replace +from types import SimpleNamespace + +import pytest +import torch + +import vllm.model_executor.layers.fused_moe.b12x_moe as b12x_moe +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle +import vllm.model_executor.layers.fused_moe.oracle.nvfp4 as nvfp4_oracle +from tests.kernels.moe.utils import make_dummy_moe_config +from tests.kernels.quantization.nvfp4_utils import ( + FLOAT4_E2M1_MAX, + FLOAT8_E4M3_MAX, + break_fp4_bytes, +) +from tests.kernels.utils import torch_moe +from vllm import _custom_ops as ops +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import fused_topk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.b12x_moe import B12xExperts +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + FusedMoEQuantConfig, + mxfp4_w4a16_moe_quant_config, + nvfp4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + Mxfp4MoeBackend, + select_deepseek_v4_mxfp4_moe_backend, + select_mxfp4_moe_backend, +) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + select_nvfp4_moe_backend, +) +from vllm.model_executor.layers.quantization.utils.b12x_moe import ( + prepare_nvfp4_moe_layer_for_b12x, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Static, + kMxfp8Dynamic, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + + +def _quantize_nvfp4_linear( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + weights_q = [] + scales = [] + global_scales = [] + for expert_weight in weight: + global_scale = ( + FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / expert_weight.abs().max() + ).to(torch.float32) + weight_q, scale = ops.scaled_fp4_quant( + expert_weight, + global_scale, + is_sf_swizzled_layout=False, + ) + weights_q.append(weight_q) + scales.append(scale) + global_scales.append(global_scale) + return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales) + + +def _dequantize_nvfp4_linear( + tensor_fp4: torch.Tensor, + tensor_sf: torch.Tensor, + global_scale: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + rows, packed_cols = tensor_fp4.shape + cols = packed_cols * 2 + values = break_fp4_bytes(tensor_fp4, torch.float32) + values = values.reshape(rows, cols // 16, 16) + scales = tensor_sf.view(torch.float8_e4m3fn).to(torch.float32) + return ( + (values * (scales[:, : cols // 16] / global_scale).unsqueeze(-1)) + .reshape(rows, cols) + .to(dtype) + ) + + +def _nvfp4_activation_reference( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + a1_scale: torch.Tensor, + a2_scale: torch.Tensor, +) -> torch.Tensor: + tokens, hidden_size = hidden_states.shape + topk = topk_ids.shape[1] + routed_input = ( + hidden_states[:, None, :] + .expand(-1, topk, -1) + .reshape(tokens * topk, hidden_size) + ) + routed_output = torch.zeros( + tokens * topk, + hidden_size, + dtype=torch.float32, + device=hidden_states.device, + ) + flat_ids = topk_ids.reshape(-1) + + for expert in range(w1.shape[0]): + mask = flat_ids == expert + if not mask.any(): + continue + a1_q, a1_block_scale = ops.scaled_fp4_quant( + routed_input[mask], + a1_scale[expert], + is_sf_swizzled_layout=False, + ) + a1 = _dequantize_nvfp4_linear( + a1_q, + a1_block_scale, + a1_scale[expert], + torch.float32, + ) + fc1 = a1 @ w1[expert].float().t() + gate, up = fc1.chunk(2, dim=-1) + intermediate = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16) + a2_q, a2_block_scale = ops.scaled_fp4_quant( + intermediate, + a2_scale[expert], + is_sf_swizzled_layout=False, + ) + a2 = _dequantize_nvfp4_linear( + a2_q, + a2_block_scale, + a2_scale[expert], + torch.float32, + ) + routed_output[mask] = a2 @ w2[expert].float().t() + + return ( + routed_output.view(tokens, topk, hidden_size) + .mul(topk_weights[..., None]) + .sum(dim=1) + .to(hidden_states.dtype) + ) + + +def _e8m0_bytes_to_float(scale_bytes: torch.Tensor) -> torch.Tensor: + return torch.exp2(scale_bytes.view(torch.uint8).to(torch.float32) - 127.0) + + +def _e8m0_scale_bytes_from_amax(amax: torch.Tensor) -> torch.Tensor: + scale = amax.to(torch.float32) / 6.0 + scale = torch.where(scale > 0, scale, torch.full_like(scale, 2.0**-127)) + exponent = torch.ceil(torch.log2(scale)).clamp(min=-127.0, max=120.0) + return (exponent + 127.0).to(torch.uint8) + + +def _mxfp4_decode_packed(packed: torch.Tensor, cols: int) -> torch.Tensor: + fp4_lut = torch.tensor( + [ + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, + ], + dtype=torch.float32, + device=packed.device, + ) + packed = packed.view(torch.uint8) + lo = (packed & 0x0F).to(torch.long) + hi = ((packed >> 4) & 0x0F).to(torch.long) + codes = torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + return fp4_lut[codes[:, :cols]] + + +def _mxfp4_encode_values(values: torch.Tensor) -> torch.Tensor: + mag = values.abs() + codes = torch.zeros_like(mag, dtype=torch.uint8) + codes = torch.where((mag > 0.25) & (mag < 0.75), 1, codes) + codes = torch.where((mag >= 0.75) & (mag <= 1.25), 2, codes) + codes = torch.where((mag > 1.25) & (mag < 1.75), 3, codes) + codes = torch.where((mag >= 1.75) & (mag <= 2.5), 4, codes) + codes = torch.where((mag > 2.5) & (mag < 3.5), 5, codes) + codes = torch.where((mag >= 3.5) & (mag <= 5.0), 6, codes) + codes = torch.where(mag > 5.0, 7, codes) + return codes | ((values < 0).to(torch.uint8) << 3) + + +def _quantize_mxfp4(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + rows, cols = weight.shape[-2:] + blocks = weight.float().reshape(-1, cols // 32, 32) + scale_bytes = _e8m0_scale_bytes_from_amax(blocks.abs().amax(dim=-1)) + scales = _e8m0_bytes_to_float(scale_bytes).unsqueeze(-1) + codes = _mxfp4_encode_values(blocks / scales).reshape(-1, cols) + packed = codes[:, 0::2] | (codes[:, 1::2] << 4) + return ( + packed.reshape(*weight.shape[:-2], rows, cols // 2), + scale_bytes.reshape(*weight.shape[:-2], rows, cols // 32), + ) + + +def _dequantize_mxfp4( + packed: torch.Tensor, + scale_bytes: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + rows, packed_cols = packed.shape + cols = packed_cols * 2 + values = _mxfp4_decode_packed(packed, cols) + scales = _e8m0_bytes_to_float(scale_bytes).repeat_interleave(32, dim=1) + return (values * scales[:, :cols]).to(dtype) + + +def _has_b12x_moe() -> bool: + return ( + torch.cuda.is_available() + and current_platform.is_device_capability_family(120) + and B12xExperts._supports_current_device() + ) + + +def _count_fp4_negative_zeros(packed: torch.Tensor) -> int: + low = (packed & 0x0F) == 0x08 + high = (packed & 0xF0) == 0x80 + return int(low.sum().item() + high.sum().item()) + + +def _make_b12x_moe_kernel( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk: int, + activation: MoEActivation, + quant_config: FusedMoEQuantConfig, +) -> mk.FusedMoEKernel: + num_experts = w1.shape[0] + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=hidden_states.shape[1], + intermediate_size=w2.shape[2] * 2, + in_dtype=hidden_states.dtype, + activation=activation, + ) + experts = B12xExperts(moe_config, quant_config) + experts.process_weights_after_loading( + SimpleNamespace( + activation=activation, + w13_weight=w1, + w2_weight=w2, + ) + ) + return mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts, + ) + + +def _run_b12x_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + score: torch.Tensor, + topk: int, + activation: MoEActivation, + quant_config: FusedMoEQuantConfig, +) -> torch.Tensor: + num_experts = w1.shape[0] + kernel = _make_b12x_moe_kernel( + hidden_states, + w1, + w2, + topk, + activation, + quant_config, + ) + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + return kernel.apply( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + + +def _quant_config(weight_dtype: str, activation_dtype: str | None): + scale = torch.ones(1, dtype=torch.float32) + return FusedMoEQuantConfig.make( + quant_dtype=activation_dtype, + weight_dtype=weight_dtype, + w1_scale=scale, + w2_scale=scale, + g1_alphas=scale, + g2_alphas=scale, + a1_gscale=scale, + a2_gscale=scale, + ) + + +@pytest.mark.parametrize( + "weight_dtype,activation_dtype,mode,source_format,w13_layout", + [ + ("mxfp4", "mxfp8", "w4a8_mx", "fp4_e8m0_k32", "w31"), + ("mxfp4", None, "w4a16", "fp4_e8m0_k32", "w31"), + ("nvfp4", "nvfp4", "nvfp4", "modelopt_nvfp4", "w31"), + ("nvfp4", "mxfp8", "w4a8_nvfp4", "modelopt_nvfp4", "w31"), + ("nvfp4", None, "w4a16", "modelopt_nvfp4", "w13"), + ], +) +def test_b12x_moe_quant_mode_contract( + weight_dtype: str, + activation_dtype: str | None, + mode: str, + source_format: str, + w13_layout: str, +) -> None: + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=128, intermediate_size=64), + _quant_config(weight_dtype, activation_dtype), + ) + + assert experts._quant_mode() == mode + assert experts._source_format() == source_format + assert experts._w13_layout() == w13_layout + + +def test_b12x_moe_supports_only_tensor_parallel() -> None: + parallel = FusedMoEParallelConfig.make_no_parallel() + + assert B12xExperts._supports_parallel_config(parallel) + assert not B12xExperts._supports_parallel_config( + replace(parallel, use_ep=True, ep_size=2) + ) + all2all = replace(parallel, use_ep=True, dp_size=2) + assert all2all.use_all2all_kernels + assert not B12xExperts._supports_parallel_config(all2all) + assert not B12xExperts._supports_parallel_config( + replace(parallel, enable_eplb=True) + ) + + +def test_b12x_moe_rejects_unsupported_input_dtype( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=256, + intermediate_size=64, + in_dtype=torch.float32, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == "kernel does not support torch.float32 input/output dtype" + + +def test_b12x_moe_rejects_interleaved_swigluoai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=128, + intermediate_size=64, + activation=MoEActivation.SWIGLUOAI, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == "kernel does not support MoEActivation.SWIGLUOAI activation" + + +@pytest.mark.parametrize("activation_key", [kMxfp8Dynamic, kNvfp4Dynamic]) +def test_b12x_moe_rejects_uninterleaved_swigluoai_for_w4a8( + monkeypatch: pytest.MonkeyPatch, + activation_key, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=128, + intermediate_size=64, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + weight_key = kMxfp4Static if activation_key == kMxfp8Dynamic else kNvfp4Static + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + weight_key, + activation_key, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == "kernel does not support swigluoai_uninterleave with W4A8" + + +def test_b12x_moe_supports_uninterleaved_swigluoai_for_w4a16( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=128, + intermediate_size=64, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + + assert supported + assert reason is None + + +def test_b12x_moe_rejects_relu2_for_mxfp4_w4a8( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=256, + intermediate_size=64, + activation=MoEActivation.RELU2_NO_MUL, + ) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + kMxfp8Dynamic, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == "MXFP4 W4A8 supports only SiLU and SiTU" + + +def test_b12x_moe_rejects_unaligned_mxfp4_w4a8( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + kMxfp8Dynamic, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == ( + "MXFP4 W4A8 requires hidden size divisible by 256 and per-rank " + "intermediate size divisible by 32" + ) + + +def test_b12x_moe_rejects_mxfp4_scale_groups_crossing_tp_shards( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config(hidden_dim=128, intermediate_size=48) + config.intermediate_size_per_partition = 64 + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == ( + "MXFP4 requires the per-rank intermediate size to be divisible by 32" + ) + + +@pytest.mark.parametrize( + "beta,linear_beta", + [(3.0, 25.0), (4.0, 24.0), (None, None)], +) +def test_b12x_moe_rejects_nonstandard_situ_parameters( + monkeypatch: pytest.MonkeyPatch, + beta: float | None, + linear_beta: float | None, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=256, + intermediate_size=64, + activation=MoEActivation.SITU, + ) + config.activation_situ_beta = beta + config.activation_situ_linear_beta = linear_beta + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + + assert not supported + assert reason == "kernel supports only SiTU beta=4 and linear_beta=25" + + +def test_b12x_moe_supports_standard_situ_parameters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + hidden_dim=256, + intermediate_size=64, + activation=MoEActivation.SITU, + ) + config.activation_situ_beta = 4.0 + config.activation_situ_linear_beta = 25.0 + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + + assert supported + assert reason is None + + +@pytest.mark.parametrize( + "activation_key,force_a16,expected_backend", + [ + (kMxfp8Dynamic, False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8), + (None, False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8), + (kMxfp8Dynamic, True, Mxfp4MoeBackend.B12X_MXFP4_BF16), + (None, True, Mxfp4MoeBackend.B12X_MXFP4_BF16), + ], +) +def test_explicit_b12x_mxfp4_selection( + monkeypatch: pytest.MonkeyPatch, + activation_key, + force_a16: bool, + expected_backend: Mxfp4MoeBackend, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + monkeypatch.setattr( + mxfp4_oracle.envs, + "VLLM_B12X_MOE_FORCE_A16", + force_a16, + ) + config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64) + config.moe_backend = "b12x" + + backend, experts_cls = select_mxfp4_moe_backend( + config, + activation_key=activation_key, + ) + + assert backend == expected_backend + assert experts_cls is B12xExperts + + +@pytest.mark.parametrize( + "force_a16,expected_backend", + [ + (False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8), + (True, Mxfp4MoeBackend.B12X_MXFP4_BF16), + ], +) +def test_deepseek_v4_b12x_activation_selection( + monkeypatch: pytest.MonkeyPatch, + force_a16: bool, + expected_backend: Mxfp4MoeBackend, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr( + mxfp4_oracle.envs, + "VLLM_B12X_MOE_FORCE_A16", + force_a16, + ) + config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64) + config.moe_backend = "b12x" + + backend, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config) + + assert backend == expected_backend + assert experts_cls is B12xExperts + + +def test_b12x_mxfp4_falls_back_to_a16( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) + monkeypatch.setattr( + mxfp4_oracle.envs, + "VLLM_B12X_MOE_FORCE_A16", + False, + ) + config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64) + config.moe_backend = "b12x" + + backend, experts_cls = select_mxfp4_moe_backend(config) + + assert backend == Mxfp4MoeBackend.B12X_MXFP4_BF16 + assert experts_cls is B12xExperts + + +@pytest.mark.parametrize( + "activation_key,force_a16,expected_activation_key", + [ + (kNvfp4Dynamic, False, kNvfp4Dynamic), + (kMxfp8Dynamic, False, kMxfp8Dynamic), + (None, False, None), + (kNvfp4Dynamic, True, None), + (kMxfp8Dynamic, True, None), + ], +) +def test_explicit_b12x_nvfp4_selection( + monkeypatch: pytest.MonkeyPatch, + activation_key, + force_a16: bool, + expected_activation_key, +) -> None: + selected_activation_keys = [] + + def is_supported_config(cls, config, weight_key, activation_key, activation_format): + selected_activation_keys.append(activation_key) + return True, None + + monkeypatch.setattr(B12xExperts, "is_supported_config", is_supported_config) + monkeypatch.setattr( + nvfp4_oracle.envs, + "VLLM_B12X_MOE_FORCE_A16", + force_a16, + ) + config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64) + config.moe_backend = "b12x" + + backend, experts_cls = select_nvfp4_moe_backend( + config, + weight_key=kNvfp4Static, + activation_key=activation_key, + ) + + assert backend == NvFp4MoeBackend.B12X + assert experts_cls is B12xExperts + assert selected_activation_keys == [expected_activation_key] + + +@pytest.mark.parametrize( + "force_a16,expected_quant_dtype", [(False, "nvfp4"), (True, None)] +) +def test_b12x_nvfp4_force_a16_updates_quant_config( + monkeypatch: pytest.MonkeyPatch, + force_a16: bool, + expected_quant_dtype, +) -> None: + monkeypatch.setattr( + nvfp4_oracle.envs, + "VLLM_B12X_MOE_FORCE_A16", + force_a16, + ) + scale = torch.ones(1) + + quant_config = nvfp4_oracle.make_nvfp4_moe_quant_config( + backend=NvFp4MoeBackend.B12X, + w13_scale=scale, + w2_scale=scale, + w13_scale_2=scale, + w2_scale_2=scale, + a13_scale=scale, + a2_scale=scale, + ) + + assert quant_config.quant_dtype == expected_quant_dtype + + +def test_b12x_nvfp4_force_a16_updates_weight_preparation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + nvfp4_oracle.envs, + "VLLM_B12X_MOE_FORCE_A16", + True, + ) + reorder_w13 = None + + def prepare_for_b12x(**kwargs): + nonlocal reorder_w13 + reorder_w13 = kwargs["reorder_w13"] + return ( + kwargs["w13"], + kwargs["w13_scale"], + kwargs["w13_scale_2"], + kwargs["a13_scale"], + kwargs["w2"], + kwargs["w2_scale"], + kwargs["w2_scale_2"], + kwargs["a2_scale"], + ) + + monkeypatch.setattr( + nvfp4_oracle, + "prepare_nvfp4_moe_layer_for_b12x", + prepare_for_b12x, + ) + tensor = torch.ones(1) + + nvfp4_oracle.convert_to_nvfp4_moe_kernel_format( + nvfp4_backend=NvFp4MoeBackend.B12X, + layer=SimpleNamespace(), + w13=tensor, + w13_scale=tensor, + w13_scale_2=tensor, + a13_scale=tensor, + w2=tensor, + w2_scale=tensor, + w2_scale_2=tensor, + a2_scale=tensor, + is_act_and_mul=True, + ) + + assert reorder_w13 is True + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_b12x_nvfp4_preparation_pads_each_gated_half() -> None: + device = torch.device("cuda") + num_experts, hidden_size, intermediate_size = 2, 64, 48 + w13 = torch.ones( + num_experts, + 2 * intermediate_size, + hidden_size // 2, + dtype=torch.uint8, + device=device, + ) + w13_scale = torch.ones( + num_experts, + 2 * intermediate_size, + hidden_size // 16, + dtype=torch.float8_e4m3fn, + device=device, + ) + w2 = torch.ones( + num_experts, + hidden_size, + intermediate_size // 2, + dtype=torch.uint8, + device=device, + ) + w2_scale = torch.ones( + num_experts, + hidden_size, + intermediate_size // 16, + dtype=torch.float8_e4m3fn, + device=device, + ) + global_scale = torch.ones(num_experts, device=device) + input_scale = torch.tensor([[1.0, 2.0], [3.0, 1.0]], device=device) + + prepared = prepare_nvfp4_moe_layer_for_b12x( + w13, + w13_scale, + global_scale, + input_scale, + w2, + w2_scale, + global_scale, + input_scale, + is_act_and_mul=True, + ) + + prepared_w13, prepared_w13_scale, _, prepared_a13 = prepared[:4] + prepared_w2, prepared_w2_scale, _, prepared_a2 = prepared[4:] + assert prepared_w13.shape == (num_experts, 128, hidden_size // 2) + assert prepared_w13_scale.shape == (num_experts, 128, hidden_size // 16) + assert prepared_w2.shape == (num_experts, hidden_size, 32) + assert prepared_w2_scale.shape == (num_experts, 128, 4) + torch.testing.assert_close(prepared_a13, torch.tensor([2.0, 3.0], device=device)) + torch.testing.assert_close(prepared_a2, torch.tensor([2.0, 3.0], device=device)) + + +def test_b12x_moe_warmup_counts_cover_serving_range() -> None: + assert b12x_moe._b12x_moe_warmup_token_counts( + max_tokens=10, + token_counts=(3, 8, 12, 0), + ) == (1, 2, 3, 4, 8, 10) + + +def test_b12x_moe_uses_minimax_swiglu_parameters() -> None: + config = make_dummy_moe_config( + hidden_dim=128, + intermediate_size=64, + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ) + config.swiglu_limit = 7.0 + config.swiglu_alpha = 1.702 + config.swiglu_beta = 1.0 + experts = B12xExperts(config, _quant_config("mxfp4", None)) + + assert experts._swiglu_params(config.activation) == (7.0, 1.702, 1.0) + + +def test_b12x_moe_warmup_runs_each_planner_regime_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=128, intermediate_size=64), + _quant_config("mxfp4", None), + ) + meta = SimpleNamespace( + w1=torch.empty(0), + w2=torch.empty(0), + activation=MoEActivation.SILU, + quant_mode="w4a16", + num_experts=4, + hidden_size=128, + device=torch.device("cpu"), + dtype=torch.bfloat16, + topk=2, + apply_router_weight_on_input=False, + swiglu_limit=None, + swiglu_alpha=None, + swiglu_beta=None, + ) + prepared = SimpleNamespace() + planned_tokens = [] + launched_tokens = [] + + monkeypatch.setattr(experts, "_warmup_metadata", lambda layer: meta) + monkeypatch.setattr(experts, "_prepare_experts", lambda **kwargs: prepared) + + def fake_execution_plan(**kwargs): + tokens = kwargs["tokens"] + if tokens <= 2: + signature = ("micro", "decode") + elif tokens <= 4: + signature = ("dynamic", "small") + else: + signature = ("dynamic", "large") + return SimpleNamespace( + implementation=signature[0], + execution=signature[1], + ) + + def fake_plan(**kwargs): + planned_tokens.append(kwargs["tokens"]) + return SimpleNamespace( + scratch_specs=lambda: [SimpleNamespace(dtype=torch.uint8, shape=(64,))] + ) + + def fake_run(**kwargs): + launched_tokens.append(kwargs["hidden_states"].shape[0]) + + monkeypatch.setattr(b12x_moe, "_b12x_moe_execution_plan", fake_execution_plan) + monkeypatch.setattr(b12x_moe, "_run_b12x_moe_plan", fake_run) + monkeypatch.setattr(experts, "_plan", fake_plan) + + warmed = experts.warmup_launches( + SimpleNamespace(), + token_counts=(1, 2, 3, 4, 8), + ) + + assert warmed == 3 + assert planned_tokens == [1, 3, 8] + assert launched_tokens == planned_tokens + + +def test_b12x_moe_warmup_deduplicates_identical_experts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=128, intermediate_size=64), + _quant_config("mxfp4", None), + ) + calls = [] + + class RoutedExpertsStub(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.quant_method = SimpleNamespace( + moe_kernel=SimpleNamespace(fused_experts=experts) + ) + + class Holder(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.routed_experts = RoutedExpertsStub() + + monkeypatch.setattr( + B12xExperts, + "warmup_signature", + lambda self, layer: ("identical",), + ) + + def fake_warmup(self, layer, *, token_counts): + calls.append(tuple(token_counts)) + return 2 + + monkeypatch.setattr(B12xExperts, "warmup_launches", fake_warmup) + model = torch.nn.Sequential(Holder(), Holder()) + + warmed = b12x_moe.warmup_b12x_moe( + model, + max_tokens=4, + token_counts=(3,), + ) + + assert warmed == 2 + assert calls == [(1, 2, 3, 4)] + + +def test_b12x_moe_warmup_distinguishes_intermediate_sizes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = [] + + class RoutedExpertsStub(torch.nn.Module): + def __init__(self, intermediate_size: int) -> None: + super().__init__() + experts = B12xExperts( + make_dummy_moe_config( + hidden_dim=128, + intermediate_size=intermediate_size, + ), + _quant_config("mxfp4", None), + ) + self.quant_method = SimpleNamespace( + moe_kernel=SimpleNamespace(fused_experts=experts) + ) + self.w13_weight = torch.empty( + (4, 2 * intermediate_size, 64), + dtype=torch.uint8, + ) + self.w2_weight = torch.empty( + (4, 128, intermediate_size // 2), + dtype=torch.uint8, + ) + + class Holder(torch.nn.Module): + def __init__(self, intermediate_size: int) -> None: + super().__init__() + self.routed_experts = RoutedExpertsStub(intermediate_size) + + def fake_warmup(self, layer, *, token_counts): + calls.append(self.moe_config.intermediate_size) + return 1 + + monkeypatch.setattr(B12xExperts, "warmup_launches", fake_warmup) + model = torch.nn.Sequential(Holder(64), Holder(128)) + + warmed = b12x_moe.warmup_b12x_moe( + model, + max_tokens=4, + token_counts=(3,), + ) + + assert warmed == 2 + assert calls == [64, 128] + + +def test_b12x_source_release_preserves_prepared_storage_owner() -> None: + layer = torch.nn.Module() + for name, shape in ( + ("w13_weight", (4, 32, 16)), + ("w2_weight", (4, 64, 8)), + ("w13_weight_scale", (4, 32, 2)), + ("w2_weight_scale", (4, 64, 1)), + ): + layer.register_parameter( + name, + torch.nn.Parameter( + torch.empty(shape, dtype=torch.uint8), + requires_grad=False, + ), + ) + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=128, intermediate_size=64), + _quant_config("mxfp4", None), + ) + owner = SimpleNamespace( + w1_fp4=layer.w13_weight, + w2_fp4=layer.w2_weight, + w1_blockscale=layer.w13_weight_scale, + w2_blockscale=layer.w2_weight_scale, + ) + experts._prepared_experts = owner + owner_tensors = ( + owner.w1_fp4, + owner.w2_fp4, + owner.w1_blockscale, + owner.w2_blockscale, + ) + owner_ptrs = tuple(tensor.untyped_storage().data_ptr() for tensor in owner_tensors) + + experts._release_source_parameters(layer) + experts._release_source_parameters(layer) + + assert layer.w13_weight.numel() == 0 + assert layer.w2_weight.numel() == 0 + assert layer.w13_weight_scale.numel() == 0 + assert layer.w2_weight_scale.numel() == 0 + assert ( + tuple(tensor.untyped_storage().data_ptr() for tensor in owner_tensors) + == owner_ptrs + ) + + +def test_b12x_moe_reload_reuses_prepared_tensor_addresses() -> None: + @dataclass(frozen=True) + class Prepared: + weight: torch.Tensor + contract: tuple[str, ...] + + layer = torch.nn.Module() + previous = Prepared(torch.tensor([1.0, 2.0]), ("w4a16", "bf16")) + replacement = Prepared(torch.tensor([3.0, 4.0]), ("w4a16", "bf16")) + layer._b12x_prepared_experts = previous + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=128, intermediate_size=64), + _quant_config("mxfp4", None), + ) + weight_ptr = previous.weight.data_ptr() + + reused = experts._reuse_prepared_storage(layer, replacement) + + assert reused is previous + assert experts._prepared_experts is previous + assert layer._b12x_prepared_experts is previous + assert previous.weight.data_ptr() == weight_ptr + torch.testing.assert_close(previous.weight, replacement.weight) + + +def test_b12x_moe_rejects_router_weight_on_input_for_w4a8() -> None: + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=256, intermediate_size=64), + _quant_config("mxfp4", "mxfp8"), + ) + layer = SimpleNamespace( + activation=MoEActivation.SILU, + apply_router_weight_on_input=True, + ) + + with pytest.raises( + ValueError, + match="apply_router_weight_on_input only with W4A16", + ): + experts.process_weights_after_loading(layer) + + +def test_b12x_moe_workspace_uses_prepared_router_weight_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + experts = B12xExperts( + make_dummy_moe_config(hidden_dim=128, intermediate_size=64), + _quant_config("mxfp4", None), + ) + prepared = SimpleNamespace( + plan=SimpleNamespace(discards_source_parameters=False), + ) + layer = SimpleNamespace( + activation=MoEActivation.SILU, + apply_router_weight_on_input=True, + w13_weight=torch.empty(0), + w2_weight=torch.empty(0), + ) + monkeypatch.setattr(experts, "_prepare_experts", lambda **kwargs: prepared) + planned = [] + + def fake_plan(**kwargs): + planned.append(kwargs) + return SimpleNamespace( + scratch_specs=lambda: [SimpleNamespace(dtype=torch.uint8, shape=(64,))] + ) + + monkeypatch.setattr(experts, "_plan", fake_plan) + + experts.process_weights_after_loading(layer) + experts.workspace_shapes( + 8, + 128, + 128, + 2, + 4, + 4, + None, + MoEActivation.SILU, + ) + + assert planned == [ + { + "tokens": 8, + "topk": 2, + "activation": MoEActivation.SILU, + "apply_router_weight_on_input": True, + } + ] + + +@pytest.mark.skipif(not _has_b12x_moe(), reason="requires B12X MoE on SM120") +@pytest.mark.parametrize( + "activation", + [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL], +) +@torch.inference_mode() +def test_b12x_nvfp4_w4a16_matches_torch( + activation: MoEActivation, + workspace_init, +) -> None: + set_random_seed(7) + tokens, intermediate_size, hidden_size = 16, 128, 512 + num_experts, topk = 4, 2 + dtype = torch.bfloat16 + + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + hidden_states = ( + torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10 + ) + w1_rows = 2 * intermediate_size if activation.is_gated else intermediate_size + w1 = ( + torch.randn( + (num_experts, w1_rows, hidden_size), + device="cuda", + dtype=dtype, + ) + / 15 + ) + w2 = ( + torch.randn( + (num_experts, hidden_size, intermediate_size), + device="cuda", + dtype=dtype, + ) + / 15 + ) + w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) + w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) + unit_scale = torch.ones(num_experts, device="cuda", dtype=torch.float32) + + prepared = prepare_nvfp4_moe_layer_for_b12x( + w1_q, + w1_scale, + 1.0 / w1_global_scale, + unit_scale, + w2_q, + w2_scale, + 1.0 / w2_global_scale, + unit_scale, + is_act_and_mul=activation.is_gated, + reorder_w13=activation.is_gated, + ) + w1_b12x, w1_scale_b12x, w1_alpha = prepared[:3] + w2_b12x, w2_scale_b12x, w2_alpha = prepared[4:7] + assert _count_fp4_negative_zeros(w1_b12x) > 0 + assert _count_fp4_negative_zeros(w2_b12x) > 0 + quant_config = nvfp4_w4a16_moe_quant_config( + g1_alphas=w1_alpha, + g2_alphas=w2_alpha, + w1_scale=w1_scale_b12x, + w2_scale=w2_scale_b12x, + ) + score = torch.randn((tokens, num_experts), device="cuda", dtype=dtype) + output = _run_b12x_moe( + hidden_states, + w1_b12x, + w2_b12x, + score, + topk, + activation, + quant_config, + ) + assert _count_fp4_negative_zeros(w1_b12x) == 0 + assert _count_fp4_negative_zeros(w2_b12x) == 0 + + w1_ref = torch.empty_like(w1) + w2_ref = torch.empty_like(w2) + for expert in range(num_experts): + w1_ref[expert] = _dequantize_nvfp4_linear( + w1_q[expert], + w1_scale[expert], + w1_global_scale[expert], + dtype, + ) + w2_ref[expert] = _dequantize_nvfp4_linear( + w2_q[expert], + w2_scale[expert], + w2_global_scale[expert], + dtype, + ) + reference = torch_moe( + hidden_states, + w1_ref, + w2_ref, + score, + topk, + activation=activation, + ) + + torch.testing.assert_close(output, reference, atol=2e-1, rtol=2e-1) + cosine = torch.nn.functional.cosine_similarity( + output.flatten().float(), reference.flatten().float(), dim=0 + ) + assert cosine > 0.99, ( + f"cosine={cosine.item():.4f}, " + f"output_norm={output.float().norm().item():.4f}, " + f"reference_norm={reference.float().norm().item():.4f}" + ) + + +@pytest.mark.skipif(not _has_b12x_moe(), reason="requires B12X MoE on SM120") +@pytest.mark.parametrize( + "weight_dtype,activation_dtype", + [ + ("mxfp4", "mxfp8"), + ("nvfp4", "nvfp4"), + ("nvfp4", "mxfp8"), + ], +) +@torch.inference_mode() +def test_b12x_dynamic_fp4_modes_match_torch( + weight_dtype: str, + activation_dtype: str, + workspace_init, +) -> None: + set_random_seed(19) + tokens, intermediate_size, hidden_size = 16, 128, 512 + num_experts, topk = 4, 2 + dtype = torch.bfloat16 + + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + hidden_states = ( + torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10 + ) + w1 = ( + torch.randn( + (num_experts, 2 * intermediate_size, hidden_size), + device="cuda", + dtype=dtype, + ) + / 15 + ) + w2 = ( + torch.randn( + (num_experts, hidden_size, intermediate_size), + device="cuda", + dtype=dtype, + ) + / 15 + ) + nvfp4_input_scale = torch.full( + (num_experts,), + 1.0 / 1024.0, + device="cuda", + dtype=torch.float32, + ) + if weight_dtype == "mxfp4": + w1_q, w1_scale = _quantize_mxfp4(w1) + w2_q, w2_scale = _quantize_mxfp4(w2) + w1_ref = torch.stack( + [ + _dequantize_mxfp4(w1_q[e], w1_scale[e], dtype) + for e in range(num_experts) + ] + ) + w2_ref = torch.stack( + [ + _dequantize_mxfp4(w2_q[e], w2_scale[e], dtype) + for e in range(num_experts) + ] + ) + quant_config = FusedMoEQuantConfig.make( + quant_dtype=activation_dtype, + weight_dtype=weight_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + ) + else: + w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) + w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) + w1_ref = torch.stack( + [ + _dequantize_nvfp4_linear( + w1_q[e], + w1_scale[e], + w1_global_scale[e], + dtype, + ) + for e in range(num_experts) + ] + ) + w2_ref = torch.stack( + [ + _dequantize_nvfp4_linear( + w2_q[e], + w2_scale[e], + w2_global_scale[e], + dtype, + ) + for e in range(num_experts) + ] + ) + prepared = prepare_nvfp4_moe_layer_for_b12x( + w1_q, + w1_scale, + 1.0 / w1_global_scale, + nvfp4_input_scale, + w2_q, + w2_scale, + 1.0 / w2_global_scale, + nvfp4_input_scale, + is_act_and_mul=True, + ) + w1_q, w1_scale, w1_alpha, a1_scale = prepared[:4] + w2_q, w2_scale, w2_alpha, a2_scale = prepared[4:] + quant_config = FusedMoEQuantConfig.make( + quant_dtype=activation_dtype, + weight_dtype=weight_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + g1_alphas=w1_alpha, + g2_alphas=w2_alpha, + a1_gscale=1.0 / a1_scale, + a2_gscale=1.0 / a2_scale, + ) + + score = torch.randn((tokens, num_experts), device="cuda", dtype=dtype) + reference = torch_moe(hidden_states, w1_ref, w2_ref, score, topk) + output = _run_b12x_moe( + hidden_states, + w1_q, + w2_q, + score, + topk, + MoEActivation.SILU, + quant_config, + ) + + if activation_dtype == "nvfp4": + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + reference = _nvfp4_activation_reference( + hidden_states, + w1_ref, + w2_ref, + topk_weights, + topk_ids, + quant_config.a1_gscale, + quant_config.a2_gscale, + ) + + torch.testing.assert_close(output, reference, atol=2e-1, rtol=2e-1) + cosine = torch.nn.functional.cosine_similarity( + output.flatten().float(), + reference.flatten().float(), + dim=0, + ) + assert cosine > 0.99 + + +@pytest.mark.skipif(not _has_b12x_moe(), reason="requires B12X MoE on SM120") +@torch.inference_mode() +def test_b12x_mxfp4_w4a16_matches_torch(workspace_init) -> None: + set_random_seed(11) + tokens, intermediate_size, hidden_size = 16, 128, 512 + num_experts, topk = 4, 2 + dtype = torch.bfloat16 + + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + hidden_states = ( + torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10 + ) + w1 = ( + torch.randn( + (num_experts, 2 * intermediate_size, hidden_size), + device="cuda", + dtype=dtype, + ) + / 15 + ) + w2 = ( + torch.randn( + (num_experts, hidden_size, intermediate_size), + device="cuda", + dtype=dtype, + ) + / 15 + ) + w1_q, w1_scale = _quantize_mxfp4(w1) + w2_q, w2_scale = _quantize_mxfp4(w2) + w1_ref = torch.empty_like(w1) + w2_ref = torch.empty_like(w2) + for expert in range(num_experts): + w1_ref[expert] = _dequantize_mxfp4(w1_q[expert], w1_scale[expert], dtype) + w2_ref[expert] = _dequantize_mxfp4(w2_q[expert], w2_scale[expert], dtype) + quant_config = mxfp4_w4a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + ) + score = torch.randn((tokens, num_experts), device="cuda", dtype=dtype) + reference = torch_moe(hidden_states, w1_ref, w2_ref, score, topk) + output = _run_b12x_moe( + hidden_states, + w1_q, + w2_q, + score, + topk, + MoEActivation.SILU, + quant_config, + ) + + torch.testing.assert_close(output, reference, atol=2e-1, rtol=2e-1) + cosine = torch.nn.functional.cosine_similarity( + output.flatten().float(), reference.flatten().float(), dim=0 + ) + assert cosine > 0.99, ( + f"cosine={cosine.item():.4f}, " + f"output_norm={output.float().norm().item():.4f}, " + f"reference_norm={reference.float().norm().item():.4f}" + ) + + +@pytest.mark.skipif(not _has_b12x_moe(), reason="requires B12X MoE on SM120") +@pytest.mark.parametrize( + "weight_dtype,activation_dtype", + [ + ("mxfp4", None), + ("mxfp4", "mxfp8"), + ("nvfp4", None), + ("nvfp4", "nvfp4"), + ("nvfp4", "mxfp8"), + ], +) +@torch.inference_mode() +def test_b12x_moe_cuda_graph_replay( + weight_dtype: str, + activation_dtype: str | None, + workspace_init, +) -> None: + from vllm.v1.worker.workspace import lock_workspace + + set_random_seed(23) + tokens = 4 + if weight_dtype == "nvfp4" and activation_dtype is not None: + intermediate_size, hidden_size = 1024, 4096 + else: + intermediate_size, hidden_size = 128, 512 + num_experts, topk = 4, 2 + hidden_states = ( + torch.randn( + (tokens, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + / 10 + ) + w1 = ( + torch.randn( + (num_experts, 2 * intermediate_size, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + w2 = ( + torch.randn( + (num_experts, hidden_size, intermediate_size), + device="cuda", + dtype=torch.bfloat16, + ) + / 15 + ) + if weight_dtype == "mxfp4": + w1_q, w1_scale = _quantize_mxfp4(w1) + w2_q, w2_scale = _quantize_mxfp4(w2) + if activation_dtype is None: + quant_config = mxfp4_w4a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + ) + else: + quant_config = FusedMoEQuantConfig.make( + quant_dtype=activation_dtype, + weight_dtype=weight_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + ) + else: + w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1) + w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2) + input_scale = torch.full( + (num_experts,), + 1.0 if activation_dtype is None else 1.0 / 1024.0, + device="cuda", + dtype=torch.float32, + ) + prepared = prepare_nvfp4_moe_layer_for_b12x( + w1_q, + w1_scale, + 1.0 / w1_global_scale, + input_scale, + w2_q, + w2_scale, + 1.0 / w2_global_scale, + input_scale, + is_act_and_mul=True, + reorder_w13=activation_dtype is None, + ) + w1_q, w1_scale, w1_alpha, a1_scale = prepared[:4] + w2_q, w2_scale, w2_alpha, a2_scale = prepared[4:] + if activation_dtype is None: + quant_config = nvfp4_w4a16_moe_quant_config( + g1_alphas=w1_alpha, + g2_alphas=w2_alpha, + w1_scale=w1_scale, + w2_scale=w2_scale, + ) + else: + quant_config = FusedMoEQuantConfig.make( + quant_dtype=activation_dtype, + weight_dtype=weight_dtype, + w1_scale=w1_scale, + w2_scale=w2_scale, + g1_alphas=w1_alpha, + g2_alphas=w2_alpha, + a1_gscale=1.0 / a1_scale, + a2_gscale=1.0 / a2_scale, + ) + + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + kernel = _make_b12x_moe_kernel( + hidden_states, + w1_q, + w2_q, + topk, + MoEActivation.SILU, + quant_config, + ) + score = torch.randn( + (tokens, num_experts), + device="cuda", + dtype=torch.bfloat16, + ) + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score, topk, renormalize=False + ) + assert topk_weights.dtype == torch.float32 and topk_weights.is_contiguous() + assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous() + + def apply() -> torch.Tensor: + return kernel.apply( + hidden_states=hidden_states, + w1=w1_q, + w2=w2_q, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=num_experts, + expert_map=None, + apply_router_weight_on_input=False, + ) + + expected = apply().clone() + lock_workspace() + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + with torch.cuda.graph(graph, stream=stream): + actual = apply() + graph.replay() + torch.accelerator.synchronize() + + assert torch.isfinite(expected).all() + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 86bd8579239b..c78e5b626072 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -665,7 +665,13 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel( ) -> None: captured = {} expected_quant_config = object() - expected_kernel = object() + expected_kernel = SimpleNamespace( + fused_experts=SimpleNamespace( + process_weights_after_loading=lambda layer: captured.update( + {"processed_layer": layer} + ) + ) + ) expected_experts_cls = object() monkeypatch.setattr( @@ -728,6 +734,7 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel( assert captured["kernel_kwargs"]["moe_config"] is expected_moe_config assert captured["kernel_kwargs"]["experts_cls"] is expected_experts_cls assert captured["kernel_kwargs"]["routing_tables"] == "routing-tables" + assert captured["processed_layer"] is layer assert method.moe_kernel is expected_kernel diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 61bbfe33e8eb..dcf07bbcb326 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -130,6 +130,7 @@ def with_default( "flashinfer_cutlass", "flashinfer_cutedsl", "flashinfer_b12x", + "b12x", "marlin", "humming", "triton_unfused", @@ -141,13 +142,13 @@ def with_default( LinearBackend = Literal[ "auto", - "b12x", "cutlass", "flashinfer_cutlass", "flashinfer_cutedsl", "flashinfer_trtllm", "flashinfer_cudnn", "flashinfer_b12x", + "b12x", "marlin", "humming", "triton", @@ -193,7 +194,6 @@ class KernelConfig: """Backend for MoE expert computation kernels. Available options: - "auto": Automatically select the best backend based on model and hardware - - "b12x": Use B12X kernels for SM12x FP8 and FP4 linear layers - "triton": Use Triton-based fused MoE kernels - "batched_triton": Use batched Triton experts (moe_mmk) on the batched activation format ([E_local, max_num_tokens, K]) @@ -205,6 +205,7 @@ class KernelConfig: - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) + - "b12x": Use native B12X FP4 MoE kernels on SM12x - "marlin": Use Marlin kernels (weight-only quantization) - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels @@ -225,6 +226,7 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+) + - "b12x": Use native B12X FP8 and FP4 linear kernels on SM12x - "marlin": Use Marlin kernels - "triton": Use Triton-based kernels - "deep_gemm": Use DeepGEMM kernels diff --git a/vllm/envs.py b/vllm/envs.py index 657c1cd21068..17b64ee96a40 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -184,6 +184,7 @@ VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None VLLM_USE_B12X_FP8_GEMM: bool = False VLLM_USE_B12X_FP4_GEMM: bool = False + VLLM_B12X_MOE_FORCE_A16: bool = False VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False VLLM_TPU_USING_PATHWAYS: bool = False @@ -1577,6 +1578,10 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_USE_B12X_FP4_GEMM": lambda: bool( int(os.getenv("VLLM_USE_B12X_FP4_GEMM", "0")) ), + # Force B12X FP4 MoE to use BF16 activations. + "VLLM_B12X_MOE_FORCE_A16": lambda: bool( + int(os.getenv("VLLM_B12X_MOE_FORCE_A16", "0")) + ), # Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops. "VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool( int(os.getenv("VLLM_USE_FLASHINFER_MOE_INT4", "0")) diff --git a/vllm/model_executor/layers/fused_moe/b12x_moe.py b/vllm/model_executor/layers/fused_moe/b12x_moe.py new file mode 100644 index 000000000000..cd96bae85979 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/b12x_moe.py @@ -0,0 +1,871 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""B12X modular tensor-parallel fused MoE backend.""" + +from collections.abc import Iterable +from types import SimpleNamespace +from typing import Any, cast + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.kernels.b12x_utils import reuse_packed_weight_storage +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, + kMxfp8Dynamic, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def _b12x_activation_name(activation: MoEActivation) -> str: + if activation == MoEActivation.SILU: + return "silu" + if activation in (MoEActivation.RELU2, MoEActivation.RELU2_NO_MUL): + return "relu2" + return activation.value + + +def _b12x_scratch_nbytes(plan: Any) -> int: + specs = plan.scratch_specs() + if len(specs) != 1: + raise RuntimeError(f"expected one B12X MoE scratch buffer, got {len(specs)}") + spec = specs[0] + if spec.dtype != torch.uint8: + raise TypeError(f"expected B12X MoE scratch dtype uint8, got {spec.dtype}") + return int(spec.shape[0]) + + +def _b12x_moe_warmup_token_counts( + *, + max_tokens: int, + token_counts: Iterable[int] = (), +) -> tuple[int, ...]: + """Return powers of two plus serving sizes supplied by vLLM.""" + max_tokens = max(int(max_tokens), 1) + counts = { + int(token_count) + for token_count in token_counts + if 0 < int(token_count) <= max_tokens + } + token_count = 1 + while token_count < max_tokens: + counts.add(token_count) + token_count *= 2 + counts.add(max_tokens) + return tuple(sorted(counts)) + + +def _b12x_moe_execution_plan( + *, + tokens: int, + topk: int, + prepared: Any, + quant_mode: str, + apply_router_weight_on_input: bool, + swiglu_limit: float | None, + swiglu_alpha: float | None, + swiglu_beta: float | None, +) -> Any: + from b12x.moe import fused_moe + + return fused_moe.plan_execution( + num_tokens=max(int(tokens), 1), + num_topk=int(topk), + device=prepared.w1_fp4.device, + weight_plan=prepared.plan, + quant_mode=quant_mode, + apply_router_weight_on_input=apply_router_weight_on_input, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + + +def _run_b12x_moe_plan( + *, + plan: Any, + scratch: torch.Tensor, + hidden_states: torch.Tensor, + prepared: Any, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + output: torch.Tensor, + unit_scale_contract: bool, +) -> None: + from b12x.moe import fused_moe + + binding = fused_moe.bind( + plan, + scratch=scratch, + a=hidden_states, + experts=prepared, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + input_scales_static=True, + unit_scale_contract=unit_scale_contract, + ) + fused_moe.run(binding=binding) + + +def _is_current_stream_capturing() -> bool: + is_capturing = getattr(torch.cuda, "is_current_stream_capturing", None) + return bool(is_capturing is not None and is_capturing()) + + +def _normalize_topk_ids(topk_ids: torch.Tensor) -> torch.Tensor: + if topk_ids.dtype == torch.int32 and topk_ids.is_contiguous(): + return topk_ids + if _is_current_stream_capturing(): + raise RuntimeError( + "B12X MoE topk_ids normalization would allocate during CUDA capture" + ) + return topk_ids.to(dtype=torch.int32).contiguous() + + +def _normalize_topk_weights(topk_weights: torch.Tensor) -> torch.Tensor: + if topk_weights.dtype == torch.float32 and topk_weights.is_contiguous(): + return topk_weights + if _is_current_stream_capturing(): + raise RuntimeError( + "B12X MoE topk_weights normalization would allocate during CUDA capture" + ) + return topk_weights.to(dtype=torch.float32).contiguous() + + +def _workspace_as_b12x_scratch( + workspace: torch.Tensor | None, + plan: Any, +) -> torch.Tensor: + if workspace is None: + raise RuntimeError("B12X MoE requires workspace2 scratch") + if not workspace.is_contiguous(): + raise ValueError("B12X MoE workspace2 must be contiguous") + scratch = workspace.view(-1).view(torch.uint8) + required_nbytes = _b12x_scratch_nbytes(plan) + if scratch.numel() < required_nbytes: + raise ValueError( + "B12X MoE workspace2 is too small: " + f"have={scratch.numel()} bytes, need={required_nbytes} bytes" + ) + return scratch + + +def _replace_parameter_with_empty( + layer: torch.nn.Module, + name: str, +) -> torch.Tensor | None: + parameter = getattr(layer, name, None) + if not isinstance(parameter, torch.Tensor): + return None + empty = torch.empty((0,), dtype=parameter.dtype, device=parameter.device) + replace_parameter(layer, name, empty) + return getattr(layer, name) + + +def _set_quant_config_scale( + quant_config: FusedMoEQuantConfig, + descriptor_name: str, + scale: torch.Tensor, +) -> None: + descriptor = getattr(quant_config, descriptor_name) + descriptor.scale = scale + + +def _normalize_expert_scale(scale: torch.Tensor) -> torch.Tensor: + if scale.ndim == 2: + if scale.shape[1] not in (1, 2): + raise ValueError( + "expected an expert scale with one or two columns, got " + f"{tuple(scale.shape)}" + ) + scale = scale[:, 0] + return scale.to(dtype=torch.float32).contiguous() + + +def _canonicalize_fp4_zero_signs_(packed: torch.Tensor) -> None: + """Clear sign bits from packed FP4 zero values in place.""" + packed = packed.view(torch.uint8) + magnitude = packed & 0x77 + nonzero = (magnitude | (magnitude >> 1) | (magnitude >> 2)) & 0x11 + packed.bitwise_and_(0x77 | (nonzero << 3)) + + +class B12xExperts(mk.FusedMoEExpertsModular): + """FP4 MoE experts backed by the B12X SM12x planned API.""" + + def __init__( + self, + moe_config: mk.FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + if quant_config.weight_quant_dtype not in ("mxfp4", "nvfp4"): + raise ValueError( + "B12X MoE requires MXFP4 or NVFP4 weights, got " + f"{quant_config.weight_quant_dtype}" + ) + self._prepared_experts: Any | None = None + self._source_parameters_released = False + self._unit_scales: dict[torch.device, torch.Tensor] = {} + self._plans: dict[tuple[int, int, MoEActivation, bool], Any] = {} + self._apply_router_weight_on_input = False + + def _quant_mode(self) -> str: + scheme: tuple[str, str | None] = ( + cast(str, self.quant_config.weight_quant_dtype), + cast(str | None, self.quant_config.quant_dtype), + ) + modes = { + ("mxfp4", "mxfp8"): "w4a8_mx", + ("mxfp4", None): "w4a16", + ("nvfp4", "nvfp4"): "nvfp4", + ("nvfp4", "mxfp8"): "w4a8_nvfp4", + ("nvfp4", None): "w4a16", + } + try: + return modes[scheme] + except KeyError as exc: + raise ValueError( + f"unsupported B12X MoE quantization scheme {scheme}" + ) from exc + + def _source_format(self) -> str: + if self.quant_config.weight_quant_dtype == "nvfp4": + return "modelopt_nvfp4" + return "fp4_e8m0_k32" + + def _w13_layout(self) -> str: + if self._source_format() == "modelopt_nvfp4" and self._quant_mode() == "w4a16": + return "w13" + return "w31" + + def _unit_scale(self, device: torch.device, num_experts: int) -> torch.Tensor: + scale = self._unit_scales.get(device) + if scale is None or scale.numel() != num_experts: + scale = torch.ones(num_experts, dtype=torch.float32, device=device) + self._unit_scales[device] = scale + return scale + + def _weight_global_scale( + self, + device: torch.device, + num_experts: int, + scale: torch.Tensor | None, + name: str, + ) -> torch.Tensor: + if self._source_format() != "modelopt_nvfp4": + return self._unit_scale(device, num_experts) + if scale is None: + raise ValueError(f"B12X NVFP4 MoE requires {name}") + scale = _normalize_expert_scale(scale) + if scale.numel() != num_experts: + raise ValueError( + f"B12X NVFP4 MoE expected {num_experts} {name} values, " + f"got {scale.numel()}" + ) + return scale.to(device=device) + + def _swiglu_params( + self, + activation: MoEActivation, + ) -> tuple[float | None, float | None, float | None]: + if activation in ( + MoEActivation.SITU, + MoEActivation.RELU2, + MoEActivation.RELU2_NO_MUL, + ): + return None, None, None + + limit = self.quant_config.gemm1_clamp_limit + if limit is None: + limit = self.moe_config.swiglu_limit + if activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE: + return limit, None, None + + alpha = self.quant_config.gemm1_alpha + if alpha is None: + alpha = self.moe_config.swiglu_alpha + beta = self.quant_config.gemm1_beta + if beta is None: + beta = self.moe_config.swiglu_beta + return limit, alpha, beta + + def _prepare_experts( + self, + *, + w1: torch.Tensor, + w2: torch.Tensor, + activation: MoEActivation, + params_dtype: torch.dtype, + ) -> Any: + quant_mode = self._quant_mode() + if self._prepared_experts is not None: + plan = self._prepared_experts.plan + requested_dtype = str(params_dtype).removeprefix("torch.") + if ( + quant_mode in plan.quant_modes + and requested_dtype == plan.io_dtype + and _b12x_activation_name(activation) == plan.activation + ): + return self._prepared_experts + raise RuntimeError("B12X MoE prepared weights do not match this invocation") + if self._source_parameters_released: + raise RuntimeError("B12X MoE source parameters were already released") + if _is_current_stream_capturing(): + raise RuntimeError( + "B12X MoE weights must be prepared before CUDA graph capture" + ) + if self.w1_scale is None or self.w2_scale is None: + raise ValueError("B12X MoE requires w1 and w2 block scales") + + _canonicalize_fp4_zero_signs_(w1) + _canonicalize_fp4_zero_signs_(w2) + + from b12x.moe import fused_moe + + num_experts = int(w1.shape[0]) + hidden_size = int(w2.shape[1]) + intermediate_size = int(w2.shape[2]) * 2 + unit_scale = self._unit_scale(w1.device, num_experts) + w1_global_scale = self._weight_global_scale( + w1.device, num_experts, self.g1_alphas, "w1 global scales" + ) + w2_global_scale = self._weight_global_scale( + w2.device, num_experts, self.g2_alphas, "w2 global scales" + ) + + if quant_mode in ("nvfp4", "w4a8_nvfp4"): + if self.a1_gscale is None or self.a2_gscale is None: + raise ValueError("B12X NVFP4 MoE requires activation global scales") + a1_gscale = _normalize_expert_scale(self.a1_gscale).to(w1.device) + a2_gscale = _normalize_expert_scale(self.a2_gscale).to(w2.device) + else: + a1_gscale = unit_scale + a2_gscale = unit_scale + + weight_plan = fused_moe.plan_weights( + quant_modes=quant_mode, + source_format=self._source_format(), + activation=_b12x_activation_name(activation), + params_dtype=params_dtype, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + w13_layout=self._w13_layout(), + ) + self._prepared_experts = fused_moe.prepare_weights( + plan=weight_plan, + w1_fp4=w1, + w1_blockscale=self.w1_scale, + w1_global_scale=w1_global_scale, + a1_gscale=a1_gscale, + w2_fp4=w2, + w2_blockscale=self.w2_scale, + w2_global_scale=w2_global_scale, + a2_gscale=a2_gscale, + params_dtype=params_dtype, + ) + return self._prepared_experts + + def _release_source_parameters(self, layer: torch.nn.Module) -> None: + if self._source_parameters_released: + return + w1_scale = _replace_parameter_with_empty(layer, "w13_weight_scale") + w2_scale = _replace_parameter_with_empty(layer, "w2_weight_scale") + if w1_scale is not None: + _set_quant_config_scale(self.quant_config, "_w1", w1_scale) + if w2_scale is not None: + _set_quant_config_scale(self.quant_config, "_w2", w2_scale) + _replace_parameter_with_empty(layer, "w13_weight") + _replace_parameter_with_empty(layer, "w2_weight") + self._source_parameters_released = True + + def _reuse_prepared_storage(self, layer: torch.nn.Module, prepared: Any) -> Any: + previous = getattr(layer, "_b12x_prepared_experts", None) + prepared = reuse_packed_weight_storage(previous, prepared) + self._prepared_experts = prepared + layer._b12x_prepared_experts = prepared + return prepared + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self._apply_router_weight_on_input = bool( + getattr(layer, "apply_router_weight_on_input", False) + ) + if self._apply_router_weight_on_input and self._quant_mode() != "w4a16": + raise ValueError( + "B12X MoE supports apply_router_weight_on_input only with W4A16" + ) + activation = getattr(layer, "activation", self.moe_config.activation) + if isinstance(activation, str): + activation = MoEActivation.from_str(activation) + prepared = self._prepare_experts( + w1=layer.w13_weight, + w2=layer.w2_weight, + activation=cast(MoEActivation, activation), + params_dtype=self.moe_config.in_dtype, + ) + prepared = self._reuse_prepared_storage(layer, prepared) + if prepared.plan.discards_source_parameters: + self._release_source_parameters(layer) + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: mk.FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + if moe_config.has_bias: + return False, "kernel does not support expert biases" + if moe_config.in_dtype not in (torch.float16, torch.bfloat16): + return ( + False, + f"kernel does not support {moe_config.in_dtype} input/output dtype", + ) + if moe_config.activation == MoEActivation.SITU and ( + moe_config.activation_situ_beta != 4.0 + or moe_config.activation_situ_linear_beta != 25.0 + ): + return False, "kernel supports only SiTU beta=4 and linear_beta=25" + if ( + activation_key is not None + and moe_config.activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE + ): + return ( + False, + "kernel does not support swigluoai_uninterleave with W4A8", + ) + unpadded_intermediate_size = ( + moe_config.intermediate_size_per_partition_unpadded + or moe_config.intermediate_size_per_partition + ) + if weight_key == kMxfp4Static and unpadded_intermediate_size % 32 != 0: + return ( + False, + "MXFP4 requires the per-rank intermediate size to be divisible by 32", + ) + if weight_key == kMxfp4Static and activation_key == kMxfp8Dynamic: + if moe_config.activation not in ( + MoEActivation.SILU, + MoEActivation.SITU, + ): + return False, "MXFP4 W4A8 supports only SiLU and SiTU" + if ( + moe_config.hidden_dim % 256 != 0 + or moe_config.intermediate_size_per_partition % 32 != 0 + ): + return ( + False, + "MXFP4 W4A8 requires hidden size divisible by 256 and " + "per-rank intermediate size divisible by 32", + ) + return mk.FusedMoEExperts.is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + + @staticmethod + def _supports_current_device() -> bool: + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(120) + ): + return False + try: + from b12x.moe import fused_moe + except ImportError: + return False + return fused_moe.is_supported() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return True + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) in ( + (kMxfp4Static, kMxfp8Dynamic), + (kMxfp4Static, None), + (kNvfp4Static, kNvfp4Dynamic), + (kNvfp4Static, kMxfp8Dynamic), + (kNvfp4Static, None), + ) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in ( + MoEActivation.SILU, + MoEActivation.SITU, + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + MoEActivation.RELU2_NO_MUL, + ) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_ep + and moe_parallel_config.ep_size == 1 + and not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def _prepared(self) -> Any: + if self._prepared_experts is None: + raise RuntimeError( + "B12X MoE weights must be prepared before workspace planning" + ) + return self._prepared_experts + + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + if w1.numel() and w2.numel(): + return super().moe_problem_size(a1, w1, w2, topk_ids) + prepared = self._prepared() + tokens = int(a1.shape[0] if a1.ndim == 2 else a1.shape[1]) + return ( + int(prepared.num_experts), + tokens, + int(prepared.intermediate_size) * 2, + int(a1.shape[-1]), + int(topk_ids.shape[1]), + ) + + def _plan( + self, + *, + tokens: int, + topk: int, + activation: MoEActivation, + apply_router_weight_on_input: bool = False, + ) -> Any: + from b12x.moe import fused_moe + + key = ( + max(int(tokens), 1), + int(topk), + activation, + bool(apply_router_weight_on_input), + ) + plan = self._plans.get(key) + if plan is not None: + return plan + if _is_current_stream_capturing(): + raise RuntimeError("B12X MoE plans must be created before CUDA capture") + + limit, alpha, beta = self._swiglu_params(activation) + prepared = self._prepared() + plan = fused_moe.plan( + fused_moe.Caps( + max_tokens=key[0], + num_topk=key[1], + device=prepared.w1_fp4.device, + weight_plan=prepared.plan, + core_token_counts=(key[0],), + route_num_experts=0, + quant_mode=self._quant_mode(), + apply_router_weight_on_input=key[3], + swiglu_limit=limit, + swiglu_alpha=alpha, + swiglu_beta=beta, + frozen=True, + ) + ) + self._plans[key] = plan + return plan + + def _warmup_metadata(self, layer: torch.nn.Module) -> SimpleNamespace | None: + w1 = getattr(layer, "w13_weight", None) + w2 = getattr(layer, "w2_weight", None) + if not isinstance(w1, torch.Tensor) or not isinstance(w2, torch.Tensor): + return None + + activation = getattr(layer, "activation", self.moe_config.activation) + if isinstance(activation, str): + activation = MoEActivation.from_str(activation) + activation = cast(MoEActivation, activation) + prepared = self._prepared_experts + if (w1.numel() == 0 or w2.numel() == 0) and prepared is None: + return None + if prepared is not None: + num_experts = int(prepared.num_experts) + hidden_size = int(prepared.hidden_size) + intermediate_size = int(prepared.intermediate_size) + device = prepared.w1_fp4.device + else: + num_experts = int(w1.shape[0]) + hidden_size = int(w2.shape[1]) + intermediate_size = int(w2.shape[2]) * 2 + device = w1.device + limit, alpha, beta = self._swiglu_params(activation) + return SimpleNamespace( + w1=w1, + w2=w2, + activation=activation, + activation_name=_b12x_activation_name(activation), + quant_mode=self._quant_mode(), + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + dtype=self.moe_config.in_dtype, + topk=int(self.moe_config.experts_per_token), + apply_router_weight_on_input=bool( + getattr(layer, "apply_router_weight_on_input", False) + ), + swiglu_limit=limit, + swiglu_alpha=alpha, + swiglu_beta=beta, + ) + + def warmup_signature(self, layer: torch.nn.Module) -> tuple[Any, ...] | None: + meta = self._warmup_metadata(layer) + if meta is None: + return None + return ( + meta.device.type, + meta.device.index, + meta.dtype, + meta.quant_mode, + self._source_format(), + self._w13_layout(), + meta.num_experts, + meta.hidden_size, + meta.intermediate_size, + meta.topk, + meta.activation_name, + meta.apply_router_weight_on_input, + meta.swiglu_limit, + meta.swiglu_alpha, + meta.swiglu_beta, + ) + + @torch.inference_mode() + def warmup_launches( + self, + layer: torch.nn.Module, + *, + token_counts: Iterable[int], + ) -> int: + """Compile one representative launch for every planned regime.""" + meta = self._warmup_metadata(layer) + if meta is None: + return 0 + prepared = self._prepare_experts( + w1=meta.w1, + w2=meta.w2, + activation=meta.activation, + params_dtype=meta.dtype, + ) + launch_tokens: dict[tuple[Any, ...], int] = {} + for tokens in sorted({int(count) for count in token_counts if int(count) > 0}): + execution_plan = _b12x_moe_execution_plan( + tokens=tokens, + topk=meta.topk, + prepared=prepared, + quant_mode=meta.quant_mode, + apply_router_weight_on_input=meta.apply_router_weight_on_input, + swiglu_limit=meta.swiglu_limit, + swiglu_alpha=meta.swiglu_alpha, + swiglu_beta=meta.swiglu_beta, + ) + signature = (execution_plan.implementation, execution_plan.execution) + launch_tokens.setdefault(signature, tokens) + + for tokens in launch_tokens.values(): + hidden_states = torch.zeros( + (tokens, meta.hidden_size), + dtype=meta.dtype, + device=meta.device, + ) + output = torch.empty_like(hidden_states) + topk_ids = ( + torch.arange(meta.topk, device=meta.device, dtype=torch.int32) + .unsqueeze(0) + .expand(tokens, -1) + .contiguous() + ) + topk_ids.remainder_(meta.num_experts) + topk_weights = torch.full( + (tokens, meta.topk), + 1.0 / meta.topk, + dtype=torch.float32, + device=meta.device, + ) + plan = self._plan( + tokens=tokens, + topk=meta.topk, + activation=meta.activation, + apply_router_weight_on_input=meta.apply_router_weight_on_input, + ) + scratch = torch.empty( + (_b12x_scratch_nbytes(plan),), + dtype=torch.uint8, + device=meta.device, + ) + _run_b12x_moe_plan( + plan=plan, + scratch=scratch, + hidden_states=hidden_states, + prepared=prepared, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + unit_scale_contract=meta.quant_mode == "w4a16", + ) + return len(launch_tokens) + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + del N, global_num_experts, local_num_experts, expert_tokens_meta + plan = self._plan( + tokens=M, + topk=topk, + activation=activation, + apply_router_weight_on_input=self._apply_router_weight_on_input, + ) + itemsize = self.moe_config.in_dtype.itemsize + scratch_elements = max( + 1, (_b12x_scratch_nbytes(plan) + itemsize - 1) // itemsize + ) + return (0,), (scratch_elements,), (M, K) + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor | None, + workspace2: torch.Tensor | None, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool | None, + ) -> None: + del global_num_experts, a1q_scale, a2_scale, workspace13, expert_tokens_meta + if expert_map is not None: + raise ValueError("B12X TP MoE does not support expert maps") + if bool(apply_router_weight_on_input) != self._apply_router_weight_on_input: + raise ValueError( + "apply_router_weight_on_input does not match the prepared B12X MoE plan" + ) + prepared = self._prepare_experts( + w1=w1, + w2=w2, + activation=activation, + params_dtype=hidden_states.dtype, + ) + topk_ids = _normalize_topk_ids(topk_ids) + topk_weights = _normalize_topk_weights(topk_weights) + plan = self._plan( + tokens=int(hidden_states.shape[0]), + topk=int(topk_ids.shape[1]), + activation=activation, + apply_router_weight_on_input=bool(apply_router_weight_on_input), + ) + scratch = _workspace_as_b12x_scratch(workspace2, plan) + + _run_b12x_moe_plan( + plan=plan, + scratch=scratch, + hidden_states=hidden_states, + prepared=prepared, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output, + unit_scale_contract=self._quant_mode() == "w4a16", + ) + + def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: + raise NotImplementedError("LoRA is not supported for B12xExperts") + + +def warmup_b12x_moe( + model: torch.nn.Module, + *, + max_tokens: int, + token_counts: Iterable[int] = (), +) -> int: + """Warm unique B12X MoE planner regimes in a loaded model.""" + candidates = _b12x_moe_warmup_token_counts( + max_tokens=max_tokens, + token_counts=token_counts, + ) + seen: set[tuple[Any, ...]] = set() + warmed = 0 + for module in model.modules(): + routed_experts = getattr(module, "routed_experts", None) + quant_method = getattr(routed_experts, "quant_method", None) + moe_kernel = getattr(quant_method, "moe_kernel", None) + fused_experts = getattr(moe_kernel, "fused_experts", None) + if not isinstance(fused_experts, B12xExperts): + continue + signature = fused_experts.warmup_signature(routed_experts) + if signature is None or signature in seen: + continue + seen.add(signature) + warmed += fused_experts.warmup_launches( + routed_experts, + token_counts=candidates, + ) + if warmed: + logger.info( + "Warmed up %d B12X MoE launch variant(s) across %d expert signature(s).", + warmed, + len(seen), + ) + return warmed diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 054240271fb4..eee0a7d4fa32 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -865,6 +865,8 @@ def nvfp4_w4a16_moe_quant_config( g2_alphas: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: """ @@ -877,6 +879,8 @@ def nvfp4_w4a16_moe_quant_config( g1_alphas=g1_alphas, g2_alphas=g2_alphas, weight_dtype="nvfp4", + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3722528c2126..d62c9aba436b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -6,6 +6,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import envs from vllm.config import get_current_vllm_config from vllm.config.kernel import MoEBackend from vllm.config.quantization import QuantizationConfigArgs @@ -100,6 +101,8 @@ def _pack_deepgemm_mxfp4_scales( class Mxfp4MoeBackend(Enum): NONE = "None" + B12X_MXFP4_BF16 = "B12X_MXFP4_BF16" + B12X_MXFP4_MXFP8 = "B12X_MXFP4_MXFP8" # DeepGEMM FP8xFP4 backend (SM100+) DEEPGEMM_MXFP4 = "DEEPGEMM_MXFP4" # FlashInfer TRTLLM backends @@ -149,11 +152,21 @@ class Mxfp4MoeBackend(Enum): Mxfp4MoeBackend.TRITON_UNFUSED, ) +B12X_BACKENDS = ( + Mxfp4MoeBackend.B12X_MXFP4_MXFP8, + Mxfp4MoeBackend.B12X_MXFP4_BF16, +) + def backend_to_kernel_cls( backend: Mxfp4MoeBackend, ) -> list[type[mk.FusedMoEExperts]]: - if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: + if backend in B12X_BACKENDS: + from vllm.model_executor.layers.fused_moe.b12x_moe import B12xExperts + + return [B12xExperts] + + elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( DeepGemmFP4Experts, ) @@ -277,6 +290,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: via ``activation_key`` and ``is_supported_config``. """ mapping: dict[str, list[Mxfp4MoeBackend]] = { + "b12x": list(B12X_BACKENDS), "deep_gemm": [Mxfp4MoeBackend.DEEPGEMM_MXFP4], "flashinfer_trtllm": [ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, @@ -363,6 +377,8 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: """Map backend to its activation key (FP8, MXFP8, or None for BF16).""" if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: return kFp8Dynamic128Sym + if backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8: + return kMxfp8Dynamic if backend in ( Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, @@ -449,6 +465,19 @@ def _filter_by_activation( return bf16 if bf16 else backends +def _get_requested_backends( + runner_backend: MoEBackend, + requested_activation_key: QuantKey | None, +) -> list[Mxfp4MoeBackend]: + backends = map_mxfp4_backend(runner_backend) + if runner_backend == "b12x": + if envs.VLLM_B12X_MOE_FORCE_A16: + return [Mxfp4MoeBackend.B12X_MXFP4_BF16] + if requested_activation_key is None: + return backends + return _filter_by_activation(backends, requested_activation_key) + + def select_mxfp4_moe_backend( config: FusedMoEConfig, activation_key: QuantKey | None = None, @@ -464,6 +493,7 @@ def select_mxfp4_moe_backend( Note: Shape-specific fallbacks may still occur at runtime. """ + runner_backend = config.moe_backend requested_activation_key = _resolve_activation_key(activation_key) activation_format = ( @@ -472,23 +502,22 @@ def select_mxfp4_moe_backend( else mk.FusedMoEActivationFormat.Standard ) - runner_backend = config.moe_backend if runner_backend != "auto": - requested_backends = map_mxfp4_backend(runner_backend) + requested_backends = _get_requested_backends( + runner_backend, requested_activation_key + ) if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: requested_backends = [ Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b for b in requested_backends ] - candidates = _filter_by_activation(requested_backends, requested_activation_key) - if not candidates: + if not requested_backends: raise ValueError( f"moe_backend={runner_backend!r} does not support " - f"activation={requested_activation_key}; supported variants: " - f"{[b.name for b in requested_backends]}" + f"activation={requested_activation_key}" ) last_error: Exception | None = None - for requested_backend in candidates: + for requested_backend in requested_backends: act_key = ( requested_activation_key if requested_activation_key is not None @@ -585,7 +614,11 @@ def select_deepseek_v4_mxfp4_moe_backend( # falling back to the auto priority list. runner_backend = config.moe_backend if runner_backend != "auto": - requested_backends = map_mxfp4_backend(runner_backend) + requested_backends = ( + _get_requested_backends(runner_backend, None) + if runner_backend == "b12x" + else map_mxfp4_backend(runner_backend) + ) if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: requested_backends = [ Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b @@ -644,6 +677,8 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( activation: MoEActivation | None = None, ) -> tuple[int, int]: """Round up hidden_size and intermediate_size based on backend requirements.""" + if backend in B12X_BACKENDS: + return hidden_size, intermediate_size if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SITU: # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a @@ -1292,6 +1327,16 @@ def convert_weight_to_mxfp4_moe_kernel_format( is_gfx1250 = on_gfx1250() + if mxfp4_backend in B12X_BACKENDS: + return ( + w13_weight.data, + w2_weight.data, + w13_weight_scale.data, + w2_weight_scale.data, + w13_bias, + w2_bias, + ) + if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( w13_weight, @@ -1687,6 +1732,16 @@ def make_mxfp4_moe_quant_config( layer: "RoutedExperts | None" = None, ) -> FusedMoEQuantConfig | None: """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" + if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8: + return mxfp4_mxfp8_moe_quant_config( + w1_bias=w1_bias, + w2_bias=w2_bias, + w1_scale=w1_scale, + w2_scale=w2_scale, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, + gemm1_clamp_limit=swiglu_limit, + ) if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, @@ -1754,6 +1809,7 @@ def make_mxfp4_moe_quant_config( gemm1_clamp_limit=swiglu_limit, ) elif mxfp4_backend in ( + Mxfp4MoeBackend.B12X_MXFP4_BF16, Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, Mxfp4MoeBackend.TRITON, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index f48e79e13da3..436e7e6cae58 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -18,6 +18,9 @@ nvfp4_w4a16_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.quantization.utils.b12x_moe import ( + prepare_nvfp4_moe_layer_for_b12x, +) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( nvfp4_swizzled_scale_to_cutedsl_mma_view, prepare_nvfp4_moe_layer_for_fi_or_cutlass, @@ -37,6 +40,7 @@ class NvFp4MoeBackend(Enum): + B12X = "B12X" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL" @@ -68,7 +72,12 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool: def backend_to_kernel_cls( backend: NvFp4MoeBackend, ) -> list[type[mk.FusedMoEExperts]]: - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: + if backend == NvFp4MoeBackend.B12X: + from vllm.model_executor.layers.fused_moe.b12x_moe import B12xExperts + + return [B12xExperts] + + elif backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsModular, TrtLlmNvFp4ExpertsMonolithic, @@ -146,6 +155,7 @@ def backend_to_kernel_cls( def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: """Map user's MoEBackend to NvFp4MoeBackend.""" mapping = { + "b12x": NvFp4MoeBackend.B12X, "cutlass": NvFp4MoeBackend.VLLM_CUTLASS, "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, @@ -163,6 +173,12 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: ) +def _use_a16(backend: NvFp4MoeBackend, checkpoint_uses_a16: bool) -> bool: + return checkpoint_uses_a16 or ( + backend == NvFp4MoeBackend.B12X and envs.VLLM_B12X_MOE_FORCE_A16 + ) + + def select_nvfp4_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, @@ -189,6 +205,7 @@ def select_nvfp4_moe_backend( ] NVFP4_BACKENDS_WITH_CLAMP = { + NvFp4MoeBackend.B12X, NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, @@ -250,6 +267,8 @@ def _return_or_raise( runner_backend = config.moe_backend if runner_backend != "auto": requested_backend = map_nvfp4_backend(runner_backend) + if _use_a16(requested_backend, False): + activation_key = None # For batched activation format, use batched variant if available. if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts @@ -310,6 +329,7 @@ def convert_to_nvfp4_moe_kernel_format( w2_scale_2: torch.Tensor, a2_scale: torch.Tensor | None, is_act_and_mul: bool, + use_a16: bool = False, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -320,7 +340,36 @@ def convert_to_nvfp4_moe_kernel_format( torch.Tensor, torch.Tensor, ]: - if nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: + use_a16 = _use_a16(nvfp4_backend, use_a16) + if nvfp4_backend == NvFp4MoeBackend.B12X: + if a13_scale is None or a2_scale is None: + if not use_a16: + raise ValueError("B12X NVFP4 MoE requires activation scales") + num_experts = w13.shape[0] + a13_scale = torch.ones(num_experts, dtype=torch.float32, device=w13.device) + a2_scale = torch.ones(num_experts, dtype=torch.float32, device=w2.device) + ( + w13, + w13_scale, + w13_scale_2, + a13_scale, + w2, + w2_scale, + w2_scale_2, + a2_scale, + ) = prepare_nvfp4_moe_layer_for_b12x( + w13=w13, + w13_scale=w13_scale, + w13_scale_2=w13_scale_2, + a13_scale=a13_scale, + w2=w2, + w2_scale=w2_scale, + w2_scale_2=w2_scale_2, + a2_scale=a2_scale, + is_act_and_mul=is_act_and_mul, + reorder_w13=use_a16, + ) + elif nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: ( w13, w13_scale, @@ -477,7 +526,9 @@ def make_nvfp4_moe_quant_config( swiglu_alpha: float | None = None, swiglu_beta: float | None = None, layer: torch.nn.Module | None = None, + use_a16: bool = False, ) -> FusedMoEQuantConfig: + use_a16 = _use_a16(backend, use_a16) if backend == NvFp4MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.quantization.utils.humming_utils import ( @@ -491,12 +542,16 @@ def make_nvfp4_moe_quant_config( gemm1_beta=getattr(layer, "swiglu_beta", None), gemm1_clamp_limit=swiglu_limit, ) - elif backend == NvFp4MoeBackend.MARLIN: + elif backend == NvFp4MoeBackend.MARLIN or ( + backend == NvFp4MoeBackend.B12X and use_a16 + ): return nvfp4_w4a16_moe_quant_config( g1_alphas=w13_scale_2, g2_alphas=w2_scale_2, w1_scale=w13_scale, w2_scale=w2_scale, + gemm1_alpha=swiglu_alpha, + gemm1_beta=swiglu_beta, gemm1_clamp_limit=swiglu_limit, ) elif backend == NvFp4MoeBackend.EMULATION: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index bb266c0f8248..813e4249f694 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -28,6 +28,7 @@ Mxfp4MoeBackend, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, + select_mxfp4_moe_backend, ) from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 CompressedTensorsMoEMethod, @@ -49,7 +50,12 @@ def __init__(self, moe): # use cutlass if supported, otherwise fallback to marlin for weight-only FP4 self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() self.experts_cls: type[mk.FusedMoEExperts] - if self.use_cutlass_mxfp4: + if getattr(moe, "moe_backend", "auto") == "b12x": + self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe) + assert experts_cls is not None + self.experts_cls = experts_cls + self.use_cutlass_mxfp4 = False + elif self.use_cutlass_mxfp4: logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE") self.experts_cls = CutlassExpertsMxfp4 elif current_platform.is_xpu(): @@ -138,7 +144,7 @@ def get_fused_moe_quant_config( w2_scale=layer.w2_weight_scale, ) else: - # W4A16: weight-only via Marlin + # Native packed-layout backends or weight-only Marlin. return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=layer.w13_weight_scale, @@ -208,6 +214,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: mxfp4_backend=self.mxfp4_backend, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def apply( self, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 40f389716661..92728691d2de 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -52,6 +52,7 @@ def __init__( ): super().__init__(moe) self.group_size = 16 + self.use_a16 = use_a16 # Select experts implementation. self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( @@ -231,6 +232,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2_scale_2=(1.0 / layer.w2_weight_global_scale), a2_scale=(1.0 / layer.w2_input_global_scale), is_act_and_mul=self.moe.is_act_and_mul, + use_a16=self.use_a16, ) replace_parameter(layer, "w13_weight", w13) @@ -267,6 +269,7 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon swiglu_alpha=getattr(layer, "swiglu_alpha", None), swiglu_beta=getattr(layer, "swiglu_beta", None), layer=layer, + use_a16=self.use_a16, ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py index 437a744947e8..c7903e2fdc35 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -61,7 +61,10 @@ def __init__(self, moe) -> None: self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() self.mxfp4_backend = Mxfp4MoeBackend.MARLIN self.experts_cls: type[mk.FusedMoEExperts] | None = None - if self.use_cutlass_mxfp4: + if getattr(moe, "moe_backend", "auto") == "b12x": + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) + self.use_cutlass_mxfp4 = False + elif self.use_cutlass_mxfp4: self.experts_cls = CutlassExpertsMxfp4 logger.info_once("Using CutlassExpertsMxfp4 for AutoRound MXFP4 MoE") elif current_platform.is_xpu(): @@ -146,7 +149,7 @@ def get_fused_moe_quant_config( w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, ) - # Weight-only (Marlin) or native XPU kernel. + # Native packed-layout backends or weight-only Marlin. return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=layer.w13_weight_scale, @@ -216,6 +219,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: mxfp4_backend=self.mxfp4_backend, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def apply( self, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index db98b76b3f5d..0268f0e469e6 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1558,6 +1558,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2_scale_2=layer.w2_weight_scale_2, a2_scale=layer.w2_input_scale, is_act_and_mul=self.moe.is_act_and_mul, + use_a16=self.use_a16, ) replace_parameter(layer, "w13_weight", w13) @@ -1594,6 +1595,7 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi swiglu_alpha=getattr(layer, "swiglu_alpha", None), swiglu_beta=getattr(layer, "swiglu_beta", None), layer=layer, + use_a16=self.use_a16, ) @property diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index de659e4f5084..d993e29e20a9 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -813,6 +813,7 @@ def _setup_kernel( experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def _convert_k3_situ_weight_to_kernel_format( self, layer: RoutedExperts diff --git a/vllm/model_executor/layers/quantization/utils/b12x_moe.py b/vllm/model_executor/layers/quantization/utils/b12x_moe.py new file mode 100644 index 000000000000..c2b2000db00e --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/b12x_moe.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight preparation helpers for native B12X MoE kernels.""" + +import torch + +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + swizzle_blockscale, +) +from vllm.utils.math_utils import round_up + + +def _reorder_w13_to_w31( + weight: torch.Tensor, + scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + if not weight.is_contiguous() or not scale.is_contiguous(): + raise ValueError("NVFP4 MoE weights and scales must be contiguous") + if weight.size(1) % 2 != 0 or scale.size(1) != weight.size(1): + raise ValueError("gated NVFP4 MoE weights and scales must have even rows") + + half = weight.size(1) // 2 + return ( + torch.cat((weight[:, half:], weight[:, :half]), dim=1).contiguous(), + torch.cat((scale[:, half:], scale[:, :half]), dim=1).contiguous(), + ) + + +def _pad_dim(tensor: torch.Tensor, dim: int, pad_size: int) -> torch.Tensor: + if pad_size <= 0: + return tensor + + dim %= tensor.ndim + shape = list(tensor.shape) + original_size = shape[dim] + shape[dim] += pad_size + padded = tensor.new_zeros(shape) + slices = [slice(None)] * tensor.ndim + slices[dim] = slice(0, original_size) + padded[tuple(slices)] = tensor + return padded.contiguous() + + +def _pad_gated_rows(tensor: torch.Tensor, half_pad_size: int) -> torch.Tensor: + if tensor.size(1) % 2 != 0: + raise ValueError("gated NVFP4 MoE tensors must have even row counts") + half_size = tensor.size(1) // 2 + first, second = tensor.split(half_size, dim=1) + return torch.cat( + (_pad_dim(first, 1, half_pad_size), _pad_dim(second, 1, half_pad_size)), + dim=1, + ).contiguous() + + +def _pad_gated_weights( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + pad_size = round_up(w13_scale.size(1), 128) - w13.size(1) + if pad_size <= 0: + return w13, w13_scale, w2, w2_scale + if w13_scale.size(1) != w13.size(1): + raise ValueError("w13 weight and scale row counts must match") + if w13.size(1) % 2 != 0 or pad_size % 2 != 0: + raise ValueError("gated NVFP4 MoE padding must split evenly") + + half_pad_size = pad_size // 2 + if half_pad_size % 16 != 0: + raise ValueError("NVFP4 MoE padding must preserve 16-value scale blocks") + + half_size = w13.size(1) // 2 + if w2.size(2) * 2 != half_size: + raise ValueError("w2 shape does not match gated w13") + if w2_scale.size(2) * 16 != half_size: + raise ValueError("w2 scale shape does not match gated w13") + + return ( + _pad_gated_rows(w13, half_pad_size), + _pad_gated_rows(w13_scale, half_pad_size), + _pad_dim(w2, 2, half_pad_size // 2), + _pad_dim(w2_scale, 2, half_pad_size // 16), + ) + + +def _per_expert_scale( + scale: torch.Tensor, + num_experts: int, + name: str, +) -> torch.Tensor: + scale = scale.to(torch.float32) + if scale.dim() == 0: + return scale.expand(num_experts).contiguous() + if scale.dim() == 1: + if scale.numel() != num_experts: + raise ValueError( + f"{name} must have {num_experts} elements, got {scale.numel()}" + ) + return scale.contiguous() + if scale.dim() == 2: + if scale.size(0) != num_experts: + raise ValueError( + f"{name} first dimension must be {num_experts}, got {scale.size(0)}" + ) + return scale.max(dim=1).values.contiguous() + raise ValueError(f"{name} must be scalar, 1D, or 2D, got {tuple(scale.shape)}") + + +def prepare_nvfp4_moe_layer_for_b12x( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w13_scale_2: torch.Tensor, + a13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + w2_scale_2: torch.Tensor, + a2_scale: torch.Tensor, + is_act_and_mul: bool, + reorder_w13: bool = False, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Prepare native B12X NVFP4 MoE weights and scales.""" + num_experts = w13.shape[0] + a13_scale = _per_expert_scale(a13_scale, num_experts, "a13_scale") + a2_scale = _per_expert_scale(a2_scale, num_experts, "a2_scale") + + if reorder_w13 and is_act_and_mul: + w13, w13_scale = _reorder_w13_to_w31(w13, w13_scale) + if is_act_and_mul: + w13, w13_scale, w2, w2_scale = _pad_gated_weights(w13, w13_scale, w2, w2_scale) + + w13_scale = swizzle_blockscale(w13_scale) + pad_size = w13_scale.size(1) - w13.size(1) + if pad_size > 0: + if is_act_and_mul: + raise RuntimeError("gated NVFP4 MoE padding must precede scale swizzling") + w13 = torch.nn.functional.pad(w13, (0, 0, 0, pad_size)) + w2 = torch.nn.functional.pad(w2, (0, pad_size // 2, 0, 0)) + w2_scale = torch.nn.functional.pad(w2_scale, (0, pad_size // 16)) + + w2_scale = swizzle_blockscale(w2_scale) + return w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c35e1ac30c9a..82bed644d5e1 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -12,6 +12,14 @@ import vllm.envs as envs from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.mxfp8.b12x import warmup_b12x_mxfp8_linear +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( + warmup_b12x_block_fp8_linear, +) +from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( + warmup_b12x_tensor_fp8_linear, +) +from vllm.model_executor.layers.fused_moe.b12x_moe import warmup_b12x_moe from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( @@ -113,15 +121,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) + compilation_config = worker.vllm_config.compilation_config + cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or []) + # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder # layer per token; warm them across token sizes first so the first real # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside). deepseek_v4_mhc_warmup( worker.get_model(), max_tokens=worker.scheduler_config.max_num_batched_tokens, - cudagraph_capture_sizes=( - worker.vllm_config.compilation_config.cudagraph_capture_sizes or [] - ), + cudagraph_capture_sizes=cudagraph_capture_sizes, ) # Run next so input-prep kernels JIT against pristine runner state. @@ -156,6 +165,67 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): max_tokens = worker.scheduler_config.max_num_batched_tokens deep_gemm_warmup(model, max_tokens) + b12x_linear_max_tokens = worker.scheduler_config.max_num_batched_tokens + b12x_linear_output_dtype = getattr( + getattr(worker, "model_config", None), + "dtype", + torch.bfloat16, + ) + warmed_block_fp8 = warmup_b12x_block_fp8_linear( + worker.get_model(), + max_tokens=b12x_linear_max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + output_dtype=b12x_linear_output_dtype, + ) + if warmed_block_fp8: + logger.info( + "Warmed up %d B12X block-FP8 linear GEMM signatures.", + warmed_block_fp8, + ) + + warmed_mxfp8 = warmup_b12x_mxfp8_linear( + worker.get_model(), + max_tokens=b12x_linear_max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + output_dtype=b12x_linear_output_dtype, + ) + if warmed_mxfp8: + logger.info("Warmed up %d B12X MXFP8 linear GEMM signatures.", warmed_mxfp8) + + warmed_tensor_fp8 = warmup_b12x_tensor_fp8_linear( + worker.get_model(), + max_tokens=b12x_linear_max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + output_dtype=b12x_linear_output_dtype, + ) + if warmed_tensor_fp8: + logger.info( + "Warmed up %d B12X tensor FP8 linear GEMM signatures.", + warmed_tensor_fp8, + ) + + moe_token_counts = [ + worker.scheduler_config.max_num_batched_tokens, + *cudagraph_capture_sizes, + *( + size + for size in (getattr(compilation_config, "compile_sizes", None) or []) + if isinstance(size, int) + ), + ] + max_num_scheduled_tokens = getattr( + worker.scheduler_config, + "max_num_scheduled_tokens", + None, + ) + if max_num_scheduled_tokens is not None: + moe_token_counts.append(max_num_scheduled_tokens) + warmup_b12x_moe( + worker.get_model(), + max_tokens=max(moe_token_counts), + token_counts=moe_token_counts, + ) + minimax_m3_msa_warmup(worker) enable_flashinfer_autotune = ( From 1697565c109a86d2cc2324350f5c9247e4a4f8c6 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Fri, 7 Aug 2026 12:21:48 -0700 Subject: [PATCH 4/5] feat(attention): add B12X causal paged backend Add an opt-in SM120/SM121 backend for causal paged GQA attention. Support BF16 and FP8 KV caches, decode and extend batches, sliding windows, sinks, speculative verification, and CUDA graph replay. Assisted-by: OpenAI Codex Signed-off-by: Luke Alonso --- tests/v1/attention/test_attention_backends.py | 311 ++++- tests/v1/attention/test_b12x_attn.py | 196 +++ vllm/v1/attention/backends/b12x_attn.py | 1239 +++++++++++++++++ vllm/v1/attention/backends/registry.py | 1 + 4 files changed, 1716 insertions(+), 31 deletions(-) create mode 100644 tests/v1/attention/test_b12x_attn.py create mode 100644 vllm/v1/attention/backends/b12x_attn.py diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 87a9c80942ee..f41c22fc33ef 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -3,6 +3,7 @@ """Tests for v1 attention backends without GPUModelRunner dependency.""" from functools import partial +from types import SimpleNamespace import pytest import torch @@ -30,6 +31,7 @@ AttentionType, CommonAttentionMetadata, ) +from vllm.v1.attention.backends.b12x_attn import B12XPagedAttentionBackend from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.backends.utils import ( set_kv_cache_layout, @@ -153,7 +155,12 @@ def create_and_prepopulate_kv_cache( # For an fp8 kv cache, store the cache in the fp8 dtype so that assigning # the higher-precision context tensors quantizes them, mirroring runtime. fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) - storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype + if fp8_kv_cache: + storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] + elif kv_cache_dtype == "auto": + storage_dtype = dtype + else: + storage_dtype = STR_DTYPE_TO_TORCH_DTYPE[kv_cache_dtype] kv_cache = torch.zeros( num_blocks, @@ -250,6 +257,8 @@ def run_attention_backend( attn_type: AttentionType = AttentionType.DECODER, sliding_window: int | None = None, kv_cache_dtype: str = "auto", + sinks: torch.Tensor | None = None, + use_cuda_graph: bool = False, ) -> torch.Tensor: """Run attention computation using the specified backend's AttentionImpl.""" @@ -292,13 +301,14 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): ) else: # Build metadata - builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device) - if actual_backend == AttentionBackendEnum.FLEX_ATTENTION: - builder.direct_build = use_direct_block_mask - attn_metadata = builder.build( - common_prefix_len=0, - common_attn_metadata=common_attn_metadata, - ) + with set_current_vllm_config(vllm_config): + builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device) + if actual_backend == AttentionBackendEnum.FLEX_ATTENTION: + builder.direct_build = use_direct_block_mask + attn_metadata = builder.build( + common_prefix_len=0, + common_attn_metadata=common_attn_metadata, + ) # Instantiate implementation num_heads = vllm_config.model_config.get_num_attention_heads( @@ -309,16 +319,19 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): ) head_size = vllm_config.model_config.get_head_size() scale = 1.0 / (head_size**0.5) - impl = impl_cls( - num_heads=num_heads, - head_size=head_size, - scale=scale, - num_kv_heads=num_kv_heads, - alibi_slopes=None, - sliding_window=sliding_window, - attn_type=attn_type, - kv_cache_dtype=kv_cache_dtype, - ) + extra_impl_kwargs = {"sinks": sinks} if sinks is not None else {} + with set_current_vllm_config(vllm_config): + impl = impl_cls( + num_heads=num_heads, + head_size=head_size, + scale=scale, + num_kv_heads=num_kv_heads, + alibi_slopes=None, + sliding_window=sliding_window, + attn_type=attn_type, + kv_cache_dtype=kv_cache_dtype, + **extra_impl_kwargs, + ) # Create mock layer and output buffer mock_layer = MockAttentionLayer(device) @@ -334,9 +347,28 @@ def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): impl.do_kv_cache_update( mock_layer, key, value, kv_cache, attn_metadata.slot_mapping ) - output = impl.forward( - mock_layer, query, key, value, kv_cache, attn_metadata, output=output - ) + if use_cuda_graph: + impl.forward( + mock_layer, query, key, value, kv_cache, attn_metadata, output=output + ) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + impl.forward( + mock_layer, + query, + key, + value, + kv_cache, + attn_metadata, + output=output, + ) + graph.replay() + torch.accelerator.synchronize() + else: + output = impl.forward( + mock_layer, query, key, value, kv_cache, attn_metadata, output=output + ) return output @@ -354,6 +386,11 @@ def _test_backend_correctness( rtol: float = 1e-2, tensor_parallel_size: int = 1, kv_cache_dtype: str = "auto", + sliding_window_override: int | None = None, + use_attention_sinks: bool = False, + use_cuda_graph: bool = False, + num_speculative_tokens: int = 0, + model_dtype: torch.dtype | None = None, ): """ Test that all backends produce similar outputs to a reference implementation @@ -398,10 +435,20 @@ def _test_backend_correctness( model_name=model, tensor_parallel_size=1, # Always use TP=1 to avoid multi-GPU requirements max_model_len=max(batch_spec.seq_lens), + dtype=model_dtype or "auto", block_size=block_size, num_gpu_blocks=8192, hf_config_override=hf_config_override, ) + if AttentionBackendEnum.B12X_ATTN in backend_to_test: + vllm_config.scheduler_config.max_num_seqs = batch_spec.batch_size + vllm_config.scheduler_config.max_num_batched_tokens = max( + sum(batch_spec.query_lens), 64 + ) + if num_speculative_tokens > 0: + vllm_config.speculative_config = SimpleNamespace( + num_speculative_tokens=num_speculative_tokens + ) vllm_config.cache_config.cache_dtype = kv_cache_dtype device = torch.device(f"{DEVICE_TYPE}:0") @@ -418,10 +465,19 @@ def _test_backend_correctness( vllm_config.parallel_config ) head_size = vllm_config.model_config.get_head_size() - sliding_window = vllm_config.model_config.get_sliding_window() + sliding_window = ( + sliding_window_override + if sliding_window_override is not None + else vllm_config.model_config.get_sliding_window() + ) dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype) block_size = vllm_config.cache_config.block_size scale = 1.0 / (head_size**0.5) + sinks = ( + torch.linspace(-0.3, 0.3, num_q_heads, dtype=dtype, device=device) + if use_attention_sinks + else None + ) fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype) if fp8_kv_cache: @@ -474,14 +530,45 @@ def _test_backend_correctness( block_mask = create_block_mask( final_mask_mod, B=None, H=None, Q_LEN=q_len, KV_LEN=kv_len, device=device ) - sdpa_out_i = flex_attention( - q_sdpa_in, - k_sdpa_in, - v_sdpa_in, - block_mask=block_mask, - scale=scale, - enable_gqa=True, - ) + if sinks is None: + sdpa_out_i = flex_attention( + q_sdpa_in, + k_sdpa_in, + v_sdpa_in, + block_mask=block_mask, + scale=scale, + enable_gqa=True, + ) + else: + q_idx = torch.arange(q_len, device=device).unsqueeze(1) + kv_idx = torch.arange(kv_len, device=device).unsqueeze(0) + mask = final_mask_mod( + torch.zeros((), device=device), + torch.zeros((), device=device), + q_idx, + kv_idx, + ) + scores = ( + torch.einsum( + "hqd,hkd->hqk", + q_sdpa_in.squeeze(0).float(), + k_sdpa_in.squeeze(0).float(), + ) + * scale + ) + scores = scores.masked_fill(~mask.unsqueeze(0), float("-inf")) + sink_logits = sinks.float().view(num_q_heads, 1, 1) + probabilities = torch.softmax( + torch.cat([scores, sink_logits.expand(num_q_heads, q_len, 1)], dim=-1), + dim=-1, + )[..., :kv_len] + sdpa_out_i = ( + torch.einsum( + "hqk,hkd->hqd", probabilities, v_sdpa_in.squeeze(0).float() + ) + .unsqueeze(0) + .to(dtype) + ) all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0)) @@ -543,9 +630,24 @@ def _test_backend_correctness( if backend_name == AttentionBackendEnum.FLASHINFER: set_kv_cache_layout("HND") reset_kv_cache_layout = True + elif backend_name == AttentionBackendEnum.B12X_ATTN: + set_kv_cache_layout("NHD") + reset_kv_cache_layout = True kv_cache_for_backend = kv_cache - if backend_cls is not None: + if backend_name == AttentionBackendEnum.B12X_ATTN: + cache_dtype = ( + FP8_KV_CACHE_DTYPES[kv_cache_dtype] + if is_quantized_kv_cache(kv_cache_dtype) + else kv_cache.dtype + ) + typed_cache = kv_cache.view(cache_dtype) + key_cache = typed_cache[..., :head_size].permute(0, 2, 1, 3) + value_cache = typed_cache[..., head_size:].permute(0, 2, 1, 3) + kv_cache_for_backend = torch.stack((key_cache, value_cache), dim=1) + if is_quantized_kv_cache(kv_cache_dtype): + kv_cache_for_backend = kv_cache_for_backend.view(torch.uint8) + elif backend_cls is not None: try: stride_order = backend_cls.get_kv_cache_stride_order() except (AttributeError, NotImplementedError): @@ -574,6 +676,8 @@ def _test_backend_correctness( sliding_window=sliding_window, attn_type=attn_type, kv_cache_dtype=kv_cache_dtype, + sinks=sinks, + use_cuda_graph=use_cuda_graph, ) finally: if reset_kv_cache_layout: @@ -606,6 +710,151 @@ def error_msg(msg: str, backend_name: str): ) +def _require_b12x_paged_attention() -> None: + capability = current_platform.get_device_capability() + if ( + not current_platform.is_cuda() + or capability is None + or not B12XPagedAttentionBackend.supports_compute_capability(capability) + ): + pytest.skip("B12X paged attention requires SM120 or SM121.") + + from b12x.attention import paged + + if not paged.is_supported(): + pytest.skip("B12X paged attention is not available.") + + +def _b12x_causal_mask( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + *, + context_len: int, +): + return q_idx + context_len >= kv_idx + + +def _b12x_causal_sliding_window_mask( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + kv_idx: torch.Tensor, + *, + context_len: int, + sliding_window: int, +): + causal_mask = q_idx + context_len >= kv_idx + window_mask = q_idx + context_len - kv_idx < sliding_window + return causal_mask & window_mask + + +@pytest.mark.parametrize( + "batch_spec_name", + ["small_decode", "small_prefill", "mixed_small", "medium_decode"], +) +@pytest.mark.parametrize( + ("kv_cache_dtype", "model_dtype"), + [ + ("auto", None), + ("bfloat16", torch.bfloat16), + ("fp8_e4m3", torch.bfloat16), + ], +) +@pytest.mark.parametrize("block_size", [64, 128]) +def test_b12x_causal_backend_correctness( + default_vllm_config, + workspace_init, + batch_spec_name: str, + kv_cache_dtype: str, + model_dtype: torch.dtype | None, + block_size: int, +): + """B12X causal paged attention matches the shared SDPA reference.""" + _require_b12x_paged_attention() + + _test_backend_correctness( + BATCH_SPECS[batch_spec_name], + "Qwen/Qwen3-0.6B", + [AttentionBackendEnum.B12X_ATTN], + _b12x_causal_mask, + block_size=block_size, + kv_cache_dtype=kv_cache_dtype, + model_dtype=model_dtype, + ) + + +@pytest.mark.parametrize("batch_spec_name", ["small_decode", "small_prefill"]) +def test_b12x_causal_sliding_window_and_sinks( + default_vllm_config, + workspace_init, + batch_spec_name: str, +): + """B12X preserves causal SWA and attention-sink semantics.""" + _require_b12x_paged_attention() + + sliding_window = 16 + mask = partial(_b12x_causal_sliding_window_mask, sliding_window=sliding_window) + + _test_backend_correctness( + BATCH_SPECS[batch_spec_name], + "Qwen/Qwen3-0.6B", + [AttentionBackendEnum.B12X_ATTN], + mask, + block_size=64, + atol=3e-2, + rtol=3e-2, + sliding_window_override=sliding_window, + use_attention_sinks=True, + ) + + +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"]) +def test_b12x_decode_cuda_graph_replay_with_sliding_window_and_sinks( + default_vllm_config, + workspace_init, + kv_cache_dtype: str, +): + """B12X causal metadata and sinks remain capture-safe during replay.""" + _require_b12x_paged_attention() + + sliding_window = 16 + mask = partial(_b12x_causal_sliding_window_mask, sliding_window=sliding_window) + + _test_backend_correctness( + BATCH_SPECS["small_decode"], + "Qwen/Qwen3-0.6B", + [AttentionBackendEnum.B12X_ATTN], + mask, + block_size=64, + atol=3e-2, + rtol=3e-2, + sliding_window_override=sliding_window, + use_attention_sinks=True, + use_cuda_graph=True, + kv_cache_dtype=kv_cache_dtype, + ) + + +def test_b12x_speculative_verifier_cuda_graph_replay( + default_vllm_config, + workspace_init, +): + """B12X replays uniform speculative verification through its graph plan.""" + _require_b12x_paged_attention() + + _test_backend_correctness( + BatchSpec(seq_lens=[32, 40], query_lens=[4, 4]), + "Qwen/Qwen3-0.6B", + [AttentionBackendEnum.B12X_ATTN], + _b12x_causal_mask, + block_size=128, + num_speculative_tokens=3, + use_cuda_graph=True, + ) + + @pytest.mark.parametrize( "batch_spec_name", [ diff --git a/tests/v1/attention/test_b12x_attn.py b/tests/v1/attention/test_b12x_attn.py new file mode 100644 index 000000000000..44c3c43fe21e --- /dev/null +++ b/tests/v1/attention/test_b12x_attn.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms.interface import DeviceCapability +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends import b12x_attn +from vllm.v1.attention.backends.b12x_attn import ( + B12XPagedAttentionBackend, + B12XPagedAttentionImpl, + B12XPagedMetadataBuilder, + _kv_page_size, + _max_page_table_width, +) +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.attention.backends.utils import set_kv_cache_layout + + +def test_b12x_attention_backend_is_registered() -> None: + assert AttentionBackendEnum.B12X_ATTN.get_class() is B12XPagedAttentionBackend + + +def test_b12x_attention_advertises_only_supported_contract() -> None: + assert B12XPagedAttentionBackend.get_supported_kernel_block_sizes() == [64, 128] + assert B12XPagedAttentionBackend.get_supported_head_sizes() == [64, 128, 192, 256] + assert B12XPagedAttentionBackend.supports_sliding_window() + assert B12XPagedAttentionBackend.supports_sink() + assert not B12XPagedAttentionBackend.supports_non_causal() + assert B12XPagedAttentionBackend.supports_compute_capability( + DeviceCapability(12, 0) + ) + assert B12XPagedAttentionBackend.supports_compute_capability( + DeviceCapability(12, 1) + ) + assert not B12XPagedAttentionBackend.supports_compute_capability( + DeviceCapability(12, 2) + ) + assert not B12XPagedAttentionBackend.supports_compute_capability( + DeviceCapability(10, 0) + ) + assert {"fp8", "fp8_e4m3"}.issubset( + B12XPagedAttentionBackend.supported_kv_cache_dtypes + ) + assert B12XPagedAttentionBackend.supports_dtype(torch.bfloat16) + assert not B12XPagedAttentionBackend.supports_dtype(torch.float16) + + +def test_b12x_attention_rejects_fp16_queries_with_fp8_kv( + default_vllm_config, +) -> None: + reason = B12XPagedAttentionBackend.supports_combination( + head_size=128, + dtype=torch.float16, + kv_cache_dtype="fp8_e4m3", + block_size=128, + use_mla=False, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + device_capability=DeviceCapability(12, 0), + ) + + assert reason == "B12X_ATTN currently requires bfloat16 queries" + + +def test_b12x_attention_rejects_float16_kv_dtype( + default_vllm_config, +) -> None: + reason = B12XPagedAttentionBackend.supports_combination( + head_size=128, + dtype=torch.bfloat16, + kv_cache_dtype="float16", + block_size=128, + use_mla=False, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + device_capability=DeviceCapability(12, 0), + ) + + assert reason == "B12X_ATTN does not support float16 KV cache" + + +def test_b12x_attention_uses_two_plane_nhd_cache() -> None: + assert B12XPagedAttentionBackend.get_kv_cache_shape( + num_blocks=3, + block_size=128, + num_kv_heads=4, + head_size=128, + cache_dtype_str="fp8_e4m3", + ) == (3, 2, 128, 4, 128) + set_kv_cache_layout("NHD") + try: + assert B12XPagedAttentionBackend.get_kv_cache_stride_order() == ( + 0, + 1, + 2, + 3, + 4, + ) + assert B12XPagedAttentionBackend.get_kv_cache_stride_order(True) == ( + 1, + 0, + 2, + 3, + 4, + 5, + ) + finally: + set_kv_cache_layout(None) + assert B12XPagedAttentionBackend.get_required_kv_cache_layout() == "NHD" + + +def test_b12x_attention_rejects_unsupported_page_size() -> None: + with pytest.raises(ValueError, match="block_size"): + B12XPagedAttentionBackend.get_kv_cache_shape(3, 32, 4, 128) + + +def test_b12x_attention_uses_uniform_batch_graphs() -> None: + assert ( + B12XPagedMetadataBuilder._cudagraph_support is AttentionCGSupport.UNIFORM_BATCH + ) + + +def test_b12x_attention_hybrid_cache_capacity_includes_expansion() -> None: + assert _max_page_table_width(4096, 128, 4096, False) == 32 + assert _max_page_table_width(4096, 128, 4096, True) == 64 + + +def test_b12x_attention_runtime_page_size_comes_from_cache() -> None: + key_cache = torch.empty((3, 64, 4, 128), device="meta") + value_cache = torch.empty_like(key_cache) + + assert _kv_page_size(key_cache, value_cache) == 64 + with pytest.raises(ValueError, match="matching K/V page sizes"): + _kv_page_size(key_cache, torch.empty((3, 128, 4, 128), device="meta")) + + +def test_b12x_attention_lazily_prepares_decode_bucket(monkeypatch) -> None: + impl = object.__new__(B12XPagedAttentionImpl) + plan = SimpleNamespace(layout=SimpleNamespace(nbytes=96)) + created: list[tuple[int, int]] = [] + + def create_plan(page_size: int, batch_size: int) -> SimpleNamespace: + created.append((page_size, batch_size)) + return plan + + impl._decode_plans = {} + impl._create_decode_plan = create_plan + impl._scratch_nbytes = 128 + impl._extend_plans = {} + impl._verify_q_per_req = 0 + metadata = SimpleNamespace(max_query_len=1) + monkeypatch.setattr(b12x_attn, "_capture_alloc_forbidden", lambda: False) + + assert impl._select_plan(metadata, 7, 7, 7, 64) is plan + assert impl._select_plan(metadata, 7, 7, 7, 64) is plan + assert created == [(64, 7)] + + +def test_b12x_attention_fp8_descales_follow_request_batch() -> None: + impl = object.__new__(B12XPagedAttentionImpl) + impl.kv_cache_dtype = "fp8_e4m3" + layer = SimpleNamespace( + _k_scale=torch.tensor(2.0), + _v_scale=torch.tensor([3.0, 4.0, 5.0]), + ) + + k_descale, v_descale = impl._prepare_fp8_descales( + layer, num_reqs=2, device=torch.device("cpu") + ) + + torch.testing.assert_close(k_descale, torch.tensor([2.0, 2.0])) + torch.testing.assert_close(v_descale, torch.tensor([3.0, 4.0])) + assert k_descale.stride() == (0,) + assert v_descale.stride() == (1,) + + +def test_b12x_attention_sinks_refresh_in_place_after_reload() -> None: + impl = object.__new__(B12XPagedAttentionImpl) + impl._sinks_cache = {} + source = torch.tensor([1.0, 2.0], dtype=torch.bfloat16) + + sinks = impl._prepare_sinks(source, torch.device("cpu")) + assert sinks is not None + sinks_ptr = sinks.data_ptr() + source.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16)) + refreshed = impl._prepare_sinks(source, torch.device("cpu")) + + assert refreshed is not None + assert refreshed.data_ptr() == sinks_ptr + torch.testing.assert_close(refreshed, source.float()) diff --git a/vllm/v1/attention/backends/b12x_attn.py b/vllm/v1/attention/backends/b12x_attn.py new file mode 100644 index 000000000000..0e99983679c0 --- /dev/null +++ b/vllm/v1/attention/backends/b12x_attn.py @@ -0,0 +1,1239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""B12X causal paged attention backend for SM12x.""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass +from functools import partial +from typing import Any, ClassVar + +import torch + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config.cache import CacheDType +from vllm.logger import init_logger +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability +from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + is_quantized_kv_cache, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionImpl, + AttentionLayer, + AttentionMetadata, + AttentionMetadataBuilder, + AttentionType, + CommonAttentionMetadata, + MultipleOf, +) +from vllm.v1.attention.backends.utils import ( + KVCacheLayoutType, + get_kv_cache_layout, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + +_B12X_SUPPORTED_PAGE_SIZES = (64, 128) +_B12X_PREFERRED_PAGE_SIZE = 128 +_MIN_PAGED_TILE_Q = 16 +_B12X_FP8_KV_CACHE_DTYPES = ("fp8", "fp8_e4m3") +_B12X_SUPPORTED_KV_CACHE_DTYPES = ( + "auto", + "bfloat16", + *_B12X_FP8_KV_CACHE_DTYPES, +) + + +def _max_page_table_width( + max_model_len: int, + block_size: int, + max_num_batched_tokens: int, + is_hybrid: bool, +) -> int: + width = max(cdiv(max(max_model_len, 1), block_size), 1) + if is_hybrid: + # Hybrid cache setup can enlarge the storage block after attention + # layers are initialized. Its expansion into kernel-sized blocks adds + # at most one storage block of trailing page-table capacity. + width += cdiv(max_num_batched_tokens, block_size) + return width + + +def _kv_page_size(key_cache: torch.Tensor, value_cache: torch.Tensor) -> int: + """Return the static kernel page geometry negotiated by vLLM. + + The KV manager can split the configured storage block into a smaller + kernel page when another backend shares its cache group. This notably + happens when a B12X target shares a DFlash cache group with FlashInfer. + Cache shapes are fixed before graph capture, so this is not a live-length + policy decision. + """ + if key_cache.ndim < 2 or value_cache.ndim < 2: + raise ValueError( + "B12X_ATTN expects paged K/V caches with a page dimension, got " + f"{tuple(key_cache.shape)} and {tuple(value_cache.shape)}." + ) + key_page_size = int(key_cache.shape[1]) + value_page_size = int(value_cache.shape[1]) + if key_page_size != value_page_size: + raise ValueError( + "B12X_ATTN requires matching K/V page sizes, got " + f"{key_page_size} and {value_page_size}." + ) + if key_page_size not in _B12X_SUPPORTED_PAGE_SIZES: + raise ValueError( + "B12X_ATTN requires runtime page size in " + f"{_B12X_SUPPORTED_PAGE_SIZES}, got {key_page_size}." + ) + return key_page_size + + +def _env_int(name: str, default: int) -> int: + value = os.getenv(name) + if value is None or value == "": + return default + try: + parsed = int(value) + except ValueError: + logger.warning("Ignoring invalid %s=%r; using %d", name, value, default) + return default + if parsed <= 0: + logger.warning("Ignoring non-positive %s=%r; using %d", name, value, default) + return default + return parsed + + +def _env_optional_storage_limit(name: str, *, allow_zero: bool) -> int | None: + value = os.getenv(name) + if value is None or value == "": + return None + try: + parsed = int(value) + except ValueError: + logger.warning( + "Ignoring invalid %s=%r; using B12X's planned capacity", + name, + value, + ) + return None + minimum = 0 if allow_zero else 1 + if parsed < minimum: + logger.warning( + "Ignoring %s=%r below the minimum %d; using B12X's planned capacity", + name, + value, + minimum, + ) + return None + return parsed + + +def _capture_alloc_forbidden() -> bool: + if not torch.cuda.is_available(): + return False + try: + return bool(torch.cuda.is_current_stream_capturing()) + except RuntimeError: + return False + + +def _ensure_i32_contiguous(tensor: torch.Tensor, name: str) -> torch.Tensor: + if tensor.dtype != torch.int32: + if _capture_alloc_forbidden(): + raise RuntimeError( + f"B12X_ATTN would convert {name} to int32 during CUDA graph " + "capture. Prepare int32 metadata before capture." + ) + tensor = tensor.to(torch.int32) + if not tensor.is_contiguous(): + if _capture_alloc_forbidden(): + raise RuntimeError( + f"B12X_ATTN would make {name} contiguous during CUDA graph " + "capture. Prepare contiguous metadata before capture." + ) + tensor = tensor.contiguous() + return tensor + + +def _dtype_from_cache_config( + kv_cache_dtype: str, + vllm_config: VllmConfig, +) -> torch.dtype: + if kv_cache_dtype == "bfloat16": + return torch.bfloat16 + if kv_cache_dtype in _B12X_FP8_KV_CACHE_DTYPES: + return current_platform.fp8_dtype() + if kv_cache_dtype != "auto": + raise NotImplementedError( + "B12X_ATTN currently supports only auto, bfloat16, " + "fp8, and fp8_e4m3 " + f"KV cache dtypes; got {kv_cache_dtype!r}." + ) + return vllm_config.model_config.dtype + + +def _is_b12x_fp8_kv_cache(kv_cache_dtype: str) -> bool: + return kv_cache_dtype in _B12X_FP8_KV_CACHE_DTYPES + + +class B12XPagedAttentionBackend(AttentionBackend): + """Opt-in b12x paged attention backend for regular/GQA decoder layers.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + + forward_includes_kv_cache_update: bool = False + + @staticmethod + def get_name() -> str: + return "B12X_ATTN" + + @classmethod + def get_impl_cls(cls) -> type[B12XPagedAttentionImpl]: + return B12XPagedAttentionImpl + + @staticmethod + def get_builder_cls() -> type[B12XPagedMetadataBuilder]: + return B12XPagedMetadataBuilder + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return list(_B12X_SUPPORTED_PAGE_SIZES) + + @classmethod + def supports_block_size(cls, block_size: int | None) -> bool: + return block_size is None or int(block_size) in _B12X_SUPPORTED_PAGE_SIZES + + @classmethod + def get_preferred_block_size(cls, default_block_size: int) -> int: + if int(default_block_size) in _B12X_SUPPORTED_PAGE_SIZES: + return int(default_block_size) + return _B12X_PREFERRED_PAGE_SIZE + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [64, 128, 192, 256] + + @classmethod + def supports_sink(cls) -> bool: + return True + + @classmethod + def supports_sliding_window(cls) -> bool: + return True + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + # Consumer Blackwell SM120 / SM121. The b12x paged kernels also gate + # internally, but keep vLLM selection fail-fast and explicit. + return (capability.major, capability.minor) in ((12, 0), (12, 1)) + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + if dtype != torch.bfloat16: + return "B12X_ATTN currently requires bfloat16 queries" + if kv_cache_dtype == "float16": + return "B12X_ATTN does not support float16 KV cache" + if ( + kv_cache_dtype is not None + and is_quantized_kv_cache(kv_cache_dtype) + and not _is_b12x_fp8_kv_cache(kv_cache_dtype) + ): + return ( + "B12X_ATTN currently supports only fp8/fp8_e4m3 quantized " + "KV cache dtypes" + ) + vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + if parallel_config.decode_context_parallel_size > 1: + return "B12X_ATTN does not yet support decode context parallelism" + if parallel_config.prefill_context_parallel_size > 1: + return "B12X_ATTN does not yet support prefill context parallelism" + return None + + @classmethod + def get_kv_cache_shape( + cls, + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if block_size not in _B12X_SUPPORTED_PAGE_SIZES: + raise ValueError( + "B12X_ATTN requires block_size in " + f"{_B12X_SUPPORTED_PAGE_SIZES}, got {block_size}." + ) + if cache_dtype_str not in _B12X_SUPPORTED_KV_CACHE_DTYPES: + raise ValueError( + "B12X_ATTN currently supports only auto, bfloat16, " + "fp8, and fp8_e4m3 " + f"KV cache dtypes; got {cache_dtype_str!r}." + ) + return (num_blocks, 2, block_size, num_kv_heads, head_size) + + @classmethod + def get_kv_cache_stride_order( + cls, + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + cache_layout = get_kv_cache_layout() + if cache_layout != "NHD": + raise ValueError( + f"B12X_ATTN requires NHD KV cache layout; got {cache_layout!r}." + ) + if include_num_layers_dimension: + return (1, 0, 2, 3, 4, 5) + return (0, 1, 2, 3, 4) + + @classmethod + def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None: + return "NHD" + + +@dataclass +class B12XPagedMetadata(AttentionMetadata): + num_actual_tokens: int + max_query_len: int + query_start_loc: torch.Tensor + max_seq_len: int + seq_lens: torch.Tensor + block_table: torch.Tensor + slot_mapping: torch.Tensor + causal: bool = True + + +class B12XPagedMetadataBuilder(AttentionMetadataBuilder[B12XPagedMetadata]): + """Metadata builder for B12X_ATTN. + + Decode and uniform speculative-verifier batches use preplanned graph + buckets. Extend/prefill remains eager and does not affect uniform decode + graph eligibility. + """ + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + supports_update_block_table: bool = True + + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> AttentionCGSupport: + del vllm_config, kv_cache_spec + return cls._cudagraph_support + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> B12XPagedMetadata: + del common_prefix_len, fast_build + cm = common_attn_metadata + return B12XPagedMetadata( + num_actual_tokens=cm.num_actual_tokens, + max_query_len=cm.max_query_len, + query_start_loc=cm.query_start_loc, + max_seq_len=cm.max_seq_len, + seq_lens=cm.seq_lens, + block_table=cm.block_table_tensor, + slot_mapping=cm.slot_mapping, + causal=cm.causal, + ) + + def update_block_table( + self, + metadata: B12XPagedMetadata, + blk_table: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> B12XPagedMetadata: + new_metadata = copy.copy(metadata) + new_metadata.block_table = blk_table + new_metadata.slot_mapping = slot_mapping + return new_metadata + + +class B12XPagedAttentionImpl(AttentionImpl[B12XPagedMetadata]): + """b12x paged GQA attention implementation.""" + + can_return_lse_for_decode: bool = False + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None = None, + attn_type: AttentionType = AttentionType.DECODER, + kv_sharing_target_layer_name: str | None = None, + sinks: torch.Tensor | None = None, + ) -> None: + if alibi_slopes is not None: + raise NotImplementedError("B12X_ATTN does not support ALiBi.") + if logits_soft_cap not in (None, 0): + raise NotImplementedError("B12X_ATTN does not support logits soft cap.") + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + "B12X_ATTN currently supports decoder self-attention only." + ) + if is_quantized_kv_cache(kv_cache_dtype) and not _is_b12x_fp8_kv_cache( + kv_cache_dtype + ): + raise NotImplementedError( + "B12X_ATTN currently supports only fp8/fp8_e4m3 quantized " + "KV cache dtypes." + ) + if num_heads % num_kv_heads != 0: + raise ValueError("B12X_ATTN requires q heads divisible by kv heads.") + + expected_scale = head_size**-0.5 + if not math.isclose(float(scale), expected_scale, rel_tol=1e-5, abs_tol=1e-7): + raise NotImplementedError( + "B12X_ATTN currently requires canonical softmax scale " + f"head_dim**-0.5={expected_scale}, got {scale}." + ) + if self.total_cp_world_size > 1: + raise NotImplementedError( + "B12X_ATTN does not yet support decode/prefill context parallelism." + ) + + self.num_heads = int(num_heads) + self.head_size = int(head_size) + self.output_head_size = self.head_size + self.scale = float(scale) + self.num_kv_heads = int(num_kv_heads) + self.num_queries_per_kv = self.num_heads // self.num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + self.attn_type = attn_type + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name + self.window_left = -1 if sliding_window is None else int(sliding_window) - 1 + + self.sinks = sinks + if self.sinks is not None and ( + self.sinks.ndim != 1 or int(self.sinks.shape[0]) != self.num_heads + ): + raise ValueError( + "B12X_ATTN sinks must have shape " + f"[{self.num_heads}], got {tuple(self.sinks.shape)}." + ) + self._sinks_cache: dict[tuple[Any, ...], torch.Tensor] = {} + + vllm_config = get_current_vllm_config() + scheduler_config = vllm_config.scheduler_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + spec_config = vllm_config.speculative_config + default_block_size = int(cache_config.block_size) + if default_block_size not in _B12X_SUPPORTED_PAGE_SIZES: + raise ValueError( + "B12X_ATTN requires --block-size in " + f"{_B12X_SUPPORTED_PAGE_SIZES}, got " + f"{cache_config.block_size}." + ) + + self.device = torch.device("cuda", torch.accelerator.current_device_index()) + self.dtype = model_config.dtype + self.kv_torch_dtype = _dtype_from_cache_config(kv_cache_dtype, vllm_config) + if self.dtype != torch.bfloat16: + raise NotImplementedError("B12X_ATTN currently requires bfloat16 queries.") + max_batched = int(scheduler_config.max_num_batched_tokens) + max_num_seqs = int(scheduler_config.max_num_seqs) + max_model_len = int(model_config.max_model_len) + self._max_num_seqs = max_num_seqs + max_page_table_widths = { + page_size: _max_page_table_width( + max_model_len, + page_size, + max_batched, + model_config.is_hybrid, + ) + for page_size in _B12X_SUPPORTED_PAGE_SIZES + } + + # Extend dispatch may depend on the static Q tensor capacity, but never + # on live per-request lengths. Keep a small set of capacity buckets so + # short/tail prefills do not replay the maximum 8K CTA grid. + self._extend_q_capacities = tuple( + sorted( + { + min(max_batched, q_capacity) + for q_capacity in (128, 512, 1024, 2048, 4096, max_batched) + if q_capacity > 0 + } + ) + ) + + def _extend_work_items( + page_size: int, + q_capacity: int, + batch_size: int, + ) -> int: + capacity = plan_extend_graph_capacity( + device=self.device, + q_dtype=self.dtype, + kv_dtype=self.kv_torch_dtype, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + batch=batch_size, + total_q_capacity=q_capacity, + max_cache_page_count=max_page_table_widths[page_size], + window_left=self.window_left, + ) + return _env_int( + "VLLM_B12X_PAGED_EXTEND_MAX_WORK_ITEMS", + capacity.max_work_items, + ) + + from b12x.attention.paged import ( + Caps as B12XPagedAttentionScratchCaps, + ) + from b12x.attention.paged import ( + compile as compile_paged_attention, + ) + from b12x.attention.paged import ( + decode_graph_capacity as plan_decode_graph_capacity, + ) + from b12x.attention.paged import ( + decode_graph_scratch_envelope as plan_decode_graph_scratch_envelope, + ) + from b12x.attention.paged import ( + extend_graph_capacity as plan_extend_graph_capacity, + ) + from b12x.attention.paged import ( + plan as plan_paged_attention_scratch, + ) + from b12x.attention.paged import ( + run as paged_attention_forward, + ) + from b12x.attention.paged import ( + verify_graph_capacity as plan_verify_graph_capacity, + ) + + self._compile_paged_attention = compile_paged_attention + self._paged_attention_forward = paged_attention_forward + + def _make_plan( + page_size: int, + mode: str, + max_total_q: int, + max_batch: int, + max_work_items: int, + max_partial_rows: int, + use_cuda_graph: bool, + num_cache_pages: int, + copy_runtime_metadata: bool, + ) -> Any: + return plan_paged_attention_scratch( + B12XPagedAttentionScratchCaps( + device=self.device, + mode=mode, + dtype=self.dtype, + kv_dtype=self.kv_torch_dtype, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + max_total_q=max_total_q, + max_batch=max_batch, + max_page_table_width=max_page_table_widths[page_size], + max_work_items=max_work_items, + max_partial_rows=max_partial_rows, + # Shape-only planning tensor; runtime cache shape is + # validated by head/page geometry, not page count. + num_cache_pages=num_cache_pages, + use_cuda_graph=use_cuda_graph, + copy_runtime_metadata=copy_runtime_metadata, + ) + ) + + capture_sizes = vllm_config.compilation_config.cudagraph_capture_sizes or [] + decode_plan_sizes = { + int(size) for size in capture_sizes if 0 < int(size) <= max_num_seqs + } + decode_plan_sizes.add(max_num_seqs) + if os.getenv("VLLM_B12X_PAGED_DECODE_MAX_CHUNKS_PER_REQ"): + logger.warning_once( + "VLLM_B12X_PAGED_DECODE_MAX_CHUNKS_PER_REQ is ignored; " + "B12X owns decode graph chunk policy. Use the fixed " + "work/partial capacity controls only to constrain storage." + ) + decode_work_items_limit = _env_optional_storage_limit( + "VLLM_B12X_PAGED_DECODE_MAX_WORK_ITEMS", + allow_zero=False, + ) + decode_partial_rows_limit = _env_optional_storage_limit( + "VLLM_B12X_PAGED_DECODE_MAX_PARTIAL_ROWS", + allow_zero=True, + ) + + def _create_decode_plan(page_size: int, batch_size: int) -> Any: + max_page_table_width = max_page_table_widths[page_size] + capacity = plan_decode_graph_capacity( + device=self.device, + q_dtype=self.dtype, + kv_dtype=self.kv_torch_dtype, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + batch=batch_size, + max_cache_page_count=max_page_table_width, + window_left=self.window_left, + max_work_items=decode_work_items_limit, + max_partial_rows=decode_partial_rows_limit, + ) + plan = _make_plan( + page_size, + "decode", + batch_size, + batch_size, + capacity.max_work_items, + capacity.max_partial_rows, + True, + max_page_table_width, + True, + ) + plan.prepare_decode_graph_replay_state( + batch=batch_size, + total_q_capacity=batch_size, + max_page_table_width=max_page_table_width, + max_cache_page_count=max_page_table_width, + window_left=self.window_left, + ) + return plan + + self._create_decode_plan = _create_decode_plan + self._verify_q_per_req = 0 + if spec_config is not None: + self._verify_q_per_req = 1 + int( + getattr(spec_config, "num_speculative_tokens", None) or 0 + ) + if self._verify_q_per_req <= 1: + self._verify_q_per_req = 0 + + def _create_verify_plan(page_size: int, batch_size: int) -> Any: + if self._verify_q_per_req <= 1: + raise RuntimeError( + "B12X_ATTN verifier plan requested without speculation" + ) + max_page_table_width = max_page_table_widths[page_size] + total_q = batch_size * self._verify_q_per_req + capacity = plan_verify_graph_capacity( + device=self.device, + q_dtype=self.dtype, + kv_dtype=self.kv_torch_dtype, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + batch=batch_size, + query_len=self._verify_q_per_req, + max_cache_page_count=max_page_table_width, + window_left=self.window_left, + ) + plan = _make_plan( + page_size, + "verify", + total_q, + batch_size, + capacity.max_work_items, + capacity.max_partial_rows, + True, + max_page_table_width, + True, + ) + page_ids = torch.arange( + max_page_table_width, + dtype=torch.int32, + device=self.device, + ) + max_page_table = page_ids.unsqueeze(0).expand(batch_size, -1).contiguous() + max_cache_seqlens = torch.full( + (batch_size,), + capacity.representative_cache_seqlen, + dtype=torch.int32, + device=self.device, + ) + max_cu_seqlens_q = torch.arange( + 0, + total_q + 1, + self._verify_q_per_req, + dtype=torch.int32, + device=self.device, + ) + plan.prepare_graph_replay_state( + page_table=max_page_table, + cache_seqlens=max_cache_seqlens, + cu_seqlens_q=max_cu_seqlens_q, + active_total_q=total_q, + window_left=self.window_left, + ) + return plan + + self._create_verify_plan = _create_verify_plan + + def _create_extend_plan( + page_size: int, + batch_size: int, + q_capacity: int, + ) -> Any: + """Prepare a fixed-capacity extend plan without reading live lengths.""" + max_page_table_width = max_page_table_widths[page_size] + plan = _make_plan( + page_size, + "extend", + q_capacity, + batch_size, + _extend_work_items(page_size, q_capacity, batch_size), + 0, + True, + max_page_table_width, + False, + ) + page_ids = torch.arange( + max_page_table_width, + dtype=torch.int32, + device=self.device, + ) + max_page_table = page_ids.unsqueeze(0).expand(batch_size, -1).contiguous() + max_cache_seqlens = torch.full( + (batch_size,), + min(max_model_len, max_page_table_width * page_size), + dtype=torch.int32, + device=self.device, + ) + # Put one row in every request except the last, which owns the + # remainder. This represents the full total-Q capacity while the + # replay kernel remains responsible for packing arbitrary live + # per-request lengths from device cu_seqlens_q. + max_cu_seqlens_q = torch.arange( + 0, + batch_size + 1, + dtype=torch.int32, + device=self.device, + ) + max_cu_seqlens_q[-1] = q_capacity + plan.prepare_graph_replay_state( + page_table=max_page_table, + cache_seqlens=max_cache_seqlens, + cu_seqlens_q=max_cu_seqlens_q, + active_total_q=q_capacity, + window_left=self.window_left, + ) + return plan + + self._create_extend_plan = _create_extend_plan + decode_scratch_envelopes = { + page_size: plan_decode_graph_scratch_envelope( + device=self.device, + q_dtype=self.dtype, + kv_dtype=self.kv_torch_dtype, + num_q_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim_qk=self.head_size, + head_dim_vo=self.output_head_size, + page_size=page_size, + max_batch=max_num_seqs, + max_page_table_width=max_page_table_widths[page_size], + max_cache_page_count=max_page_table_widths[page_size], + window_left=self.window_left, + max_work_items=decode_work_items_limit, + max_partial_rows=decode_partial_rows_limit, + copy_runtime_metadata=True, + ) + for page_size in _B12X_SUPPORTED_PAGE_SIZES + } + self._decode_plans: dict[tuple[int, int], Any] = {} + self._verify_plans: dict[tuple[int, int], Any] = {} + self._extend_plans: dict[tuple[int, int, int], Any] = {} + for page_size in _B12X_SUPPORTED_PAGE_SIZES: + for batch_size in sorted(decode_plan_sizes): + self._decode_plans[page_size, batch_size] = self._create_decode_plan( + page_size, batch_size + ) + if self._verify_q_per_req > 1: + for batch_size in range(1, max_num_seqs + 1): + self._verify_plans[page_size, batch_size] = ( + self._create_verify_plan(page_size, batch_size) + ) + for batch_size in range(1, max_num_seqs + 1): + for q_capacity in self._extend_q_capacities: + # Equal capacity is necessarily one query token per + # request, which is handled by the decode plan. + if batch_size >= q_capacity: + continue + self._extend_plans[page_size, batch_size, q_capacity] = ( + self._create_extend_plan( + page_size, + batch_size, + q_capacity, + ) + ) + self._scratch_nbytes = max( + *(int(envelope.nbytes) for envelope in decode_scratch_envelopes.values()), + *(int(plan.layout.nbytes) for plan in self._verify_plans.values()), + *(int(plan.layout.nbytes) for plan in self._extend_plans.values()), + ) + + current_workspace_manager().get_simultaneous( + ((self._scratch_nbytes,), torch.uint8), + ) + + self.supports_quant_query_input = False + register_cutedsl_warmup_provider(self) + + logger.info_once( + "Using B12X_ATTN with q_heads=%d kv_heads=%d head_dim_qk=%d " + "head_dim_vo=%d window_left=%d planned_page_sizes=%s " + "verify_q_per_req=%d extend_q_capacities=%s scratch=%d bytes.", + self.num_heads, + self.num_kv_heads, + self.head_size, + self.output_head_size, + self.window_left, + _B12X_SUPPORTED_PAGE_SIZES, + self._verify_q_per_req, + self._extend_q_capacities, + self._scratch_nbytes, + ) + + def _compile_paged_extend_entry(self, page_size: int) -> None: + """Compile fixed-capacity paged-prefill entries without a live plan.""" + warmup_plans: list[tuple[int, int, Any, bool]] = [] + for batch_size in range(1, self._max_num_seqs + 1): + candidates = sorted( + (q_capacity, plan) + for (plan_page_size, plan_batch, q_capacity), plan in ( + self._extend_plans.items() + ) + if plan_page_size == page_size and plan_batch == batch_size + ) + for index, (q_capacity, plan) in enumerate(candidates): + warmup_plans.append((batch_size, q_capacity, plan, index == 0)) + if not warmup_plans: + return + + max_q_rows = max( + min(q_capacity, max(64, batch_size + 1)) + for batch_size, q_capacity, _, _ in warmup_plans + ) + q = torch.zeros( + (max_q_rows, self.num_heads, self.head_size), + dtype=self.dtype, + device=self.device, + ) + output = torch.zeros( + (max_q_rows, self.num_heads, self.output_head_size), + dtype=self.dtype, + device=self.device, + ) + kv_cache = torch.zeros( + (1, 2, page_size, self.num_kv_heads, self.head_size), + dtype=self.kv_torch_dtype, + device=self.device, + ) + key_cache, value_cache = self._kv_cache_views(kv_cache) + sinks = self._prepare_sinks(self.sinks, self.device) + (scratch_storage,) = current_workspace_manager().get_simultaneous( + ((self._scratch_nbytes,), torch.uint8), + ) + for batch_size, q_capacity, plan, execute in warmup_plans: + q_rows = min(q_capacity, max(64, batch_size + 1)) + page_table = torch.zeros( + (batch_size, plan.caps.max_page_table_width), + dtype=torch.int32, + device=self.device, + ) + cache_seqlens = torch.full( + (batch_size,), page_size, dtype=torch.int32, device=self.device + ) + cu_seqlens_q = torch.arange( + 0, + batch_size + 1, + dtype=torch.int32, + device=self.device, + ) + cu_seqlens_q[-1] = q_rows + k_descale = None + v_descale = None + if _is_b12x_fp8_kv_cache(self.kv_cache_dtype): + k_descale = torch.ones( + (), dtype=torch.float32, device=self.device + ).expand(batch_size) + v_descale = torch.ones( + (), dtype=torch.float32, device=self.device + ).expand(batch_size) + binding = plan.bind( + scratch=scratch_storage, + q=q[:q_rows], + k_cache=key_cache, + v_cache=value_cache, + output=output[:q_rows], + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + window_left=self.window_left, + attention_sink_bias=sinks, + k_descale=k_descale, + v_descale=v_descale, + ) + self._compile_paged_attention(binding=binding) + if execute: + # Compile-only warmup does not launch the device-side compact + # scheduler. One execution per batch covers the capture-static + # metadata variant shared by its Q-capacity plans. + self._paged_attention_forward(binding=binding) + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + common_key = ( + "b12x_paged_extend", + str(self.device), + str(self.dtype), + str(self.kv_torch_dtype), + self.num_heads, + self.num_kv_heads, + self.head_size, + self.output_head_size, + self.window_left, + self.sinks is not None, + ) + return tuple( + CuTeDSLCompileUnit( + name="b12x_paged_extend", + key=(*common_key, page_size), + compile=partial(self._compile_paged_extend_entry, page_size), + ) + for page_size in _B12X_SUPPORTED_PAGE_SIZES + ) + + def _prepare_sinks( + self, + sinks: torch.Tensor | None, + device: torch.device, + ) -> torch.Tensor | None: + if sinks is None: + return None + if sinks.device != device: + raise RuntimeError( + "B12X_ATTN sinks must be on the same CUDA device as query." + ) + if sinks.dtype == torch.float32 and sinks.is_contiguous(): + return sinks + key = ( + int(sinks.data_ptr()), + tuple(sinks.shape), + tuple(sinks.stride()), + str(sinks.dtype), + str(sinks.device), + ) + cached = self._sinks_cache.get(key) + if cached is not None: + cached.copy_(sinks) + return cached + if _capture_alloc_forbidden(): + raise RuntimeError( + "B12X_ATTN would convert attention sinks during CUDA graph " + "capture. Warm the layer eagerly or store sinks as contiguous " + "float32." + ) + cached = sinks.to(dtype=torch.float32, device=device).contiguous() + self._sinks_cache[key] = cached + return cached + + def _prepare_fp8_descales( + self, + layer: AttentionLayer, + num_reqs: int, + device: torch.device, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if not _is_b12x_fp8_kv_cache(self.kv_cache_dtype): + return None, None + if num_reqs <= 0: + raise ValueError("B12X_ATTN fp8 KV descale request count must be positive.") + + def _prepare(scale: torch.Tensor, name: str) -> torch.Tensor: + if scale.device != device: + raise RuntimeError(f"B12X_ATTN {name} must be on the query device.") + if scale.dtype != torch.float32: + raise RuntimeError(f"B12X_ATTN {name} must be float32.") + if scale.ndim == 0: + return scale.expand(num_reqs) + if scale.ndim == 1: + if int(scale.shape[0]) == 1: + return scale.expand(num_reqs) + if int(scale.shape[0]) >= num_reqs: + return scale[:num_reqs] + raise ValueError( + f"B12X_ATTN {name} must be scalar or rank-1 with at least " + f"{num_reqs} values; got shape {tuple(scale.shape)}." + ) + + return _prepare(layer._k_scale, "k_scale"), _prepare(layer._v_scale, "v_scale") + + def _kv_cache_views( + self, + kv_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + key_cache, value_cache = kv_cache.unbind(1) + key_cache = canonicalize_singleton_dim_strides(key_cache) + value_cache = canonicalize_singleton_dim_strides(value_cache) + if _is_b12x_fp8_kv_cache(self.kv_cache_dtype): + fp8_dtype = current_platform.fp8_dtype() + if key_cache.dtype == torch.uint8: + key_cache = key_cache.view(fp8_dtype) + if value_cache.dtype == torch.uint8: + value_cache = value_cache.view(fp8_dtype) + if ( + key_cache.dtype != self.kv_torch_dtype + or value_cache.dtype != self.kv_torch_dtype + ): + raise TypeError( + f"B12X_ATTN plan expects KV dtype {self.kv_torch_dtype}, got " + f"{key_cache.dtype}/{value_cache.dtype}." + ) + return key_cache, value_cache + + def _select_plan( + self, + attn_metadata: B12XPagedMetadata, + total_q: int, + q_capacity: int, + num_reqs: int, + page_size: int, + ) -> Any: + if attn_metadata.max_query_len <= 1 and int(total_q) == int(num_reqs): + batch_size = int(total_q) + plan_key = (page_size, batch_size) + plan = self._decode_plans.get(plan_key) + if plan is None: + if _capture_alloc_forbidden(): + raise RuntimeError( + "B12X_ATTN decode plan was not prepared before CUDA graph " + f"capture for page size {page_size}, batch size " + f"{batch_size}." + ) + plan = self._create_decode_plan(page_size, batch_size) + if int(plan.layout.nbytes) > self._scratch_nbytes: + raise RuntimeError( + "B12X_ATTN lazily created decode plan exceeds reserved " + f"scratch: {int(plan.layout.nbytes)} > " + f"{self._scratch_nbytes} bytes." + ) + self._decode_plans[plan_key] = plan + return plan + elif ( + self._verify_q_per_req > 1 + and attn_metadata.max_query_len == self._verify_q_per_req + and int(total_q) == int(num_reqs) * self._verify_q_per_req + ): + plan_key = (page_size, int(num_reqs)) + plan = self._verify_plans.get(plan_key) + if plan is None: + if _capture_alloc_forbidden(): + raise RuntimeError( + "B12X_ATTN verifier plan was not prepared before CUDA " + f"graph capture for page size {page_size}, batch size " + f"{num_reqs}." + ) + plan = self._create_verify_plan(page_size, int(num_reqs)) + if int(plan.layout.nbytes) > self._scratch_nbytes: + raise RuntimeError( + "B12X_ATTN lazily created verifier plan exceeds reserved " + f"scratch: {int(plan.layout.nbytes)} > " + f"{self._scratch_nbytes} bytes." + ) + self._verify_plans[plan_key] = plan + return plan + extend_q_capacity = next( + ( + capacity + for capacity in self._extend_q_capacities + if q_capacity <= capacity + ), + None, + ) + if extend_q_capacity is None: + raise ValueError( + f"B12X_ATTN extend Q capacity {q_capacity} exceeds prepared " + f"maximum {self._extend_q_capacities[-1]}." + ) + return self._extend_plans[ + page_size, + int(num_reqs), + extend_q_capacity, + ] + + def forward( + self, + layer: AttentionLayer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: B12XPagedMetadata, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + del key, value + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError( + "B12X_ATTN does not support fused output quantization." + ) + if attn_metadata is None: + return output.fill_(0) + if output.shape[-1] != self.output_head_size: + raise ValueError( + f"B12X_ATTN expected output head dim {self.output_head_size}, got " + f"{output.shape[-1]}." + ) + if kv_cache.numel() == 0: + return output.fill_(0) + + # In FULL cudagraph mode vLLM may pad attention metadata to the graph + # bucket while still passing per-layer Q/output tensors with only the + # real rows. Use tensor capacity as the launch contract and avoid + # selecting decode graph replay for padded virtual requests. + q_capacity = min( + int(query.shape[0]), + int(output.shape[0]), + ) + num_actual_tokens = min( + int(attn_metadata.num_actual_tokens), + q_capacity, + ) + if num_actual_tokens <= 0: + return output + q = query[:num_actual_tokens] + out = output[:num_actual_tokens] + if q.dtype != self.dtype or out.dtype != self.dtype: + raise TypeError( + f"B12X_ATTN plan expects dtype {self.dtype}, got " + f"q={q.dtype}, output={out.dtype}." + ) + + key_cache, value_cache = self._kv_cache_views(kv_cache) + page_size = _kv_page_size(key_cache, value_cache) + if not attn_metadata.causal: + raise NotImplementedError("B12X_ATTN supports causal attention only.") + + page_table = _ensure_i32_contiguous(attn_metadata.block_table, "block_table") + cache_seqlens = _ensure_i32_contiguous(attn_metadata.seq_lens, "seq_lens") + cu_seqlens_q = _ensure_i32_contiguous( + attn_metadata.query_start_loc, + "query_start_loc", + ) + num_reqs = int(cache_seqlens.shape[0]) + if attn_metadata.max_query_len <= 1 and num_actual_tokens < num_reqs: + num_reqs = num_actual_tokens + page_table = page_table[:num_reqs] + cache_seqlens = cache_seqlens[:num_reqs] + cu_seqlens_q = cu_seqlens_q[: num_reqs + 1] + sinks = self._prepare_sinks(self.sinks, q.device) + k_descale, v_descale = self._prepare_fp8_descales( + layer, + num_reqs, + q.device, + ) + plan = self._select_plan( + attn_metadata, + num_actual_tokens, + q_capacity, + num_reqs, + page_size, + ) + (scratch_storage,) = current_workspace_manager().get_simultaneous( + ((self._scratch_nbytes,), torch.uint8), + ) + binding = plan.bind( + scratch=scratch_storage, + q=q, + k_cache=key_cache, + v_cache=value_cache, + output=out, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + window_left=self.window_left, + active_total_q=(None if plan.caps.mode == "extend" else num_actual_tokens), + attention_sink_bias=sinks, + k_descale=k_descale, + v_descale=v_descale, + ) + self._paged_attention_forward(binding=binding) + return output + + def do_kv_cache_update( + self, + layer: AttentionLayer, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + if kv_cache.numel() == 0: + return + from vllm.v1.attention.backends.fa_utils import reshape_and_cache_flash + + key_cache, value_cache = kv_cache.unbind(1) + reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + self.kv_cache_dtype, + layer._k_scale, + layer._v_scale, + ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 646227cc4f9c..23683f842fb7 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -94,6 +94,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): ROCM_FLASHMLA_SPARSE_DSV4 = ( "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend" ) + B12X_ATTN = "vllm.v1.attention.backends.b12x_attn.B12XPagedAttentionBackend" FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend" FLASH_ATTN_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend" From 9faf6bdfeca516076fb0edb4266b100e62e85ef0 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Fri, 7 Aug 2026 15:47:42 -0700 Subject: [PATCH 5/5] docs: document B12X backend usage Document installation, backend selection, supported configurations, and key limitations for the optional B12X integrations. Assisted-by: OpenAI Codex Signed-off-by: Luke Alonso --- docs/features/quantization/b12x.md | 52 ++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 docs/features/quantization/b12x.md diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md new file mode 100644 index 000000000000..3f1bdc35db5c --- /dev/null +++ b/docs/features/quantization/b12x.md @@ -0,0 +1,52 @@ +# B12X Backends + +[B12X](https://pypi.org/project/b12x/) provides optional CUDA kernels for +NVIDIA SM120 and SM121 GPUs. Install the dependency with: + +```bash +uv pip install "vllm[b12x]" +``` + +The B12X backends are opt-in. To select all supported B12X kernels explicitly: + +```bash +vllm serve \ + --linear-backend b12x \ + --moe-backend b12x \ + --attention-backend B12X_ATTN +``` + +Only pass `--moe-backend b12x` for a compatible NVFP4 or MXFP4 MoE model. The +linear and attention backends can be selected independently. + +B12X uses MXFP8 activations by default for MXFP4 MoE and the checkpoint's +activation format for NVFP4 MoE. MXFP4 falls back to BF16 when its A8 path does +not support the model configuration. Set `VLLM_B12X_MOE_FORCE_A16=1` to force +BF16 activations for either weight format. + +For more targeted linear selection, leave `--linear-backend` at `auto` and +enable B12X for one or both quantization families: + +```bash +VLLM_USE_B12X_FP8_GEMM=1 VLLM_USE_B12X_FP4_GEMM=1 \ + vllm serve +``` + +An explicit non-B12X `--linear-backend` takes precedence over these environment +variables. + +## Supported Configurations + +| Backend | Supported configurations | +| ------- | ------------------------ | +| Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, and MXFP4 | +| MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations | +| Attention | Causal decoder MHA/MQA/GQA with BF16 model dtype; BF16, FP16, or FP8 E4M3 KV cache; head sizes 64, 128, 192, or 256; block size 64 or 128 | + +The attention backend supports decode, prefill, mixed batches, speculative +verification, sliding-window attention, attention sinks, and CUDA graphs. + +The B12X MoE backend does not support expert parallelism, expert maps, EXL3, or +NF3. The attention backend does not support MLA, non-causal attention, ALiBi, +logits soft capping, or context parallelism. Dense W4A16 layers are not handled +by B12X and continue to use another compatible backend such as Marlin.