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