cuda: device-routed MoE prefill (on-device expert counting-sort)

[jasstrong]

Summary

The large-batch MUL_MAT_ID (MoE) path currently routes experts on the host: it copies the ids tensor device→host, then runs a CPU loop over experts × tokens × top-k to group token-slots by expert, bracketed by two cudaStreamSynchronizes. The GPU sits idle through all of it.

For a 256-expert model that loop is O(ne02 × ne12 × n_expert_used) and grows linearly with the ubatch size, while the GEMM it feeds gets more efficient per token as the batch grows. So the fixed-per-token CPU routing becomes a larger and larger slice of prefill as ubatch increases.

Measured stall (std::chrono around the loop + syncs, Qwen3-MoE 35B-A3B on gfx1100):

ubatch	host-routing stall as % of prefill
512	~9.5%
2048	~24.8%

What this changes

Replaces the host loop with an on-device counting sort (ggml/src/ggml-cuda/moe-devsort.cu): histogram of tokens per expert → exclusive scan → scatter. It writes ids_to_sorted / ids_from_sorted straight into the existing device buffer; only tokens_per_expert (one int32 per expert) is copied back to the host, which is still needed to size the per-expert GEMMs.

• Bit-exact. The device sort produces the same routing as the host loop; it only changes the order of rows within an expert, which cannot change any per-token GEMM result. Verified: greedy llama-cli generation is token-for-token identical to the host path on TQ3 (gfx1100) and TQ4_1S (gfx1030).
• Weight-type agnostic. It sorts ids; it never touches weights. Helps every MoE model that reaches this path, not just the turbo types.
• Arch-neutral. Only atomicAdd(int) and a small serial scan — no warp/wave-size assumptions — so it also carries to CDNA/MFMA parts unchanged.
• On by default, with GGML_MOE_HOST_ROUTE=1 to force the original host loop (kept intact as the else branch, GGML_ASSERT and all).

Benchmarks

llama-bench, no draft model, Qwen3-MoE 35B-A3B, ubatch 2048:

GPU	arch	quant	pp512	pp2048
RX 7900 XTX	gfx1100 / RDNA3	TQ3_1S	837 → 987 (+18%)	1384 → 1925 (+39%)
Radeon PRO V620	gfx1030 / RDNA2	TQ4_1S	415 → 438 (+5.5%)	711 → 781 (+9.9%)

(RDNA2's smaller gain is expected — no matrix cores, so its GEMM is ~2× slower and the fixed routing loop is a smaller fraction of the longer prefill.)

Notes / scope

• Correctness was validated end-to-end (bit-identical generation) rather than via test-backend-ops, because its MUL_MAT_ID cases are absorbed by the mmq/mmf paths before reaching this host-routed fallthrough, so they don't exercise it. Happy to add a targeted test if you'd like one.
• The host loop assumes distinct experts per token (its existing GGML_ASSERT enforces this); the device path counts each (token, slot) once, which agrees for all valid top-k routing.

[giveen]

@TheTom

PR #322 Review — device-routed MoE prefill (on-device expert counting-sort)

Repo/PR: #322
Author: jasstrong
Base branch: feature/turboquant-kv-cache
Files changed: ggml/src/ggml-cuda/ggml-cuda.cu, ggml/src/ggml-cuda/moe-devsort.cu (new), ggml/src/ggml-cuda/moe-devsort.cuh (new)
Reviewed: 2026-08-29

Verdict

Solid, well-targeted, correct. Recommend merge, with two minor asks for the author (below).

What it does

Replaces the CPU-side MoE expert-routing loop in the large-batch MUL_MAT_ID fallback with an
on-device counting sort (histogram → exclusive scan → scatter). The old path copied the ids
tensor device→host, ran an O(experts × tokens × top-k) CPU loop to group token-slots by expert,
and uploaded the result — bracketed by two cudaStreamSynchronize calls that left the GPU idle.
The new path only copies back tokens_per_expert (one int32 per expert), needed for per-expert
GEMM sizing. Toggle: GGML_MOE_HOST_ROUTE=1 forces the original host loop.

Scope — does this only affect TQ3_1S/TQ4_1S models?

Not hardcoded to those types, but in practice, yes, on this fork. The changed function
(ggml_cuda_mul_mat_id) tries three faster paths first — mmvq, mmq, mmf — and only falls
through to this routing code when all three decline:

• mmq supports a fixed switch of types (Q4_0/Q4_1/Q5_0/Q5_1/Q8_0, K-quants, IQ*, MXFP4/NVFP4).
  TQ3_1S/TQ4_1S are not in that list.
• mmf refuses any quantized type outright.
• mmvq explicitly excludes TQ weights by name.

Confirmed empirically: a stock Q4_K_XL MoE model on the test RTX 5090 never reaches this code at
all (mmq claims it first, since the card is Turing+). TQ3_1S/TQ4_1S are the only weight types on
this fork that always land here — any other future unlisted quant type would too, but for
mainstream quant types on a modern NVIDIA GPU this code path is effectively dead.

Static review

• Traced how tokens_per_expert / ids_to_sorted / ids_from_sorted are consumed by the
  downstream per-expert GEMM loop (ggml-cuda.cu:2267): the device path's exclusive-scan offsets
  correctly reproduce the same expert-ordered contiguous blocks the host loop produced.
  Intra-expert row order differs (atomic scatter vs. deterministic host order) but is provably
  harmless — each row's GEMM result is independent of its position, and ids_from_sorted is
  derived from the same scatter position used to write the row, so the inverse gather stays
  consistent regardless of order.
• Scratch buffers (tpe_dev, fill_dev) come from ggml_cuda_pool_alloc, which is not
  zero-initialized — the PR correctly zeros both before the atomic histogram/scatter kernels.
• The static const bool moe_host_route = getenv(...) idiom matches an existing pattern already
  used elsewhere in this codebase (common.cuh:1263), so it's stylistically consistent rather
  than a one-off convention.
• No CMakeLists changes needed or made — ggml-cuda's CMakeLists globs *.cu/*.cuh, so the two
  new files are picked up automatically. Confirmed by inspection.
• The existing comment above this fallback path notes it "should not be reached when recording
  CUDA graphs, because it requires stream synchronization" — the new code preserves exactly one
  cudaStreamSynchronize (down from two), so this invariant still holds.

Empirical verification

Built the PR in an isolated git worktree against the actual project GPU (RTX 5090, 32GB). No TQ
model existed locally, so requantized Qwen3-Coder-30B-A3B-Instruct (Q4_K_XL source) to TQ3_1S
and TQ4_1S via llama-quantize --allow-requantize to exercise the changed code.

Bit-exactness (device-sort route vs. GGML_MOE_HOST_ROUTE=1), same seed, greedy decoding:

• Short prompt (13 tokens): output identical.
• Large prefill (~3,300 tokens, AGENTS.md as input, -ub 2048): all 60 generated tokens
  identical; only the reported t/s differed.

Performance (llama-bench, TQ3_1S, -ngl 999, device route vs. host route):

ubatch	pp512	pp2048
512	+10.2%	+10.3%
2048	+10.2%	+33.7%

This reproduces the PR's own claimed pattern — larger gains at larger ubatch — on different
hardware than the author benchmarked (RTX 5090 here vs. RX 7900 XTX / Radeon PRO V620 in the PR),
which is a good independent-hardware sanity check on the underlying mechanism.

Requested changes (minor, not blocking)

1. No bounds check on expert index in the device kernels. atomicAdd(&tpe[e], 1) for an
   out-of-range e is an out-of-bounds write. The host path only had a debug-only assert
   (stripped in release builds), so this isn't a regression, but device-side OOB is worse if it's
   ever hit. A cheap if (e >= 0 && e < ne02) guard would be safer than relying on "can't happen."
2. k_moe_excl_scan is single-threaded over ne02 experts. Fine at today's expert counts
   (~128–256), but worth a one-line comment noting that assumption so it gets revisited if a
   much-higher-expert-count model shows up.

The author's own note about missing test-backend-ops coverage (their MUL_MAT_ID cases are
absorbed by mmq/mmf before reaching this path) is worth taking them up on for lasting regression
protection, though it's not a blocker given the bit-exact verification above.

[jasstrong]

Thanks for the thorough review and the independent 5090 bit-exact/perf confirmation — much appreciated.

Both addressed:

• Bounds-guarded the expert index in k_moe_hist, and also in k_moe_scatter — its offsets[e]/fill[e] and the derived scatter write have the same OOB exposure — threading ne02 through both kernels. Well-formed ids never trip it, so it's a pure safety net.
• Added the note that k_moe_excl_scan is deliberately single-threaded over the small expert count, flagged to revisit for much higher counts.

On test-backend-ops coverage: agreed a MUL_MAT_ID case that reaches this path is worth adding for lasting regression protection — I'll follow up separately.