feat(moe): add SM120 W4A16 b12x kernels (#3271)
## 📌 Description
Adds SM120 b12x MoE W4A16 support alongside the existing W4A4 path.
This PR adds activation precision plumbing for W4A4/W4A16 dispatch,
wires
W4A16 static, dynamic, and direct micro kernels, and preserves the
intended scale
semantics: fc2_input_scale is required only for
activation_precision="fp4" and
tolerated/ignored for activation_precision="bf16".
It also ports the relevant W4A4 FlashInfer integration details to W4A16,
including workspace sizing, swizzled
scale storage, current-stream direct launches, resource gating, CUDA
graph coverage, and wrapper support.
## 🔍 Related Issues
N/A
## 🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following
items are complete.
### âś… Pre-commit Checks
- [x] I have installed pre-commit by running pip install pre-commit (or
used your preferred method).
- [x] I have installed the hooks with pre-commit install.
- [x] I have run the hooks manually with pre-commit run --all-files and
fixed any reported issues.
> Ran via uvx pre-commit run --all-files; all hooks passed.
## đź§Ş Tests
- [x] Tests have been added or updated as needed.
- [x] All tests are passing (unittest, etc.).
Tested with:
uvx pre-commit run --all-files
git diff --check
/home/luke/projects/sglang/.venv/bin/python -m compileall \
flashinfer/cute_dsl/fp4_common.py \
flashinfer/cute_dsl/utils.py \
flashinfer/fused_moe/cute_dsl/b12x_moe.py \
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py \
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py
\
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_dynamic_kernel.py
\
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_micro_kernel.py
\
flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_w4a16_static_kernel.py
\
flashinfer/trace/templates/moe.py \
tests/moe/test_b12x_fused_moe.py
FLASHINFER_DISABLE_VERSION_CHECK=1
/home/luke/projects/sglang/.venv/bin/python \
-m pytest tests/moe/test_b12x_fused_moe.py -k "direct_micro" -q
FLASHINFER_DISABLE_VERSION_CHECK=1
/home/luke/projects/sglang/.venv/bin/python \
-m pytest tests/moe/test_b12x_fused_moe.py -k "activation_precision or
w4a16" -q
## Reviewer Notes
The direct W4A16 micro path is intentionally limited to cases where
local expert IDs match global expert IDs;
expert-parallel remapping falls back to the compact static W4A16 path.
The compact W4A16 micro fallback present upstream is not wired because
upstream dispatch leaves it inactive;
unsupported direct-micro cases fall back to static W4A16.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Configurable activation precision (fp4 or bf16) for fused
Mixture-of-Experts with new dynamic, static, and low‑latency micro
kernel variants.
* Exposed CUDA stream helper for interoperability.
* **Improvements**
* New GPU synchronization, atomic, non‑coherent global and relaxed
shared‑memory primitives for robust kernel coordination.
* Expanded FP4/FP8/BF16 quantize/dequantize, packing, dot‑product
helpers and vectorized scatter; precision‑aware workspace/dispatch
tuning and optional fc2 input‑scale semantics.
* **Tests**
* Broadened coverage and trace schema updates to validate
activation‑precision behaviors, W4A16 paths, and wrapper/trace
expectations.
[](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3271)
<!-- end of auto-generated comment: release notes by coderabbit.ai --> L
Luke Alonso committed
5ef7afa3905c28ca1e26f77df2042becb3f1bdfa
Parent: 0a128d1
Committed by GitHub <noreply@github.com>
on 5/11/2026, 6:55:07 AM