SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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.

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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