vLLM b12x 后端实战：在 NVIDIA SM120/SM121 GPU 上启用 MXFP4/NVFP4 线性层与 MoE 内核

b12x.md 是 vLLM 仓库中针对 b12x 可选 CUDA 内核的官方说明。本文以此文档为骨架，结合 setup.py、vllm/config/kernel.py、vllm/model_executor/layers/fused_moe/b12x.py、vllm/envs.py 等源码与配置，系统梳理 b12x Linear/MoE 后端的安装方式、启停开关、支持的量化格式、激活精度策略、配置约束与底层内核调度机制，帮助读者在 SM120/SM121（如 RTX PRO 系列、DGX Spark 等）设备上把 FP4 量化模型真正跑出最佳性能。

一、b12x 是什么：面向 SM120/SM121 的可选 CUDA 内核

b12x 是一个独立的第三方 Python 包（vLLM 中以 optional dependency 接入），为 NVIDIA SM120 / SM121 计算能力（即 Blackwell 12x 架构族）提供高性能 CUDA 内核实现。在 vLLM 中，b12x 主要覆盖两类算子的实现：

• 线性层（Linear）GEMM 内核，源码位于 vllm/model_executor/kernels/linear/ 下的 mxfp4、mxfp8、nvfp4、scaled_mm 子目录中的 b12x.py；
• MoE 专家计算内核，即 vllm/model_executor/layers/fused_moe/b12x.py 中实现的模块化（modular）张量并行 Fused MoE 后端。

从 vLLM 的视角看，b12x 属于平台加速后端插件：内核只有在显式安装 b12x 包、并运行在支持型号的 CUDA 设备上时才会被实例化。访问 b12x 包的入口封装在 vllm/utils/b12x.py，其中通过 importlib.util.find_spec("b12x") 探测包是否已安装，并惰性导入 b12x.gemm.blockscaled、b12x.gemm.mxfp8_linear、b12x.gemm.tensor_fp8_linear、b12x.moe.fused_moe、b12x.attention.paged 等子模块——模块不存在时返回 None，不会影响其余后端。

二、安装 b12x 依赖

文档给出标准安装方式，即通过 vLLM 的 extras 一键安装：

uv pip install "vllm[b12x]"

仓库中 vllm[b12x] 这个 extra 的定义位于 setup.py（第 1527 行左右）：

"b12x": ["b12x==1.3.0"],

即该 extra 当前会将 b12x 固定钉在 1.3.0 版本，以保证与仓库中已对接的 API（如 b12x.gemm.blockscaled、b12x.moe.fused_moe 的 plan 式接口）兼容。使用 pip 且希望等价安装，可执行 pip install "vllm[b12x]"。

需要说明的适用前提：

• b12x 内核仅对 CUDA 平台、且计算能力属于 Blackwell 12x 家族（SM120/SM121） 的设备生效；
• 该依赖可选：未安装时，vLLM 自动选择流程会跳过 b12x 相关后端，其余已有后端（Triton、CUTLASS、Marlin、DeepGEMM、FlashInfer、Humming 等）照常可用。

三、启用方式：显式指定 Linear / MoE 后端

b12x 的线性层内核与 MoE 内核可以独立开关、独立组合。官方文档给出的启动命令为：

vllm serve <model> \
    --linear-backend b12x \
    --moe-backend b12x

其中：

• --linear-backend b12x 显式将线性层 GEMM 内核固定为 b12x；
• --moe-backend b12x 显式将 MoE 专家计算内核固定为 b12x。

3.1 仅在 MoE 模型上使用 b12x MoE 后端

官方文档特别强调：只有当模型是兼容的 NVFP4 或 MXFP4 MoE 模型时，才需要（也才应该）单独传 --moe-backend b12x。原因在于 b12x MoE 后端只实现 FP4 权重的专家计算路径（见下文"支持的配置"表格），对权重不满足条件的 MoE 层会自动落到其他兼容实现。

两个参数默认值均为 auto，且互不绑定，因此可以只切 MoE、只切 Linear，或按模型实际情况组合。

3.2 参数对应的源码定义

--moe-backend 与 --linear-backend 在服务层解析后落到 vllm/config/kernel.py 的 KernelConfig：

• MoEBackend 与 LinearBackend 都是字符串字面量联合类型，其中都包含 "auto"、"b12x"、"emulation" 等取值（见 MoEBackend = Literal[...]、LinearBackend = Literal[...] 定义）；
• 两个字段的 docstring 都说明了 "b12x" 的含义：
  • moe_backend 的 "b12x"："Use b12x FP4 MoE kernels on SM12x"；
  • linear_backend 的 "b12x"："Use native B12X FP8 and FP4 linear kernels on SM12x"；
• 字段校验器会把用户传入的字符串先 lower() 并把 - 规范为 _（_normalize_moe_backend / _normalize_linear_backend），因此诸如 --moe-backend B12X、--linear-backend b12x 写法均可被接受；
• 默认 auto 时由"后端 oracle"根据模型量化格式与硬件能力自动挑选。

3.3 自动选择中的优先级

文档指出 b12x 线性内核在自动选择流程中参与的顺序是：在既有优化后端之后、模拟（emulation）后端之前。也就是说当 --linear-backend auto 时，若模型权重格式（FP8/FP4 各变体）没有比 b12x 更优先的既有后端可用，且设备与安装条件满足，则会选中 b12x；只有当 b12x 等真实内核均不可用时，才会退化到逐元素反量化到 BF16/FP16 后做普通 GEMM 的 emulation 路径（该路径在 KernelConfig docstring 中被明确标注为"仅用于测试"）。

3.4 FlashInfer 桥接下的 b12x 变体

除原生 b12x 后端外，KernelConfig 中还列出了经 FlashInfer 分发的两个变体：flashinfer_b12x（MoE 侧："FlashInfer CuteDSL fused MoE for SM12x"）与 flashinfer_b12x（Linear 侧："FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)"）。二者与原生 b12x 属于不同代码路径，本文以官方 b12x 文档所述的原生 b12x 后端为准。

四、支持的配置

官方文档用一张表格给出 b12x 支持的全部配置：

后端	支持的配置
Linear（线性层）	Per-tensor FP8、128x128 block FP8、MXFP8、NVFP4、MXFP4
MoE	张量并行（TP）MXFP4 权重 + BF16 或 MXFP8 激活；NVFP4 权重 + BF16、NVFP4 或 MXFP8 激活

4.1 Linear 五种量化格式的内核落点

在 vLLM 内核目录中，五种 Linear 配置分别由对应子目录下的 b12x.py 承接：

• Per-tensor FP8 / 128x128 Block FP8：走 vllm/model_executor/kernels/linear/scaled_mm/b12x.py，运行时分别调用 b12x 的 mm_block_fp8（block-scaled GEMM）与 b12x.gemm.tensor_fp8_linear 子模块，scale 经 _upcast_e8m0_to_fp32 处理；
• MXFP8：走 vllm/model_executor/kernels/linear/mxfp8/b12x.py，使用 b12x.gemm.mxfp8_linear；
• NVFP4：走 vllm/model_executor/kernels/linear/nvfp4/b12x.py；
• MXFP4：走 vllm/model_executor/kernels/linear/mxfp4/b12x.py，其核心实现把激活先用 FlashInfer 的 MXFP4 量化例程（flashinfer_mxfp4_quantize(..., backend="cute-dsl")）压成 x_packed + x_scale_swizzled，再调用 b12x 的 blockscaled.mm_mxfp4。

这些内核类的 is_supported 检查都一致要求：CUDA 平台 + current_platform.is_device_capability_family(120)（即 SM12x 家族），否则返回不可用原因，例如源码中的提示："B12X MXFP4 kernels require a Blackwell 12x device"。

4.2 MoE 五种激活/权重组合的映射

张量并行 MoE 路径的"权重格式 + 激活格式 → b12x 模式"映射定义在 vllm/model_executor/layers/fused_moe/b12x.py 的 _B12X_MOE_MODES 表中：

权重格式	激活格式	b12x quant_mode	源格式	w13 布局
MXFP4	MXFP8	w4a8_mx	fp4_e8m0_k32	w31
MXFP4	BF16	w4a16	fp4_e8m0_k32	w31
NVFP4	NVFP4	nvfp4	modelopt_nvfp4	w31
NVFP4	MXFP8	w4a8_nvfp4	modelopt_nvfp4	w31
NVFP4	BF16	w4a16	modelopt_nvfp4	w13

在 fused_moe/b12x.py 中，B12xExperts 是一个模块化 Fused MoE 专家实现（继承 mk.FusedMoEExpertsModular），其构造会校验权重量化格式必须是 mxfp4 或 nvfp4（quant_config.weight_quant_dtype not in ("mxfp4", "nvfp4") 时直接抛 ValueError）。权重在 process_weights_after_loading 阶段通过 b12x 的 plan API 做打包预处理（plan_weights + prepare_weights），把 FP4 权重整理为适合 SM12x 执行的布局，w1_blockscale、w2_blockscale 与 NVFP4 所需的全局 scale / 激活全局 scale 一并传入。

五、激活精度策略与 MXFP4 BF16 回退

5.1 默认策略

官方文档对激活精度选择做了明确说明，分两种权重格式：

• MXFP4 MoE：默认使用 MXFP8 激活（W4A8 高吞吐路径）；
• NVFP4 MoE：使用 checkpoint 里记录的激活格式（即权重随附的 activation format，w1_input_scale/w2_input_scale 决定，可为 NVFP4、MXFP8 或 BF16）；
• MXFP4 在其 A8 路径不支持当前模型配置时，自动回退到 BF16 激活（W4A16）。

5.2 源码中的选择逻辑

MXFP4 MoE 的后端挑选逻辑集中在 vllm/model_executor/layers/fused_moe/oracle/mxfp4.py：

• Mxfp4MoeBackend 定义了 B12X_MXFP4_MXFP8 与 B12X_MXFP4_BF16 两个候选；
• 当 runner_backend == "b12x" 时，_get_requested_backends 遵循文档所述策略——模型没有显式请求激活格式时优先返回 MXFP8（W4A8）候选，注释明确写道 "W4A8 is the high-throughput b12x path and is preferred when the model does not request an activation format"；
• 随后逐个用对应内核类的 is_supported_config 检查维度约束，不满足就抛出原因；也就是说"MXFP4 回退 BF16"本质上是按模型配置约束逐项筛选内核类的结果。

5.3 维度约束举例（来自 B12xExperts.is_supported_config）

从 fused_moe/b12x.py 的 is_supported_config 可以看到 W4A8 / W4A16 判定时实际校验的硬性条件：

• 不带 bias（"kernel does not support expert biases"）；
• 输入输出 dtype 仅支持 FP16 / BF16；
• MXFP4（W4A16 与 W4A8 均）要求每个 rank 的 intermediate size 能被 32 整除；
• MXFP4 + MXFP8（W4A8）额外要求 hidden size 能被 256 整除，且激活函数仅支持 SiLU / SiTU（"MXFP4 W4A8 supports only SiLU and SiTU"）；
• SiTU 激活只接受固定参数 beta=4, linear_beta=25；
• swigluoai_uninterleave 激活不支持 W4A8（activation-key 非空时）；
• 支持激活集为 SiLU、SiTU、SWIGLUOAI_UNINTERLEAVE、RELU2_NO_MUL。

如果你的 MXFP4 MoE 模型因 hidden/intermediate 维度或激活函数不满足这些约束而命中不了 W4A8，vLLM 就会走 W4A16（BF16 激活）路径——这正是文档所说"MXFP4 回退到 BF16"的落点。

5.4 强制 BF16 激活的环境变量

对两种 FP4 权重格式（MXFP4 与 NVFP4），均可通过环境变量强制使用 BF16 激活而跳过 W4A8：

VLLM_B12X_MOE_FP4_FORCE_A16=1 vllm serve <model> --moe-backend b12x

该环境变量在 vllm/envs.py 中登记（默认 False，读取字符串 "1" 判定为真），并分别作用于两条 MoE oracle 决策路径：

• MXFP4：见 vllm/model_executor/layers/fused_moe/oracle/mxfp4.py，置 1 时强制返回 B12X_MXFP4_BF16；
• NVFP4：见 vllm/model_executor/layers/fused_moe/oracle/nvfp4.py。

注：在 vLLM 环境变量命名中 False 是 VLLM 默认布尔值（环境变量通常默认关闭）；如需临时验证 W4A16 与 W4A8 的吞吐差异，用置 1 的方式强制即可。

六、边界与限制：什么情况 b12x 不接管

官方文档在末尾给出三条重要边界，这些限制都能在 fused_moe/b12x.py 的实现中得到印证：

1. 稠密 W4A16 层不由 b12x 处理，继续走其他兼容后端（如 Marlin）。这与 b12x 的定位一致：b12x Linear 覆盖的是 FP8/MXFP8/NVFP4/MXFP4 稠密 GEMM，而权重仅为 INT4（W4A16 weight-only）的层由 marlin 等后端承接，二者可共存于同一模型（因为 --linear-backend 只约束实现了该格式的层，未实现的层按 KernelConfig.linear_backend docstring 所述回落到自动选择）。

2. b12x MoE 后端不支持专家并行（Expert Parallelism）。B12xExperts._supports_parallel_config 要求 use_ep=False 且 ep_size == 1，同时关闭 use_all2all_kernels 与 EPLB。它只承担"张量并行 + 每 rank 本地专家"的 MoE 场景，这也解释了为何文档支持表把 MoE 限定为 Tensor-parallel。

3. 不支持 expert maps、EXL3、NF3。对应实现为 supports_expert_map() 返回 False，apply() 里对非空 expert_map 直接抛错（"b12x TP MoE does not support expert maps"）。

此外还有从源码中可以直接确认、文档未逐条列出的限制：

• 不支持 LoRA：B12xExperts.moe_sum 直接 raise NotImplementedError("LoRA is not supported for B12xExperts")；
• 不支持带 bias 的 MoE；
• apply_router_weight_on_input 只在 W4A16（BF16 激活）模式下可用，W4A8 组合会抛错；
• FP4 权重打包为 b12x 布局后，如果 checkpoint 声明 discards_source_parameters，vLLM 会释放原始 FP4 权重参数（源码中通过 _release_source_parameters 将其替换为空张量）以节省显存，并在 CUDA Graph 捕获前完成全部 prepare/plan 工作——b12x MoE 的权重预处理与执行 plan 不允许在 CUDA capture 期间发生（对应多个 RuntimeError 检查）。

七、执行与预热机制：plan / bind / run 与 JIT 预热

b12x MoE 使用"plan（计划）→ bind（绑定）→ run（执行）"三段式异步内核 API，封装在 fused_moe/b12x.py：

• plan：按 (token 数, topk, 激活函数, apply_router_weight_on_input) 作为 key 缓存执行计划 fused_moe.plan(...)（frozen=True），同一 key 只建一次；
• bind/run：把 scratch（来自 MoE workspace2 的 uint8 暂存区）、隐藏状态 a、已打包的 experts、topk_weights、topk_ids 绑定后执行；topk_ids / topk_weights 非标准 dtype 或非连续时会在进入内核前规整为 int32 / fp32，CUDA capture 期间禁止此类分配；
• workspace：workspace_shapes() 按计划计算并返回 (0,)（无 workspace13）与由 scratch 字节数换算的 workspace2 形状，运行时 _workspace_as_b12x_scratch 会校验 workspace2 连续性且字节数不小于 plan.scratch_specs() 要求。

预热（JIT warmup） 是 b12x 这类 CuteDSL/计划式内核低延迟运行的前提。B12xExperts 把自身注册为层的 warmup provider（layer.b12x_warmup_provider = self），对外暴露 B12xWarmupUnit（见 vllm/utils/b12x.py）。预热时（warmup_launches）会对每个 token 规模求一次执行计划，并按 (implementation, execution) 签名去重，为每种内核策略只跑一次代表性启动，把 JIT 编译耗时挡在首 token 之前。整体预热框架见 vllm/model_executor/warmup/kernel_warmup.py 与 vllm/model_executor/warmup/b12x_warmup.py。

八、实操小结：启用 b12x 的检查清单

要在当前 vLLM 仓库版本的模型服务上正确启用 b12x，按以下顺序确认即可：

1. 硬件确认：nvidia-smi 确认 GPU 为 SM120/SM121（Blackwell 12x 家族）且运行 CUDA 平台；
2. 依赖安装：uv pip install "vllm[b12x]"（等价钉住 b12x==1.3.0，见 setup.py）；
3. 后端选择：对稠密线性层场景用 --linear-backend b12x；对兼容的 NVFP4/MXFP4 MoE 模型用 --moe-backend b12x，两者独立；
4. 精度策略：默认 MXFP4 MoE 走 W4A8（MXFP8 激活）；模型配置不满足 W4A8 维度约束时自动回退 W4A16；需要对照验证时设置 VLLM_B12X_MOE_FP4_FORCE_A16=1 强制 BF16 激活；
5. 边界检查：确认模型未依赖专家并行、expert map、EXL3、NF3、MoE bias 或 LoRA，且稠密 W4A16 层落在 Marlin 等既有后端的默认自动选择范围；
6. 观察验证：启动日志会打印选中后端与 kernel oracle 的决策信息；首次启动时观察 b12x JIT 预热日志，确认计划式内核完成编译后再压测吞吐。

通过以上配置，即可在 SM120/SM121 GPU 上同时利用 b12x 的高吞吐 FP4 线性层与 MoE 内核，将 NVFP4/MXFP4 量化模型在该平台上的性能潜力完整释放。