perf(cuda): batched int8-MMQ trunk prefill gated off for MoE — OLMoE prefill 97 vs llama.cpp 17,137 t/s (176×)

[pekkah]

Part of #405 (README vs llama.cpp CUDA benchmark sweep). Biggest gap found.

OLMoE-1B-7B fully offloads to the GPU (-g -1), yet its CUDA prefill is ~176× slower than llama.cpp on the same box, and prefill is even slower than its own decode — the tell-tale sign of a per-token prefill loop.

Measured (RTX 4070 Ti, llama.cpp CUDA b9529)

	README (sharpi)	llama.cpp -ngl 99 pp2048	gap
prefill	97.4 t/s	17 137 t/s	176×
decode	126.0 t/s	426.1 t/s (tg128)	3.4× (see #408)

OLMoE prefill (97.4) < OLMoE decode (126.0) ⇒ prefill is running token-by-token, not as a batched trunk.

Root cause (verified)

CudaForwardPass.IsBatchedPrefillSupported() (src/SharpInference.Engine/CudaForwardPass.cs:2618) disables the fast int8-MMQ batched-trunk prefill (#136/#141) for any MoE model:

if (_isMoE || _tqEnabled || _hasAttnBias || !_hp.IsNeoxRope) return false;

_isMoE trips for OLMoE, so Prefill() falls through to the per-token Forward() loop (CudaForwardPass.cs:~2557-2572) — every prompt token re-reads all weights, no tensor-core MMQ, no batched attention.

The work

Extend the batched-trunk prefill to dense-attention MoE models (OLMoE: standard attention + a routed-MoE FFN). The trunk attention path is identical to the supported dense models; the new piece is a batched MoE FFN over N prompt tokens:

1. Run the batched attention trunk (already exists — PrefillBatchedTrunk).
2. For the FFN, batch the router over N tokens, then do a grouped/segmented expert GEMM (gather tokens per expert → per-expert GEMM-N → scatter), instead of the per-token expert matvec. The CPU-MoE prefill already groups tokens per expert (PrefillBatchedCpuMoe in CudaHybridGdnForwardPass) — mirror that bucketing for the on-GPU full-offload FFN.
3. Relax the _isMoE gate to allow OLMoE-shaped models (dense attention, on-GPU experts) once the batched FFN exists; keep it closed for the GDN/CPU-MoE hybrids that have their own path.

Argmax-stable (int8 MMQ) is the accepted contract, matching the dense models. Validate greedy parity vs the per-token loop and re-measure prefill.

References

• src/SharpInference.Engine/CudaForwardPass.cs:2618 (IsBatchedPrefillSupported), :2511 (prefill gate), PrefillBatchedTrunk
• MMQ prefill: Gemma 4 CUDA perf: collapse per-token kernel launches + enable batched prefill #136 (batched trunk), Gemma 4 GPU prefill is ~70x slower than llama.cpp — adopt cuBLAS GEMM for prefill #141 (Q8_0/Q4_K MMQ), perf(cuda): split-K + better-tiled int8 MMQ GEMM — 23-34% → 60-80% of TC peak (#141 follow-up) #149 (SoA repack)
• Grouped-expert-GEMM prior art: Carnice prefill toward llama.cpp 676 t/s — per-layer rewrite levers (op-offload follow-up to #387) #388 (lever 3, MegaBlocks-style segmented GEMM for routed experts)