Fused MoE Kernel Notes

This note records implementation invariants for VeOmni’s GPU fused MoE kernels. The DeepSeek-V4 hash-MoE CI failure exposed one of these invariants, but the invariant itself is generic and applies to the non-EP fused_triton MoE path.

Scatter/gather layout

The non-EP fused MoE path expands tokens by their top-k routes before running per-expert grouped GEMMs:

original hidden states: [tokens, hidden]
scattered activations: [tokens * topk, hidden]

scatter_index maps each (token, topk_slot) pair into the expert-grouped scattered tensor. Later, moe_gather() sums the routed expert outputs back to [tokens, hidden].

Backward grouped-GEMM launch bound

For group_gemm_same_nk, max_M is a per-expert launch bound. It must cover the largest expert group in the scattered tensor. The total scattered row count is a conservative safe bound:

grad_fc2_output = moe_scatter(grad_output, scatter_index)
num_scattered_tokens = grad_fc2_output.shape[0]

Use num_scattered_tokens for the downstream group_gemm_same_nk(..., max_M=...) calls in non-EP backward. The matching group_gemm_same_mn(..., max_K=...) value is kept consistent with the scattered input row count, but the stale-row correctness issue is about group_gemm_same_nk under-covering its output rows.

Do not use the original token count as the backward grouped-GEMM bound. It is only safe when every expert group has at most tokens scattered rows.

Why ordinary top-k routing may hide the bug

Consider tokens = 128 and topk = 2.

If each token routes to two different experts in a balanced pattern:

token0 -> expert0, expert1
token1 -> expert0, expert1
...
token127 -> expert0, expert1

After grouping by expert:

expert0: 128 rows
expert1: 128 rows

Each expert group has at most tokens rows, so a stale bound like max_M = tokens may still cover every group. The implementation would still be relying on the wrong invariant, but this routing pattern may not expose it.

Duplicate top-k routes

Some routers can assign multiple top-k slots of the same token to the same expert. DeepSeek-V4 hash-MoE can do this when the token-to-expert lookup maps both slots to the same expert.

For tokens = 128 and topk = 2:

token0 -> expert0, expert0
token1 -> expert0, expert0
...
token127 -> expert0, expert0

After grouping by expert:

expert0: 256 rows

With BLOCK_M = 128, a stale max_M = 128 launches enough blocks for only the first 128 rows of that expert group. Rows 128-255 can remain uncomputed or stale, and later gather/gradient operations may consume arbitrary values, including inf or nan. A visible symptom is a non-finite gradient norm such as gnorm = nan.

Regression test pattern

A focused regression test should explicitly construct duplicate top-k routes:

selected_experts = zeros([tokens, topk])
routing_weights = full([tokens, topk], weight)

Then compare split-fc1 fused, merged-fc1 fused, and eager MoE backward results, and assert all produced gradients are finite.