From 2740c817ffbe6c1980507c7cbd27b4bc94d2c856 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Fri, 21 Aug 2026 08:04:36 -0700 Subject: [PATCH] [Kernel] Add b12x FP4 MoE backend (#52018) Signed-off-by: Luke Alonso Signed-off-by: mgoin Signed-off-by: Misha Goin Co-authored-by: mgoin --- .buildkite/test_areas/kernels.yaml | 25 +- docs/features/quantization/b12x.md | 27 +- setup.py | 2 +- tests/kernels/moe/test_b12x.py | 1245 +++++++++++++++++ tests/kernels/quantization/nvfp4_utils.py | 17 +- tests/model_executor/test_b12x_warmup.py | 12 +- tests/quantization/test_auto_round.py | 35 +- vllm/config/kernel.py | 2 + vllm/envs.py | 5 + vllm/model_executor/layers/fused_moe/b12x.py | 770 ++++++++++ .../model_executor/layers/fused_moe/config.py | 4 + .../layers/fused_moe/oracle/mxfp4.py | 78 +- .../layers/fused_moe/oracle/nvfp4.py | 64 +- .../compressed_tensors_moe_w4a4_mxfp4.py | 18 +- .../compressed_tensors_moe_w4a4_nvfp4.py | 3 + .../quantization/inc/schemes/inc_mxfp4_moe.py | 33 +- .../layers/quantization/modelopt.py | 2 + .../layers/quantization/mxfp4.py | 2 + .../layers/quantization/online/mxfp4.py | 1 + .../layers/quantization/quark/quark_moe.py | 1 + .../layers/quantization/utils/b12x_moe.py | 108 ++ vllm/model_executor/warmup/b12x_warmup.py | 19 +- vllm/utils/b12x.py | 4 + 23 files changed, 2417 insertions(+), 60 deletions(-) create mode 100644 tests/kernels/moe/test_b12x.py create mode 100644 vllm/model_executor/layers/fused_moe/b12x.py create mode 100644 vllm/model_executor/layers/quantization/utils/b12x_moe.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index e1fddcd7cb..e3ec0d17e2 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -343,9 +343,9 @@ steps: # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py -- label: ":nvidia: (DGX) Spark B12X Linear Kernels Nightly" +- label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly" key: b12x-linear-kernels-dgx-spark-nightly - timeout_in_minutes: 30 + timeout_in_minutes: 60 device: dgx-spark optional: true num_devices: 1 @@ -353,17 +353,34 @@ steps: - arm64-image-build source_file_dependencies: - setup.py + - vllm/config/kernel.py + - vllm/envs.py - vllm/model_executor/kernels/linear/ + - vllm/model_executor/layers/fused_moe/b12x.py + - vllm/model_executor/layers/fused_moe/config.py + - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py + - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py + - vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/ + - vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py + - vllm/model_executor/layers/quantization/modelopt.py + - vllm/model_executor/layers/quantization/mxfp4.py + - vllm/model_executor/layers/quantization/utils/b12x_moe.py - vllm/model_executor/warmup/b12x_warmup.py - vllm/utils/b12x.py + - tests/kernels/moe/test_b12x.py + - tests/kernels/quantization/nvfp4_utils.py - tests/model_executor/kernels/test_b12x_linear.py - tests/model_executor/test_b12x_warmup.py - tests/kernels/quantization/test_block_fp8.py + - tests/quantization/reference_mxfp4.py + - tests/quantization/test_auto_round.py commands: - - uv pip install --system b12x==1.2.4 + - uv pip install --system b12x==1.2.6 - pytest -v -s model_executor/kernels/test_b12x_linear.py model_executor/test_b12x_warmup.py - kernels/quantization/test_block_fp8.py -k b12x + kernels/quantization/test_block_fp8.py + kernels/moe/test_b12x.py + quantization/test_auto_round.py -k b12x - label: ":nvidia: (H100) Helion Kernels Shard %N" key: kernels-helion-test diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md index 534bc6eb60..b5f56d71cb 100644 --- a/docs/features/quantization/b12x.md +++ b/docs/features/quantization/b12x.md @@ -1,24 +1,37 @@ -# B12X Linear Backend +# b12x Linear and MoE Backends -[B12X](https://pypi.org/project/b12x/) provides optional CUDA kernels for +[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]" ``` -B12X participates in automatic kernel selection after established optimized -backends and before emulation. Select it explicitly with: +b12x linear kernels participate in automatic selection after established +optimized backends and before emulation. Select the linear and MoE backends +explicitly with: ```bash -vllm serve --linear-backend b12x +vllm serve \ + --linear-backend b12x \ + --moe-backend b12x ``` +Only pass `--moe-backend b12x` for a compatible NVFP4 or MXFP4 MoE model. The +linear and MoE 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_FP4_FORCE_A16=1` to +force BF16 activations for either FP4 weight format. + ## 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 | -Dense W4A16 layers are not handled by B12X and continue to use another -compatible backend such as Marlin. +Dense W4A16 layers are not handled by b12x and continue to use another +compatible backend such as Marlin. The b12x MoE backend does not support expert +parallelism, expert maps, EXL3, or NF3. diff --git a/setup.py b/setup.py index ba0b7216de..5463dcfcaf 100644 --- a/setup.py +++ b/setup.py @@ -1530,7 +1530,7 @@ setup( # 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.4"], + "b12x": ["b12x==1.2.6"], # 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/moe/test_b12x.py b/tests/kernels/moe/test_b12x.py new file mode 100644 index 0000000000..fa1dd4410d --- /dev/null +++ b/tests/kernels/moe/test_b12x.py @@ -0,0 +1,1245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the 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 as b12x +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 ( + dequantize_nvfp4_to_dtype, + quant_nvfp4_tensor, +) +from tests.kernels.utils import torch_moe +from tests.quantization.reference_mxfp4 import dq_mxfp4_torch +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 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.mxfp4_utils import mxfp4_quantize +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: + weight_q, scale, global_scale = quant_nvfp4_tensor( + expert_weight, + 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: + return dequantize_nvfp4_to_dtype( + tensor_fp4, + tensor_sf, + global_scale, + dtype=dtype, + device=tensor_fp4.device, + is_sf_linear_layout=True, + ) + + +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 _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) + layer = SimpleNamespace( + activation=activation, + apply_router_weight_on_input=False, + w13_weight=w1, + w2_weight=w2, + w13_weight_scale=quant_config.w1_scale, + w2_weight_scale=quant_config.w2_scale, + ) + if quant_config.weight_quant_dtype == "nvfp4": + layer.w13_weight_scale_2 = quant_config.g1_alphas + layer.w2_weight_scale_2 = quant_config.g2_alphas + if quant_config.quant_dtype is not None: + assert quant_config.a1_gscale is not None + assert quant_config.a2_gscale is not None + layer.w13_input_scale = 1.0 / quant_config.a1_gscale + layer.w2_input_scale = 1.0 / quant_config.a2_gscale + experts.process_weights_after_loading(layer) + 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, + ) + + +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) + ) + + +_SITU_REASON = "kernel supports only SiTU beta=4 and linear_beta=25" +_UNINTERLEAVED_W4A8_REASON = "kernel does not support swigluoai_uninterleave with W4A8" + + +@pytest.mark.parametrize( + "config_kwargs,overrides,weight_key,activation_key,expected_reason", + [ + pytest.param( + {"hidden_dim": 128, "activation": MoEActivation.SWIGLUOAI}, + {}, + kMxfp4Static, + None, + "kernel does not support MoEActivation.SWIGLUOAI activation", + id="interleaved-swigluoai", + ), + pytest.param( + { + "hidden_dim": 128, + "activation": MoEActivation.SWIGLUOAI_UNINTERLEAVE, + }, + {}, + kMxfp4Static, + kMxfp8Dynamic, + _UNINTERLEAVED_W4A8_REASON, + id="mxfp4-w4a8-uninterleaved-swigluoai", + ), + pytest.param( + { + "hidden_dim": 128, + "activation": MoEActivation.SWIGLUOAI_UNINTERLEAVE, + }, + {}, + kMxfp4Static, + None, + None, + id="w4a16-uninterleaved-swigluoai", + ), + pytest.param( + {"activation": MoEActivation.RELU2_NO_MUL}, + {}, + kMxfp4Static, + kMxfp8Dynamic, + "MXFP4 W4A8 supports only SiLU and SiTU", + id="mxfp4-w4a8-relu2", + ), + pytest.param( + {"hidden_dim": 128}, + {}, + kMxfp4Static, + kMxfp8Dynamic, + ( + "MXFP4 W4A8 requires hidden size divisible by 256 and per-rank " + "intermediate size divisible by 32" + ), + id="mxfp4-w4a8-alignment", + ), + pytest.param( + {"hidden_dim": 128, "intermediate_size": 48}, + {"intermediate_size_per_partition": 64}, + kMxfp4Static, + None, + "MXFP4 requires the per-rank intermediate size to be divisible by 32", + id="mxfp4-tp-scale-groups", + ), + pytest.param( + {"activation": MoEActivation.SITU}, + {"activation_situ_beta": 3.0, "activation_situ_linear_beta": 25.0}, + kMxfp4Static, + None, + _SITU_REASON, + id="situ-beta", + ), + pytest.param( + {"activation": MoEActivation.SITU}, + {"activation_situ_beta": 4.0, "activation_situ_linear_beta": 25.0}, + kMxfp4Static, + None, + None, + id="situ-standard-parameters", + ), + ], +) +def test_b12x_moe_config_support( + monkeypatch: pytest.MonkeyPatch, + config_kwargs, + overrides, + weight_key, + activation_key, + expected_reason: str | None, +) -> None: + monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True) + config = make_dummy_moe_config( + **{"hidden_dim": 256, "intermediate_size": 64, **config_kwargs} + ) + for name, value in overrides.items(): + setattr(config, name, value) + + supported, reason = B12xExperts.is_supported_config( + B12xExperts, + config, + weight_key, + activation_key, + mk.FusedMoEActivationFormat.Standard, + ) + + assert (supported, reason) == (expected_reason is None, expected_reason) + + +@pytest.mark.parametrize( + "activation_key,force_a16,expected_backend", + [ + (kMxfp8Dynamic, False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8), + (None, False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8), + ], +) +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_FP4_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 + + +def test_explicit_b12x_mxfp4_force_a16_uses_a16_contract( + 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_FP4_FORCE_A16", + True, + ) + config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64) + config.moe_backend = "b12x" + + backend, experts_cls = select_mxfp4_moe_backend( + config, + activation_key=kMxfp8Dynamic, + ) + + assert backend == Mxfp4MoeBackend.B12X_MXFP4_BF16 + 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_FP4_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_compressed_tensors_mxfp4_preserves_checkpoint_packing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + compressed_tensors_moe_w4a4_mxfp4 as ct_mxfp4, + ) + + monkeypatch.setattr( + ct_mxfp4.CutlassExpertsMxfp4, + "_supports_current_device", + lambda: False, + ) + monkeypatch.setattr( + ct_mxfp4, + "select_mxfp4_moe_backend", + lambda moe: (Mxfp4MoeBackend.B12X_MXFP4_MXFP8, B12xExperts), + ) + monkeypatch.setattr( + ct_mxfp4, + "prepare_moe_fp4_layer_for_marlin", + lambda layer: pytest.fail("b12x must not use Marlin packing"), + ) + moe_config = SimpleNamespace(w13_num_shards=2, moe_backend="b12x") + method = ct_mxfp4.CompressedTensorsW4A4Mxfp4MoEMethod(moe_config) + processed_layers: list[torch.nn.Module] = [] + fake_experts = SimpleNamespace( + process_weights_after_loading=processed_layers.append + ) + kernel = SimpleNamespace(fused_experts=fake_experts) + monkeypatch.setattr(method, "get_fused_moe_quant_config", lambda _: object()) + monkeypatch.setattr(ct_mxfp4, "make_mxfp4_moe_kernel", lambda **_: kernel) + layer = torch.nn.Module() + layer._expert_routing_tables = lambda: () + method.create_weights( + layer, + num_experts=2, + hidden_size=64, + intermediate_size_per_partition=32, + params_dtype=torch.bfloat16, + ) + w13_packed_data = layer.w13_weight_packed.data + w2_packed_data = layer.w2_weight_packed.data + + method.process_weights_after_loading(layer) + + assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr() + assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr() + assert processed_layers == [layer] + + +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_FP4_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), + (kNvfp4Dynamic, 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_FP4_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_FP4_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_FP4_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_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( + num_experts=4, + experts_per_token=2, + hidden_dim=128, + intermediate_size=64, + ), + _quant_config("mxfp4", None), + ) + prepared = SimpleNamespace( + num_experts=4, + hidden_size=128, + intermediate_size=64, + w1_fp4=torch.empty(0), + ) + layer = SimpleNamespace( + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + ) + planned_tokens = [] + launched_tokens = [] + + with pytest.raises(RuntimeError, match="process_weights_after_loading"): + experts.warmup_launches(layer, token_counts=(1,)) + experts._prepared_experts = 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, "_b12x_moe_execution_plan", fake_execution_plan) + monkeypatch.setattr(b12x, "_run_b12x_moe_plan", fake_run) + monkeypatch.setattr(experts, "_plan", fake_plan) + + warmed = experts.warmup_launches(layer, 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_reload_reprepares_current_parameters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + plan = SimpleNamespace( + discards_source_parameters=False, + quant_modes=("w4a16",), + io_dtype="bfloat16", + activation="silu", + ) + prepared_inputs: list[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = [] + + def prepare_weights(**kwargs): + prepared_inputs.append( + ( + kwargs["w1_fp4"], + kwargs["w2_fp4"], + kwargs["w1_blockscale"], + kwargs["w2_blockscale"], + ) + ) + return SimpleNamespace(plan=plan) + + extension = SimpleNamespace( + plan_weights=lambda **_: plan, prepare_weights=prepare_weights + ) + monkeypatch.setattr(b12x, "_require_b12x_fused_moe", lambda: extension) + experts = B12xExperts( + make_dummy_moe_config(num_experts=2, hidden_dim=4, intermediate_size=8), + _quant_config("mxfp4", None), + ) + layer = SimpleNamespace( + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + w13_weight=torch.full((2, 8, 2), 1, dtype=torch.uint8), + w2_weight=torch.full((2, 4, 4), 1, dtype=torch.uint8), + w13_weight_scale=torch.full((2, 8, 1), 1, dtype=torch.uint8), + w2_weight_scale=torch.full((2, 4, 1), 1, dtype=torch.uint8), + ) + + experts.process_weights_after_loading(layer) + layer.w13_weight = torch.full_like(layer.w13_weight, 2) + layer.w2_weight = torch.full_like(layer.w2_weight, 2) + layer.w13_weight_scale = torch.full_like(layer.w13_weight_scale, 3) + layer.w2_weight_scale = torch.full_like(layer.w2_weight_scale, 3) + experts.process_weights_after_loading(layer) + + assert len(prepared_inputs) == 2 + assert prepared_inputs[-1][0] is layer.w13_weight + assert prepared_inputs[-1][1] is layer.w2_weight + assert prepared_inputs[-1][2] is layer.w13_weight_scale + assert prepared_inputs[-1][3] is layer.w2_weight_scale + + +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_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), + w13_weight_scale=torch.empty(0), + w2_weight_scale=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) + assert layer.b12x_warmup_provider is experts + 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, + } + ] + + +@dataclass +class _B12xMoeCase: + hidden_states: torch.Tensor + score: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w1_ref: torch.Tensor + w2_ref: torch.Tensor + quant_config: FusedMoEQuantConfig + activation: MoEActivation + activation_dtype: str | None + topk: int = 2 + + +def _make_b12x_moe_case( + weight_dtype: str, + activation_dtype: str | None, + *, + activation: MoEActivation = MoEActivation.SILU, + tokens: int = 16, + seed: int = 19, +) -> _B12xMoeCase: + set_random_seed(seed) + num_experts, hidden_size, intermediate_size = 4, 512, 128 + dtype = torch.bfloat16 + 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 + ) + + if weight_dtype == "mxfp4": + w1_q, w1_scale = mxfp4_quantize(w1) + w2_q, w2_scale = mxfp4_quantize(w2) + w1_ref = torch.stack( + [dq_mxfp4_torch(w1_q[e], w1_scale[e], dtype) for e in range(num_experts)] + ) + w2_ref = torch.stack( + [dq_mxfp4_torch(w2_q[e], w2_scale[e], dtype) for e in range(num_experts)] + ) + 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, + ) + elif weight_dtype == "nvfp4": + 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) + ] + ) + 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=activation.is_gated, + reorder_w13=activation_dtype is None and activation.is_gated, + ) + 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, + ) + else: + raise ValueError(f"unsupported test weight dtype {weight_dtype}") + + return _B12xMoeCase( + hidden_states=hidden_states, + score=torch.randn((tokens, num_experts), device="cuda", dtype=dtype), + w1=w1_q, + w2=w2_q, + w1_ref=w1_ref, + w2_ref=w2_ref, + quant_config=quant_config, + activation=activation, + activation_dtype=activation_dtype, + ) + + +@pytest.mark.skipif(not _has_b12x_moe(), reason="requires b12x MoE on SM120") +@pytest.mark.parametrize( + "weight_dtype,activation_dtype,activation", + [ + pytest.param("mxfp4", "mxfp8", MoEActivation.SILU, id="mxfp4-mxfp8"), + pytest.param("mxfp4", None, MoEActivation.SILU, id="mxfp4-bf16"), + pytest.param("nvfp4", "nvfp4", MoEActivation.SILU, id="nvfp4-nvfp4"), + pytest.param("nvfp4", "mxfp8", MoEActivation.SILU, id="nvfp4-mxfp8"), + pytest.param("nvfp4", None, MoEActivation.SILU, id="nvfp4-bf16-silu"), + pytest.param( + "nvfp4", + None, + MoEActivation.RELU2_NO_MUL, + id="nvfp4-bf16-relu2", + ), + ], +) +@torch.inference_mode() +def test_b12x_moe_matches_torch( + weight_dtype: str, + activation_dtype: str | None, + activation: MoEActivation, + workspace_init, +) -> None: + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + case = _make_b12x_moe_case( + weight_dtype, + activation_dtype, + activation=activation, + ) + reference = torch_moe( + case.hidden_states, + case.w1_ref, + case.w2_ref, + case.score, + case.topk, + activation=case.activation, + ) + if activation_dtype == "nvfp4": + topk_weights, topk_ids, _ = fused_topk( + case.hidden_states, + case.score, + case.topk, + renormalize=False, + ) + reference = _nvfp4_activation_reference( + case.hidden_states, + case.w1_ref, + case.w2_ref, + topk_weights, + topk_ids, + case.quant_config.a1_gscale, + case.quant_config.a2_gscale, + ) + + checks_zero_canonicalization = ( + weight_dtype == "nvfp4" and activation_dtype is None + ) + if checks_zero_canonicalization: + assert _count_fp4_negative_zeros(case.w1) > 0 + assert _count_fp4_negative_zeros(case.w2) > 0 + output = _run_b12x_moe( + case.hidden_states, + case.w1, + case.w2, + case.score, + case.topk, + case.activation, + case.quant_config, + ) + if checks_zero_canonicalization: + assert _count_fp4_negative_zeros(case.w1) == 0 + assert _count_fp4_negative_zeros(case.w2) == 0 + + 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") +@pytest.mark.parametrize( + "weight_dtype,activation_dtype", + [ + pytest.param("mxfp4", "mxfp8", id="w4a8"), + pytest.param("nvfp4", None, id="w4a16"), + ], +) +@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 + + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + case = _make_b12x_moe_case( + weight_dtype, + activation_dtype, + tokens=4, + seed=23, + ) + kernel = _make_b12x_moe_kernel( + case.hidden_states, + case.w1, + case.w2, + case.topk, + case.activation, + case.quant_config, + ) + topk_weights, topk_ids, _ = fused_topk( + case.hidden_states, + case.score, + case.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=case.hidden_states, + w1=case.w1, + w2=case.w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=case.activation, + global_num_experts=case.w1.shape[0], + 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/kernels/quantization/nvfp4_utils.py b/tests/kernels/quantization/nvfp4_utils.py index df6513c131..21d0bebf15 100644 --- a/tests/kernels/quantization/nvfp4_utils.py +++ b/tests/kernels/quantization/nvfp4_utils.py @@ -43,6 +43,7 @@ def dequantize_nvfp4_to_dtype( device, block_size=16, is_sf_128x4_layout=True, + is_sf_linear_layout=False, ): """Dequantize the fp4 tensor back to high precision.""" # Two fp4 values are packed into one uint8. @@ -52,7 +53,9 @@ def dequantize_nvfp4_to_dtype( tensor_f32 = break_fp4_bytes(tensor_fp4, dtype) tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) tensor_sf = tensor_sf.view(torch.float8_e4m3fn) - if is_sf_128x4_layout: + if is_sf_linear_layout: + tensor_sf = tensor_sf[:m, : k // block_size] + elif is_sf_128x4_layout: tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) else: tensor_sf = convert_swizzled_8x4_layout_to_linear(tensor_sf, m, k, block_size) @@ -146,7 +149,15 @@ def get_nvfp4_global_scale(a: torch.Tensor): return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32) -def quant_nvfp4_tensor(a: torch.Tensor): +def quant_nvfp4_tensor( + a: torch.Tensor, + *, + is_sf_swizzled_layout: bool = True, +): a_global_scale = get_nvfp4_global_scale(a) - a_quant, a_block_scale = scaled_fp4_quant(a, a_global_scale) + a_quant, a_block_scale = scaled_fp4_quant( + a, + a_global_scale, + is_sf_swizzled_layout=is_sf_swizzled_layout, + ) return a_quant, a_block_scale, a_global_scale diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index 26069e0a9a..71e22e3023 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -181,8 +181,14 @@ def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None: worker = SimpleNamespace( get_model=lambda: SimpleNamespace(modules=modules), - scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=8, + max_num_scheduled_tokens=32, + ), model_config=SimpleNamespace(dtype=torch.float32), + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace(compile_sizes=[4, 16]) + ), ) platform = SimpleNamespace( is_cuda=lambda: True, @@ -200,7 +206,7 @@ def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None: assert scans == 1 assert calls == [ - ("first", (1, 2, 8), torch.bfloat16), - ("second", (1, 2, 8), torch.bfloat16), + ("first", (1, 2, 4, 8, 16, 32), torch.bfloat16), + ("second", (1, 2, 4, 8, 16, 32), torch.bfloat16), ] assert synchronized == [True] diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 9d48926853..a0319ab20d 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -663,12 +663,28 @@ def test_inc_mxfp4_linear_method_registers_and_processes_weights( assert captured["processed_layer"] is layer -def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel( +@pytest.mark.parametrize( + ("moe_backend", "is_xpu", "mxfp4_backend"), + [ + ("auto", True, Mxfp4MoeBackend.XPU), + ("b12x", False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8), + ], +) +def test_inc_mxfp4_moe_method_preserves_checkpoint_packing( monkeypatch, + moe_backend: str, + is_xpu: bool, + mxfp4_backend: Mxfp4MoeBackend, ) -> 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( @@ -676,11 +692,16 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel( "CutlassExpertsMxfp4._supports_current_device", lambda: False, ) - monkeypatch.setattr(current_platform, "is_xpu", lambda: True) + monkeypatch.setattr(current_platform, "is_xpu", lambda: is_xpu) monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." "select_mxfp4_moe_backend", - lambda moe: (Mxfp4MoeBackend.XPU, expected_experts_cls), + lambda moe: (mxfp4_backend, expected_experts_cls), + ) + monkeypatch.setattr( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." + "prepare_moe_fp4_layer_for_marlin", + lambda layer: pytest.fail("packed backends must not use Marlin packing"), ) monkeypatch.setattr( "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe." @@ -698,7 +719,10 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel( INCMxfp4MoEMethod, ) - expected_moe_config = SimpleNamespace(w13_num_shards=2) + expected_moe_config = SimpleNamespace( + w13_num_shards=2, + moe_backend=moe_backend, + ) method = INCMxfp4MoEMethod(moe=cast(Any, expected_moe_config)) layer = torch.nn.Module() layer._expert_routing_tables = lambda: "routing-tables" @@ -731,6 +755,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 7287c1ce92..8bd8830570 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -130,6 +130,7 @@ MoEBackend = Literal[ "flashinfer_cutlass", "flashinfer_cutedsl", "flashinfer_b12x", + "b12x", "marlin", "humming", "triton_unfused", @@ -204,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 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 diff --git a/vllm/envs.py b/vllm/envs.py index e33e2fd272..5973b58507 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -190,6 +190,7 @@ if TYPE_CHECKING: 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_B12X_MOE_FP4_FORCE_A16: bool = False VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False VLLM_TPU_USING_PATHWAYS: bool = False @@ -1621,6 +1622,10 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1")) ), + # Force b12x FP4 MoE to use BF16 activations. + "VLLM_B12X_MOE_FP4_FORCE_A16": lambda: bool( + int(os.getenv("VLLM_B12X_MOE_FP4_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.py b/vllm/model_executor/layers/fused_moe/b12x.py new file mode 100644 index 0000000000..d31ba1489d --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -0,0 +1,770 @@ +# 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 typing import Any + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +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 +from vllm.utils.b12x import ( + B12xWarmupUnit, + get_b12x_fused_moe, + reuse_packed_weight_storage, +) + +logger = init_logger(__name__) + +_B12X_MOE_MODES: dict[ + tuple[torch.dtype | str | None, torch.dtype | str | None], + tuple[str, str, str], +] = { + ("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 _require_b12x_fused_moe() -> Any: + fused_moe = get_b12x_fused_moe() + assert fused_moe is not None + return fused_moe + + +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_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: + fused_moe = _require_b12x_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: + fused_moe = _require_b12x_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 _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}" + ) + scheme = ( + quant_config.weight_quant_dtype, + quant_config.quant_dtype, + ) + try: + self._quant_mode, self._source_format, self._w13_layout = _B12X_MOE_MODES[ + scheme + ] + except KeyError as exc: + raise ValueError( + f"unsupported b12x MoE quantization scheme {scheme}" + ) from exc + 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 _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 _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) + + fused_moe = _require_b12x_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, + ) + return 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, + ) + + def _refresh_quant_config(self, layer: torch.nn.Module) -> None: + self.quant_config._w1.scale = layer.w13_weight_scale + self.quant_config._w2.scale = layer.w2_weight_scale + if self._source_format != "modelopt_nvfp4": + return + + self.quant_config._w1.alpha_or_gscale = layer.w13_weight_scale_2 + self.quant_config._w2.alpha_or_gscale = layer.w2_weight_scale_2 + if self._quant_mode in ("nvfp4", "w4a8_nvfp4"): + self.quant_config._a1.alpha_or_gscale = 1.0 / layer.w13_input_scale + self.quant_config._a2.alpha_or_gscale = 1.0 / layer.w2_input_scale + + 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: + self.quant_config._w1.scale = w1_scale + if w2_scale is not None: + self.quant_config._w2.scale = 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) + if prepared is not previous: + self._plans.clear() + 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 = layer.apply_router_weight_on_input + 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" + ) + self._source_parameters_released = False + self._refresh_quant_config(layer) + prepared = self._prepare_experts( + w1=layer.w13_weight, + w2=layer.w2_weight, + activation=layer.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) + layer.b12x_warmup_provider = self + + @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 + fused_moe = get_b12x_fused_moe() + if fused_moe is None: + 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 by process_weights_after_loading" + ) + 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: + fused_moe = _require_b12x_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 get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + assert output_dtype == self.moe_config.in_dtype + prepared = self._prepared() + activation = layer.activation + limit, alpha, beta = self._swiglu_params(activation) + + def compile() -> None: + self.warmup_launches(layer, token_counts=token_counts) + + return B12xWarmupUnit( + name="MoE", + key=( + type(self), + prepared.w1_fp4.device, + output_dtype, + self._quant_mode, + self._source_format, + self._w13_layout, + int(prepared.num_experts), + int(prepared.hidden_size), + int(prepared.intermediate_size), + int(self.moe_config.experts_per_token), + _b12x_activation_name(activation), + bool(layer.apply_router_weight_on_input), + limit, + alpha, + beta, + ), + compile=compile, + ) + + @torch.inference_mode() + def warmup_launches( + self, + layer: torch.nn.Module, + *, + token_counts: Iterable[int], + ) -> int: + """Compile one representative launch for every planned regime.""" + activation = layer.activation + dtype = self.moe_config.in_dtype + topk = int(self.moe_config.experts_per_token) + apply_router_weight_on_input = bool(layer.apply_router_weight_on_input) + limit, alpha, beta = self._swiglu_params(activation) + prepared = self._prepared() + device = prepared.w1_fp4.device + 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=topk, + prepared=prepared, + quant_mode=self._quant_mode, + apply_router_weight_on_input=apply_router_weight_on_input, + swiglu_limit=limit, + swiglu_alpha=alpha, + swiglu_beta=beta, + ) + signature = (execution_plan.implementation, execution_plan.execution) + launch_tokens.setdefault(signature, tokens) + + for tokens in launch_tokens.values(): + hidden_states = torch.zeros( + (tokens, int(prepared.hidden_size)), + dtype=dtype, + device=device, + ) + output = torch.empty_like(hidden_states) + topk_ids = ( + torch.arange(topk, device=device, dtype=torch.int32) + .unsqueeze(0) + .expand(tokens, -1) + .contiguous() + ) + topk_ids.remainder_(int(prepared.num_experts)) + topk_weights = torch.full( + (tokens, topk), + 1.0 / topk, + dtype=torch.float32, + device=device, + ) + plan = self._plan( + tokens=tokens, + topk=topk, + activation=activation, + apply_router_weight_on_input=apply_router_weight_on_input, + ) + scratch = torch.empty( + (_b12x_scratch_nbytes(plan),), + dtype=torch.uint8, + device=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=self._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 w1, w2, global_num_experts + del 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._prepared() + 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") diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 456d6c5793..350792aff9 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 e118331970..4a8c723278 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 @@ from typing import TYPE_CHECKING, Literal, Union 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,9 @@ def _pack_deepgemm_mxfp4_scales( class Mxfp4MoeBackend(Enum): NONE = "None" + # b12x backends + B12X_MXFP4_MXFP8 = "B12X_MXFP4_MXFP8" + B12X_MXFP4_BF16 = "B12X_MXFP4_BF16" # DeepGEMM FP8xFP4 backend (SM100+) DEEPGEMM_MXFP4 = "DEEPGEMM_MXFP4" # FlashInfer TRTLLM backends @@ -141,11 +145,21 @@ TRITON_BACKENDS = ( 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 import B12xExperts + + return [B12xExperts] + + elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( DeepGemmFP4Experts, ) @@ -269,6 +283,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, @@ -355,6 +370,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, @@ -441,6 +458,22 @@ 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_FP4_FORCE_A16: + return [Mxfp4MoeBackend.B12X_MXFP4_BF16] + # W4A8 is the high-throughput b12x path and is preferred when the + # model does not request an activation format. Keep this as backend + # policy instead of exposing a separate backend name per precision. + 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, @@ -456,6 +489,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 = ( @@ -464,26 +498,25 @@ 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 + if requested_backend == Mxfp4MoeBackend.EMULATION else _backend_activation_key(requested_backend) ) try: @@ -577,7 +610,7 @@ 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 activation_format == mk.FusedMoEActivationFormat.BatchedExperts: requested_backends = [ Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b @@ -636,6 +669,10 @@ 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: + # b12x plans for the exact model dimensions. B12xExperts validates the + # required MXFP4 block alignment before selecting the backend. + return hidden_size, intermediate_size if backend == Mxfp4MoeBackend.EMULATION: # Emulation has no kernel tile; it only needs OCP MX block alignment so the # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated @@ -1288,6 +1325,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, @@ -1711,6 +1758,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, @@ -1778,6 +1835,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 3f23ecd619..a860003775 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -5,6 +5,7 @@ from enum import Enum import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import envs from vllm.config.kernel import MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -17,6 +18,9 @@ from vllm.model_executor.layers.fused_moe.config import ( 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, @@ -36,6 +40,7 @@ logger = init_logger(__name__) class NvFp4MoeBackend(Enum): + B12X = "B12X" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL" @@ -67,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 import B12xExperts + + return [B12xExperts] + + elif backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsModular, TrtLlmNvFp4ExpertsMonolithic, @@ -145,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, @@ -162,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_FP4_FORCE_A16 + ) + + def select_nvfp4_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, @@ -188,6 +205,7 @@ def select_nvfp4_moe_backend( ] NVFP4_BACKENDS_WITH_CLAMP = { + NvFp4MoeBackend.B12X, NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, @@ -249,6 +267,8 @@ def select_nvfp4_moe_backend( 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 @@ -264,7 +284,7 @@ def select_nvfp4_moe_backend( f"explicitly requested moe_backend={runner_backend!r} does " f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', " f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', " - f"'marlin', or 'humming' instead." + f"'b12x', 'marlin', or 'humming' instead." ) return _return_or_raise( requested_backend, config, weight_key, activation_key, activation_format @@ -303,6 +323,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, @@ -313,7 +334,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, @@ -470,7 +520,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 ( @@ -484,12 +536,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 bb266c0f82..f7d12c8635 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 @@ -25,9 +25,11 @@ from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsMxFp4, ) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + B12X_BACKENDS, 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, @@ -46,10 +48,17 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): super().__init__(moe) self.group_size = 32 self.mxfp4_backend = Mxfp4MoeBackend.MARLIN - # use cutlass if supported, otherwise fallback to marlin for weight-only FP4 + # Backend selection must match the weight preparation below: CUTLASS + # swizzles scales, b12x and XPU consume checkpoint packing, and Marlin + # repacks weights and scales. self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() self.experts_cls: type[mk.FusedMoEExperts] - if self.use_cutlass_mxfp4: + if moe.moe_backend == "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 +147,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): w2_scale=layer.w2_weight_scale, ) else: - # W4A16: weight-only via Marlin + # b12x selects W4A8 or W4A16; XPU uses W4A4; Marlin uses W4A16. return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=layer.w13_weight_scale, @@ -188,7 +197,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight_scale = torch.nn.Parameter( torch.stack(swizzled_w2), requires_grad=False ) - elif current_platform.is_xpu(): + elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu(): pass else: logger.warning_once( @@ -208,6 +217,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): 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 40f3897166..92728691d2 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 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ): 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 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): 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 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): 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 437a744947..b1a8a294f1 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 @@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + B12X_BACKENDS, Mxfp4MoeBackend, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, @@ -42,26 +43,29 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase): """W4A4 MXFP4 group MoE for AutoRound ``auto_round:llm_compressor`` exports. Registers the packed MXFP4 layout (uint8 ``weight_packed`` + uint8 E8M0 - ``weight_scale``, ``group_size=32``) and dispatches the fused MoE to the - best backend for the current device: CUTLASS (true W4A4 on supported - GPUs), the native XPU kernel, or Marlin weight-only as a fallback. The - per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are folded - into the stacked ``w13`` / ``w2`` parameters by ``make_expert_params_mapping``. + ``weight_scale``, ``group_size=32``) and dispatches the fused MoE to b12x + when requested; otherwise it uses CUTLASS W4A4, XPU W4A4, or Marlin W4A16. + The per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are + folded into the stacked ``w13`` / ``w2`` parameters by + ``make_expert_params_mapping``. """ def __init__(self, moe) -> None: super().__init__(moe) self.group_size = 32 # Backend selection must stay consistent with the weight preparation in - # process_weights_after_loading / get_fused_moe_quant_config, which only - # implement three layouts: CUTLASS swizzle (true W4A4), the native XPU - # kernel (packed passthrough), and Marlin weight-only. XPU dispatch is - # deferred to the shared oracle; every other non-CUTLASS device falls - # back to Marlin (mirroring CompressedTensorsW4A4Mxfp4MoEMethod). + # process_weights_after_loading / get_fused_moe_quant_config, which use + # three layouts: CUTLASS swizzle (true W4A4), b12x and XPU packed + # passthrough, and Marlin weight-only. XPU dispatch is deferred to the + # shared oracle; b12x dispatch is explicit; every remaining non-CUTLASS + # device falls back to Marlin. 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 moe.moe_backend == "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 +150,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase): w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, ) - # Weight-only (Marlin) or native XPU kernel. + # b12x selects W4A8 or W4A16; XPU uses W4A4; Marlin uses W4A16. return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=layer.w13_weight_scale, @@ -194,9 +198,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase): layer.w2_weight_scale = torch.nn.Parameter( torch.stack(swizzled_w2), requires_grad=False ) - elif current_platform.is_xpu(): - # The XPU fused-MoE kernel consumes the packed layout directly; no - # swizzle / repack / transpose is required. + elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu(): pass else: logger.warning_once( @@ -216,6 +218,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase): 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 185a55d00b..2adf21b46f 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1574,6 +1574,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): 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) @@ -1610,6 +1611,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): 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 6506c1fce1..926b4312d5 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -376,6 +376,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def process_weights_after_loading(self, layer: RoutedExperts) -> None: w13 = layer.w13_weight @@ -769,6 +770,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def process_weights_after_loading(self, layer): w13 = layer.w13_weight diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py index 3607a4bdf7..d09e9ceb6b 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp4.py +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -255,6 +255,7 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def get_fused_moe_quant_config( self, layer: torch.nn.Module diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 7ef64e97a9..752d2e423a 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1334,6 +1334,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def get_fused_moe_quant_config( 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 0000000000..d78a8db541 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/b12x_moe.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Weight preparation helpers for b12x MoE kernels.""" + +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + swap_w13_to_w31, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + swizzle_blockscale, +) +from vllm.utils.math_utils import round_up + + +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]: + if w13_scale.size(1) != w13.size(1): + raise ValueError("w13 weight and scale row counts must match") + half_size = w13.size(1) // 2 + if half_size * 2 != w13.size(1): + raise ValueError("gated NVFP4 MoE weights must have even row counts") + 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") + + half_pad_size = round_up(half_size, 64) - half_size + if half_pad_size == 0: + return w13, w13_scale, w2, w2_scale + + def pad_rows(tensor: torch.Tensor) -> torch.Tensor: + shape = tensor.shape + return F.pad( + tensor.reshape(shape[0], 2, half_size, *shape[2:]), + (0, 0, 0, half_pad_size), + ).flatten(1, 2) + + return ( + pad_rows(w13), + pad_rows(w13_scale), + F.pad(w2, (0, half_pad_size // 2)), + F.pad(w2_scale, (0, half_pad_size // 16)), + ) + + +def _per_expert_scale( + scale: torch.Tensor, + num_experts: int, + name: str, +) -> torch.Tensor: + if scale.dim() == 0: + scale = scale.expand(num_experts) + elif scale.dim() == 2: + scale = scale.amax(dim=1) + if scale.dim() != 1 or scale.numel() != num_experts: + raise ValueError(f"{name} must contain one value per expert") + return scale.to(torch.float32).contiguous() + + +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 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 = swap_w13_to_w31(w13) + w13_scale = swap_w13_to_w31(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 = F.pad(w13, (0, 0, 0, pad_size)) + w2 = F.pad(w2, (0, pad_size // 2, 0, 0)) + w2_scale = F.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/b12x_warmup.py b/vllm/model_executor/warmup/b12x_warmup.py index 5e0e4a0fa4..ae645ea0a5 100644 --- a/vllm/model_executor/warmup/b12x_warmup.py +++ b/vllm/model_executor/warmup/b12x_warmup.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm B12X JIT kernels used by a loaded model.""" +"""Warm b12x JIT kernels used by a loaded model.""" from collections import Counter from collections.abc import Iterable @@ -61,9 +61,20 @@ def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None: ) if output_dtype not in (torch.bfloat16, torch.float16): output_dtype = torch.bfloat16 + compile_sizes = worker.vllm_config.compilation_config.compile_sizes or [] + max_num_batched_tokens = worker.scheduler_config.max_num_batched_tokens + serving_sizes = [ + max_num_batched_tokens, + *cudagraph_capture_sizes, + *(size for size in compile_sizes if isinstance(size, int)), + ] + max_tokens = max_num_batched_tokens + max_num_scheduled_tokens = worker.scheduler_config.max_num_scheduled_tokens + if max_num_scheduled_tokens is not None: + max_tokens = max(max_tokens, max_num_scheduled_tokens) token_counts = b12x_warmup_token_counts( - max_tokens=worker.scheduler_config.max_num_batched_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, + max_tokens=max_tokens, + cudagraph_capture_sizes=serving_sizes, ) units = _collect_warmup_units( worker.get_model(), @@ -72,7 +83,7 @@ def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None: ) for name, count in _compile_warmup_units(units).items(): logger.info_once( - "Warmed up %d B12X %s linear GEMM signatures.", + "Warmed up %d b12x %s kernel signature(s).", count, name, ) diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 8e2d7e938b..d7882fb6f3 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -52,6 +52,10 @@ def get_b12x_tensor_fp8_linear() -> ModuleType | None: return _get_submodule("b12x.gemm.tensor_fp8_linear") +def get_b12x_fused_moe() -> ModuleType | None: + return _get_submodule("b12x.moe.fused_moe") + + def b12x_warmup_token_counts( *, max_tokens: int,