ggml-cuda: size MMQ tile to MoE tokens-per-expert - #39
Merged
Conversation
For MUL_MAT_ID, choose the tile width from ~2x the average tokens-per-expert (ncols_dst/nchannels_y) instead of ncols_max, so the common per-expert tiles are filled instead of mostly empty at mmq_x=128. Assisted-by: Claude Opus 4 (1M context)
Author
|
I tried a variety of real texts and improvements remain. I'll keep the changes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this changes
An MoE prefill optimization for the MMQ (quantized matmul) path. For
MUL_MAT_ID,mul_mat_q_casepicks the tile width (mmq_x) by tiling overncols_max, which assumes every token can land in a single tile column range - the worst case of all tokens routed to one expert. In practice each expert receives only aboutncols_dst / nchannels_ytokens, so at largencols_maxthe selection lands on wide tiles (e.g.mmq_x=128) that are mostly empty per expert.This computes the tile width from
~2xthe average tokens-per-expert instead: whenexpert_bounds != nullptr,ncols_to_tile = min(2 * ceil(ncols_dst / nchannels_y), ncols_max). The2xcovers routing imbalance, and larger batches converge back toncols_max, so dense (non-MoE) matmuls are unaffected. The heuristic is restricted to RDNA3.5 (GGML_CUDA_CC_IS_RDNA3_5), the only architecture it has been tuned and validated on.Benchmarks
Measured on gfx1151 (Radeon 8060S), Qwen3.6-35B-A3B Q4_K_M,
-ngl 999 -r 1, built with-DGGML_HIP_ROCWMMA_FATTN=OFF. Baseline isgfx11at the same commit base.Prefill throughput improves substantially; decode is untouched (the change only affects batched-matmul tiling).
-r 1has no error bars, but the prefill deltas are far outside run-to-run noise.