SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

fix(sm12x): fix micro-kernel workspace sizing when routed_rows > num_local_experts (#3191)

## Summary
Two bugs in the `b12x_fused_moe` micro-kernel path (SM120/SM121,
triggered when `routed_rows <= micro_cutover`, typically 20–40):
- **`allocate_sm120_static_workspace`** (`moe_dispatch.py`):
`compact_topk_ids` was sized `state_E` (num local experts), but the
micro-kernel fills it with `flat_ids` of length `routed_rows =
num_tokens * num_topk`. When `num_tokens * num_topk >
num_local_experts` (e.g. 2 tokens × 8 topk = 16 pairs, 8 local experts),
this caused an assertion failure: `compact_topk_ids buffer too small: 8
< 16`. Fix: size as `max(state_E, max_rows)`.
- **`compact_topk_ids`** (`triton_compact.py`): validation required
`weight_expert_ids.numel() >= total_pairs`. This was wrong — the Triton
kernel writes to `weight_expert_ids` only at indices
`0..active_expert_count-1`, bounded by `state_E` (num
unique active experts), not `total_pairs`. The check rejected valid
calls where `total_pairs > state_E`. Fix: remove the check (with
explanatory comment).
Both bugs surface together whenever the batch is small enough to hit the
micro-kernel path but `num_tokens * num_topk > num_local_experts`.
## Test plan
- [ ] Unit tests in `tests/kernels/moe/test_flashinfer_b12x_moe.py`
cover the small-batch micro-kernel path — 24/24 pass with this fix on
DGX Spark (SM121)

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

* **Bug Fixes**
* Increased pre-pass compaction workspace to prevent micro-kernel buffer
overruns and related runtime assertions.

* **Documentation**
* Clarified compaction behavior: the compacting kernel writes only up to
the active expert count, so the prior strict size validation for the
weight-expert ID buffer was removed.

* **Tests**
* Added regression tests for cases where token count × top-k exceeds
local expert count.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Meenakshi Venkataraman <meenakshiv@nvidia.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Ka-Hyun Nam <knam@nvidia.com>
M
meena-at-work committed
14f2beef319eed9950bad5f2f88f9fea2bf0fd1d
Parent: 42bf79e
Committed by GitHub <noreply@github.com> on 5/7/2026, 7:38:47 AM