CUDA: let any expert count use the fast mm_ids_helper path

[ServeurpersoCom]

Overview

CUDA: unlock the fast mm_ids_helper path for any expert count, enable it for 10 (Qwen3.8-Flash-Next)

Additional information

mm_ids_helper groups tokens by expert before the MoE matmuls. It has a fast path that spreads warp lanes across tokens, and a slow generic one that walks a single token per iteration with a warp reduction each time.

The fast path only accepted expert counts dividing 32, plus a hardcoded case for 6, so anything else quietly fell back to the slow path. Qwen3.8-Flash-Next uses 10, and nsys put this kernel at 13.3% of prefill time, ahead of flash attention.

The real constraint isn't the expert count itself but the number of lanes per token, and the loop already guarded the extra lanes. So the padding just rounds up to the next power of two. Every count already dispatched gets the same padding as before and is unchanged. I added 10; any other count is one line away.

RTX PRO 6000 at 55k context, first run discarded: prefill 2334 to 2600 t/s 2139 -> 2351 t/s, +9.9% (current master 62acc89). Generation doesn't move, with a single token there's nothing to walk. Only tested on CUDA with 32 wide warps.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[ServeurpersoCom]

I checked exact head af8ae552a8914a4c84f6931c9eb8debcbf207615 on an RTX 3080 Ti (sm_86, CUDA 13.3). I compiled the proposed power-of-two helper and an sm_86 CUDA probe for the optimized-path n_expert_used = 10 lane guard, then ran the probe: neu_padded=16, two groups per warp, 20 valid-lane loads, 12 padded-lane skips, bad=0 (padded_ok=1). This is bounded compile/pad-guard evidence, not an end-to-end llama.cpp runtime or output-correctness test. I did not reproduce the RTX PRO 6000 2334 -> 2600 t/s result, so I cannot validate that performance claim. The patch adds optimized dispatch only for 10; other undispatched expert counts still use the existing generic path.

The perf claim needs a model whose expert count was not dispatched, so Qwen3.8-Flash-Next with 512 experts and 10 used. On anything using 2, 4, 6, 8, 16 or 32 there is nothing to measure, those already took the fast path.

I re-ran it on current master, reverting only f1793c1 for the control, two separate builds, alternated in both orders with the first pair discarded:

without 2161.5 2138.3 2137.8 2130.8
with 2418.6 2374.1 2349.1 2340.8

2139 -> 2351 t/s, +9.9% prefill at 55k context. A bit below the 2334 -> 2600 in the description, which came from a branch carrying other changes. (#28023 not merged for now)