SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

Replace SM120 W4A16 MoE kernels (#3336)

## ๐Ÿ“Œ Description

Replace the previous SM120 W4A16 MoE kernel split with a packed-route
design.
The new W4A16 path groups routed token rows by expert before GEMM work,
then drives the
fused FC1/FC2 kernel from packed route indices, block expert IDs, and a
route count.

This removes the old static, dynamic, and micro W4A16 kernel files and
adds W4A16-specific
prepared-weight, route-packing, FP4 helper, host utility, and
fused-kernel modules under the
moe_w4a16_ prefix. The FlashInfer b12x MoE dispatch, workspace
allocation, tracing, and tests
are updated to use the new W4A16 backend while preserving the existing
caller-facing interfaces.

  ## ๐Ÿ” Related Issues

  None.

  ## ๐Ÿš€ 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.

  ## ๐Ÿงช Tests

  - [x] Tests have been added or updated as needed.
  - [x] All tests are passing (unittest, etc.).
  
 ## Reviewer Notes
 
Please focus review on the W4A16 kernel replacement itself:

- the packed-route layout and route-block metadata consumed by the fused
kernel
  - the FC1/activation/FC2 dataflow over that packed routed layout
  - workspace sizing and reuse for graph-captured execution
- preservation of the existing FlashInfer b12x MoE API surface for W4A16
callers

The intended behavior for the full-expert path is that topk_ids are
valid dense expert IDs. Skipped/nonlocal routing goes through the mapped
route path.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added explicit quantization mode (quant_mode) and source_format
inputs, W4A16 route-packing execution path, and W4A16
weight-preparation/packing utilities.

* **Improvements**
* Unified quant_mode-aware workspace allocation and dispatch for
SM120/SM121 MoE kernels; wrapper and trace templates surface new inputs
and validation.

* **Bug Fixes / Stability**
* CUDA-graph capture now requires preallocated outputs and rejects
incompatible quant_mode/source_format combinations.

* **Tests**
* Expanded tests to cover quant_mode, W4A16 routing, workspace
contracts, and CUDA-graph guards.

<!-- review_stack_entry_start -->

[![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/3336)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
L
Luke Alonso committed
16a8e274c4b7b2c4501894e765c2e4a1d82c2fa5
Parent: 77477e2
Committed by GitHub <noreply@github.com> on 5/22/2026, 4:23:55 PM