iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel] Add b12x FP4 MoE backend (#52018)

Browse Source 
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: Misha Goin <mgoin64@gmail.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Luke Alonso
2026-08-21 08:04:36 -07:00
committed by GitHub
co-authored by mgoin
parent 5ee84d3c52
commit 2740c817ff
23 changed files with 2417 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+21 -4
View File
@@ -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
docs/features/quantization/b12x.md
+20 -7
View File
@@ -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 <model> --linear-backend b12x
vllm serve <model> \
--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.
setup.py
+1 -1
View File
@@ -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
tests/kernels/moe/test_b12x.py
+1245
View File
File diff suppressed because it is too large Load Diff
tests/kernels/quantization/nvfp4_utils.py
+14 -3
View File
@@ -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
tests/model_executor/test_b12x_warmup.py
+9 -3
View File
@@ -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]
tests/quantization/test_auto_round.py
+30 -5
View File
@@ -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
vllm/config/kernel.py
+2
View File
@@ -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
vllm/envs.py
+5
View File
@@ -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"))
vllm/model_executor/layers/fused_moe/b12x.py
+770
View File
@@ -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")
vllm/model_executor/layers/fused_moe/config.py
+4
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+68 -10
View File
@@ -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,
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+60 -4
View File
@@ -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:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+14 -4
View File
@@ -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,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+3
View File
@@ -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(
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
+18 -15
View File
@@ -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,
vllm/model_executor/layers/quantization/modelopt.py
+2
View File
@@ -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
vllm/model_executor/layers/quantization/mxfp4.py
+2
View File
@@ -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
vllm/model_executor/layers/quantization/online/mxfp4.py
+1
View File
@@ -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
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
View File
@@ -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
vllm/model_executor/layers/quantization/utils/b12x_moe.py
+108
View File
@@ -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
vllm/model_executor/warmup/b12x_warmup.py
+15 -4
View File
@@ -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,
)
vllm/utils/b12x.py
+4
View File
@@ -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,