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 --> [](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