Repository navigation
CUDA: faster MoE prompt processing (compact MMQ tiling, device expert sort, GET_ROWS) - #264
danielhanchen wants to merge 22 commits into
Conversation
The MUL_MAT_ID grid had ceil(ne12/J) column tiles for every expert. At prompt sizes each expert sees a small part of the tokens (2048 tokens, 512 experts, 10 used: ~40 rows per expert), so ~15 of 16 blocks exited at once, while each still occupied a whole SM (MMQ runs one 256-thread block per SM). Expert e now owns the virtual tiles starting at expert_bounds[e]/J + e, which leaves room for its ceil(rows/J) tiles; there are expert_bounds[n]/J + n of them in total. Each warp maps a virtual tile to (expert, tile) with a warp-wide search over expert_bounds. The tiles and their arithmetic are unchanged. Used when it gives fewer tiles than the per-expert grid, so small batches keep the old grid.
get_rows launched one 256-thread block per gathered row. For rows of 1-16 elements (MoE router weights, sparse attention block selection: up to 1M rows of 1 or 4 values) almost every thread was idle. Such rows now use one thread per element (or per 16-byte vector) over all rows.
…T_ROWS rows test mode: 64 experts, 6 used, 300 tokens (empty and partial tiles), and GET_ROWS with rows of 1-65 elements. perf mode: prompt-sized MoE shapes (256 experts / 8 used, 512 experts / 10 used) and short-row GET_ROWS.
The tile width was picked for ne12 columns (128 for any prompt), so an expert with ~16-64 rows filled a quarter to half of its tile. With the compact MoE tiling, the many tiles of a narrow width no longer turn into an empty grid, so on NVIDIA pick the width for ne12*n_expert_used/n_expert rows (as RDNA3 and RDNA4 already do) whenever the compact tiling is used for it; otherwise the width for ne12 is kept. mmq_moe_compact_ntx holds the condition for both the choice and the launch.
… the compact tiling The reason to pick the tile width for ne12 was the empty per-expert grid, which the compact tiling removes on any GPU. RDNA3 and RDNA4 keep their existing choice.
The cuBLAS fallback path in ggml_cuda_mul_mat_id copied the ids tensor to the host, bucketed the tokens by expert in an O(n_experts * n_tokens * n_expert_used) triple loop, then copied the permutation back. That cost two stream synchronizations per MoE node and a host loop that grows with the expert count, which is 256 for Qwen3.6-35B-A3B. Reuse mm_ids_helper, which the MMQ and MMF paths already use to build the same permutation on the device. One call with write_inverse == false yields the gather map and the expert bounds; the scatter map is the inverse of the helper's ids_dst, obtained with a small new scatter kernel rather than a second helper call. The ordering matches the old host loop element for element, not merely group for group: mm_ids_helper assigns compact index nex_prev + it_compact, where nex_prev counts all (token, slot) pairs belonging to a strictly lower expert and it_compact increments in token order within an expert. The scatter map stays consistent because it inverts the helper's own output rather than being derived separately. One synchronization remains: the per-expert token counts size the GEMMs and are needed on the host. Removing it requires a grouped GEMM. The host loop is kept as a fallback for the cases mm_ids_helper cannot handle: n_tokens*4 exceeding the shared memory limit, a non-unit ids->nb[0], or the helper's 2^22 token and 2^10 expert-used packing limits. Qwen3.6-35B-A3B BF16 prefill, compute-only 32->128 chunk slope, warmup discarded: 2,324 tok/s -> 7,051 tok/s (median of 6 and 4 reps, base bracketing patched) Logits are byte identical on both a dense and an MoE model, and test-backend-ops -o MUL_MAT_ID passes 869/869. (cherry picked from commit 7ce0db983ad1f5666976bcb449a0a6b4ac5ad8c1) (cherry picked from commit 468799e92fba960967cc76a83aca9b538dcbbe83)
mm_ids_helper launches one block per expert and each block is a single warp that walks every token, so the work is O(n_experts * n_tokens) regardless of how many assignments there actually are. That was fine for the MMQ and MMF callers it was written for, where the token count is small. Pointing it at mul_mat_id prefill is not: with 256 experts and 32,768 tokens it does 8.4M id reads to produce 262K assignments, spread over 256 warps of a 148 SM GPU. Measured, it was 216.6 ms over 120 launches, 13.2% of all GPU kernel time on Qwen3.6-35B-A3B at -b/-ub 32768, second only to the gated delta net kernel. Replace it, above a size threshold, with a stable counting sort: per-block histogram over experts, an exclusive scan over (expert, block), then a scatter where each block rewalks its own tokens in order. That is O(assignments) and fully parallel. Stability is not incidental. An atomics-based scatter would be cheaper still, but it would produce a different permutation on every run, and the project gates these changes on byte-identical logits. Keeping the scatter stable keeps that gate meaningful even though the permutation order within an expert does not change any output value. The single-warp kernel is kept and still selected below the threshold, where its one launch beats three, and GGML_CUDA_MM_IDS_SORT pins either implementation for testing. Qwen3.6-35B-A3B, -b/-ub 32768: 216.6 ms -> 2.66 ms over the five kernels Qwen3.6-35B-A3B, -b/-ub 2048: 35.0 ms -> 6.51 ms total GPU kernel time at 32768: 1,640.2 ms -> 1,428.4 ms Bit-identical on both the MoE and the dense model; test-backend-ops passes MUL_MAT_ID 869/869 with the counting sort pinned, with the legacy kernel pinned, and on auto. (cherry picked from commit cc47f114b71c999ba7da1d8336841c468602b1dc)
…n release builds Two things in the mul_mat_id path. mm_ids_invert wrote ids_inverse[ids_dst[ic]] with no bound on the index. The counting sort fills every slot, but the legacy single-warp path leaves holes when a token selects the same expert twice, and a hole holds recycled pool memory. That turned a malformed ids tensor into an unbounded out-of-bounds device write rather than a wrong answer. The host fallback used plain assert on the expert id, which compiles out in a release build, so it checked nothing in any build worth benchmarking. Promoted to GGML_ASSERT. The comment above the expert_bounds assert also overstated what that assert catches, and is corrected: the counting sort counts every assignment, so a repeated expert still totals ne_get_rows and passes there. Also wraps the __syncwarp calls in mm_ids_warp_sync, since HIP has no __syncwarp and every kernel involved runs one warp per block, where a block barrier is an exact substitute. (cherry picked from commit 24c0824ddc794c20ca7a368a7b44da6e669d3a14)
get_rows only took the wide-access kernel when src0 and dst had the same type; a converting copy (F32 -> BF16/F16) moved one element per thread. Handle VEC elements per thread, with the wider of the two accesses 16 bytes, alignment checked per side. Same ggml_cuda_cast per element, bit-identical. Taken from staging-llama.cpp 94f70cf60 (cherry picked from 1ea166c18), applied next to the short-row path.
…s, stream-k fixup, unaligned views The one-thread-per-element get_rows kernel was used for up to UINT32_MAX - 256 elements. fastdiv adds umulhi(n, mp) + n in 32 bits, which overflows for n >= 2^31 (ne00=3 with 1073741825 rows gave wrong elements), and the grid-stride step i += gridDim.x*blockDim.x wraps for n > 2^31 + 256. Larger gathers now use the per-row kernels. Tests: MUL_MAT_ID with 256 and 512 experts, with k = 2048 so that the stream-k fixup runs over the compact tiling with empty tiles, and GET_ROWS on src0 views whose rows are not 16 byte aligned.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: def8dbddcc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // ================================================================================================= | ||
| // Counting sort: the same mapping, built with work proportional to the number of routed | ||
| // assignments rather than to n_experts * n_tokens. | ||
| // | ||
| // mm_ids_helper above gives each expert one single-warp block and has it walk every token, so it |
There was a problem hiding this comment.
This banner continues for 20 lines, includes benchmark history, and splits several sentences across comment lines. Reduce it to the non-obvious invariant, such as the stable-ordering requirement, and keep performance evidence outside the code; repository guidance requires concise comments and explicitly prohibits hard-wrapping sentences.
AGENTS.md reference: AGENTS.md:L77-L84
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Shortened the banner to one line and trimmed the other comments this PR adds to 1-2 lines, without benchmark history (68475a8).
| bool ggml_cuda_mm_ids_helper_supported(const int n_experts, const int n_tokens) { | ||
| // n_expert_used only bounds the legacy path's bit packing, and any value that fails it also | ||
| // fails the 2^10 assert there, so the counting sort is what decides support in that case. | ||
| return mm_ids_counting_sort_possible(n_experts) || mm_ids_legacy_possible(n_tokens, 1); |
There was a problem hiding this comment.
When the per-expert histogram exceeds smpb, the counting sort is unavailable, but this can still return true because it tests the legacy path with n_expert_used == 1. If the real tensor uses at least 1024 experts and its token buffer otherwise fits, ggml_cuda_mul_mat_id selects the device path, mm_ids_use_counting_sort selects the legacy kernel, and launch_mm_ids_helper aborts on its 10-bit packing assertion instead of using the existing host fallback. Pass the real n_expert_used into this predicate so unsupported shapes remain on the host path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed: with more experts than shared memory fits for the counting sort and n_expert_used >= 1024, the device path was chosen and hit the 10-bit assert. The support check now takes the real n_expert_used (044b574), and a new MUL_MAT_ID case (32768 experts, 1024 used) that aborted before now passes.
| for (int m : {128, 320}) { | ||
| test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 128, 8, b, m, 256, 2048)); |
There was a problem hiding this comment.
These cases run in the normal correctness suite, where every supported backend result is compared against the CPU reference implementation. From test_mul_mat_id::op_flops, the two shapes here execute about 3.76 billion operations per type/broadcast pair; repeating them for five types and both broadcast modes adds roughly 37.6 billion CPU-reference operations from these lines alone. This can substantially extend or time out test-backend-ops across the backend matrix; retain the empty-tile/fixup coverage with fewer representative types or smaller dimensions.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Cut the k=2048 cases to Q4_0 and Q4_K with m=128 and n=128, about 2.1 G CPU-reference ops instead of 37.6 G (23971ed). Profiling shows the stream-k fixup kernel still launches for both types.
|
@codex review |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
The single-warp helper leaves holes in the permutation for repeated ids, which the bounds checks do not catch, so the cuBLAS fallback returned wrong output where the host sort used to reject the input. Sort on the device only when the counting sort is used. An empty batch no longer selects the counting sort, which would launch an empty grid.
…ed faster The counting sort was picked for n_experts*n_tokens >= 65536, which made small float MoE ops slower: 4 launches replace the single-warp kernel, and MMF gets no tiling gain to offset them (F32 E=512 T=128 +50% on B200, f16 E=512 U=8 T=128 +20-50% on T4/L4/RTX PRO 6000). Sweeps over E 64-512, U 2/8, T 8-8192 on T4, L4 and RTX PRO 6000 put the MMQ/MMF crossover at 256-512 tokens independently of the expert count, so: - MMQ/MMF: the counting sort from 512 tokens on (or when the single-warp kernel cannot run), else the single-warp kernel as before. - cuBLAS fallback: the device sort, now always the counting sort (the only one with a complete permutation for repeated ids), from n_tokens*n_expert_used >= 64, else the host sort. The host sort no longer allocates the compact map. - Compact MMQ tiling: not for up to 16 tokens on Blackwell, nor for up to 16 tokens with at most 1 column per 16 experts, where it was measured up to 1.6x (RTX PRO 6000) and 5-8% (T4, L4) slower than the dense tiling. GGML_CUDA_MM_IDS_SORT=0/1 still forces either sort.
The two maps are adjacent in ids_buf_dev, so the cuBLAS fallback's host sort uploads them with one cudaMemcpyAsync again instead of two.
…rt_used >= 16384 The host sort loops over all ids once per expert, so its cost grows with the expert count too. With n_tokens*n_expert_used >= 64 the device sort was picked for E=64 U=8 T=16 where it measured 2-4% slower (L4, RTX PRO 6000). From E*T*U >= 16384 on it was faster in every measured cell on T4, L4 and RTX PRO 6000.
The single-warp helper counts the columns of the lower experts per slot but writes one column per token, so a token that repeats an expert in its row (valid for ggml_mul_mat_id, e.g. test-backend-ops dup=1) left unwritten columns inside that expert's range. MMQ then read stale pool contents as src1/dst indices: wrong output and illegal memory accesses (pre-existing, the abort of test-backend-ops MUL_MAT_ID dup=1 on every build). Count the matching slots too, and when they exceed the columns written, write one column per slot instead. Ids without repeats keep the old path; the repeated case is a slow per-token loop, which top-k routing never reaches. The q8_0/q4_0 dup=1 cases now pass; the f16 ones still fail in MMF, which keeps one slot per token and expert (pre-existing, no out-of-bounds access).
The host sort stopped at the first slot of a token that uses the expert, so a token repeating an expert came out short and the size assert aborted (pre-existing). With the device sort now picked only from n_experts*n_tokens*n_expert_used >= 16384, small shapes with repeated ids took that abort where the counting sort used to handle them. Give each matching slot its own row, as the device sort does; ids without repeats are unchanged. tests: MUL_MAT_ID dup=1 for IQ1_M, which goes through the cuBLAS fallback (host sort up to n = 33, device sort at n = 64).
…path does The histogram skips ids outside [0, n_experts), so the permutation comes out short and MMQ / MMF would read entries that were never written. The prefix kernel now checks the total against n_tokens*n_expert_used and traps, matching the host assert of the cuBLAS fallback.
This branch has not been deployed
1 participant
Faster prompt processing for MoE models on CUDA. Mixture-of-experts prompts spend a large share of their time in
mul_mat_id(MMQ) and in the expert gather around it. This PR fixes three things there:expert_bounds[e] / J + e).Jis picked for the rows of an average expert instead of the largest. This is on for every vendor that uses the compact tiling; RDNA3/4 keep the upstream per-expert pick.mm_ids_helpernow builds the token-to-expert permutation with a counting sort on the GPU.mul_mat_idthe src1 gather goes from 112 to 28.8 us per launch.n_flat <= INT32_MAX. Beyond 2^31 elements,fastdiv(umulhi(n, mp) + nin 32 bits) and the grid-stride step overflow, which an exhaustive probe confirmed. Larger gathers use the existing per-row kernels.Results
RTX PRO 6000 Blackwell (G4), llama-bench
-fa 1 -ngl 999, ub 2048 unless noted. Each value is the median of 3 fresh processes in rotated order; A/A (base twice) agreed within 1.1%.Correctness
The head
test-backend-opsbinary was run against both the base and head libraries, so both arms see the same cases.GGML_CUDA_MM_IDS_SORT=1/0).dup=1MUL_MAT_ID cases fail identically on base (upstream).llama_pinset.py verify --cudaof the current pin set plus this branch on b11565 passes: merge, pin contract, CPU gate, full CUDA build, test-chat, test-backend-ops, and the GPU feature matrix with all 7 features. It also merges cleanly on b11574.Not covered
Base is upstream
08ffba21e(base/upstream-08ffba21e), so the diff is only this change. Like every pin inscripts/unsloth/pr-set.json, it is merged onto the aged upstreamb####tag for the prebuilt and dropped once upstream has the equivalent.