Repository navigation
metal : fix mul_mat_id with duplicate ids - #30373
Merged
Merged
Conversation
kernel_mul_mm_id_map0 assumed at most one match per token and a fixed capacity of ne21 entries per expert, which produced wrong or dropped output rows and could overflow into the slice of the next expert when the ids contain duplicates. build a compact per-expert list of the (token, slot) entries instead: - pass 1 counts the entries of every expert, thread 0 computes the offsets of the expert slices (tpe now holds [base, count]) - pass 2 writes the entries with one unconditional store per token (conditional stores of divergent lanes are very slow), the dummy value of the tokens without a match lands on a padding slot of the slice that is never read - duplicates are handled by a rare branch - kernel_mul_mm_id reads the base/count of its expert and loops over the entry tiles that do not fit into the grid The ids buffer is sized as ne20*ne21 + ne02 (entries + padding). 883/883 test-backend-ops MUL_MAT_ID tests pass, llama-bench is within run-to-run variance of the previous behavior. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
use FOR_UNROLL for the id loops of kernel_mul_mm_id_map0 and turn ne02 into the function constant FC_mul_mm_id_map0_ne02 (FC_MUL_MM + 7), so that the threadgroup offsets of cnt/bas/sids and the prefix-sum bound fold at compile time. the pipeline name now encodes ne02 (same convention as the ssm_conv nc/silu constants) and the map0 pipeline is compiled with the constant value set. ne02 is no longer part of ggml_metal_kargs_mul_mm_id_map0 - the dispatch grid still comes from the tensor ne[2] on the host. no measurable perf change: 883/883 test-backend-ops MUL_MAT_ID tests pass, op-level and llama-bench results are within run variance. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
github-actions
Bot
added
ggml
follow the approach of the vulkan implementation (#29998): instead of launching a grid over experts x ne21 entry tiles and letting most of the threadgroups early exit, kernel_mul_mm_id_map0 now also emits a row map - the total tile count plus one (expert, first entry, valid rows) descriptor per 32-entry tile that has work. kernel_mul_mm_id is launched with a static upper bound on the number of tiles (full tiles over all entries plus at most one partial tile per expert), reads its descriptor and processes a single tile - the tail early exits. This removes the round loop that was needed to cover duplicate ids past the grid (register pressure: th_max 768 -> back to the plain mul_mm level) and duplicates are now handled implicitly, since they only produce more descriptors within the bound. The grid shrinks by ~8-16x on typical shapes (e.g. 128 experts, 8 used, 512 tokens: 128 tiles instead of 2048 entry tiles, of which 94% only launched to exit) and ne21 is no longer needed by the mm kernel kargs. 883/883 test-backend-ops MUL_MAT_ID tests pass. llama-bench (Qwen3.6-35B-A3B Q8_0) is now faster than master on all pp sizes (pp2048/4096/8192 +2%, pp512 +0.5%), tg32 unchanged. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
ggerganov
marked this pull request as ready for review
October 12, 2026 08:35
[no ci]
This was referenced Oct 12, 2026
3 of 6 tasks
1 task done
edwardyoon
pushed a commit
to edwardyoon/focus-llama
that referenced
this pull request
Oct 24, 2026
* metal : fix mul_mat_id limits and kernel name * metal : fix mul_mat_id with duplicate ids kernel_mul_mm_id_map0 assumed at most one match per token and a fixed capacity of ne21 entries per expert, which produced wrong or dropped output rows and could overflow into the slice of the next expert when the ids contain duplicates. build a compact per-expert list of the (token, slot) entries instead: - pass 1 counts the entries of every expert, thread 0 computes the offsets of the expert slices (tpe now holds [base, count]) - pass 2 writes the entries with one unconditional store per token (conditional stores of divergent lanes are very slow), the dummy value of the tokens without a match lands on a padding slot of the slice that is never read - duplicates are handled by a rare branch - kernel_mul_mm_id reads the base/count of its expert and loops over the entry tiles that do not fit into the grid The ids buffer is sized as ne20*ne21 + ne02 (entries + padding). 883/883 test-backend-ops MUL_MAT_ID tests pass, llama-bench is within run-to-run variance of the previous behavior. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * metal : make mul_mat_id map0 ne02 a function constant use FOR_UNROLL for the id loops of kernel_mul_mm_id_map0 and turn ne02 into the function constant FC_mul_mm_id_map0_ne02 (FC_MUL_MM + 7), so that the threadgroup offsets of cnt/bas/sids and the prefix-sum bound fold at compile time. the pipeline name now encodes ne02 (same convention as the ssm_conv nc/silu constants) and the map0 pipeline is compiled with the constant value set. ne02 is no longer part of ggml_metal_kargs_mul_mm_id_map0 - the dispatch grid still comes from the tensor ne[2] on the host. no measurable perf change: 883/883 test-backend-ops MUL_MAT_ID tests pass, op-level and llama-bench results are within run variance. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : clean-up + TODO * metal : dispatch mul_mat_id entry tiles from the prepass follow the approach of the vulkan implementation (ggml-org#29998): instead of launching a grid over experts x ne21 entry tiles and letting most of the threadgroups early exit, kernel_mul_mm_id_map0 now also emits a row map - the total tile count plus one (expert, first entry, valid rows) descriptor per 32-entry tile that has work. kernel_mul_mm_id is launched with a static upper bound on the number of tiles (full tiles over all entries plus at most one partial tile per expert), reads its descriptor and processes a single tile - the tail early exits. This removes the round loop that was needed to cover duplicate ids past the grid (register pressure: th_max 768 -> back to the plain mul_mm level) and duplicates are now handled implicitly, since they only produce more descriptors within the bound. The grid shrinks by ~8-16x on typical shapes (e.g. 128 experts, 8 used, 512 tokens: 128 tiles instead of 2048 entry tiles, of which 94% only launched to exit) and ne21 is no longer needed by the mm kernel kargs. 883/883 test-backend-ops MUL_MAT_ID tests pass. llama-bench (Qwen3.6-35B-A3B Q8_0) is now faster than master on all pp sizes (pp2048/4096/8192 +2%, pp512 +0.5%), tg32 unchanged. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : add smem assert
auRose94
pushed a commit
to auRose94/llama.cpp
that referenced
this pull request
Oct 28, 2026
* metal : fix mul_mat_id limits and kernel name * metal : fix mul_mat_id with duplicate ids kernel_mul_mm_id_map0 assumed at most one match per token and a fixed capacity of ne21 entries per expert, which produced wrong or dropped output rows and could overflow into the slice of the next expert when the ids contain duplicates. build a compact per-expert list of the (token, slot) entries instead: - pass 1 counts the entries of every expert, thread 0 computes the offsets of the expert slices (tpe now holds [base, count]) - pass 2 writes the entries with one unconditional store per token (conditional stores of divergent lanes are very slow), the dummy value of the tokens without a match lands on a padding slot of the slice that is never read - duplicates are handled by a rare branch - kernel_mul_mm_id reads the base/count of its expert and loops over the entry tiles that do not fit into the grid The ids buffer is sized as ne20*ne21 + ne02 (entries + padding). 883/883 test-backend-ops MUL_MAT_ID tests pass, llama-bench is within run-to-run variance of the previous behavior. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * metal : make mul_mat_id map0 ne02 a function constant use FOR_UNROLL for the id loops of kernel_mul_mm_id_map0 and turn ne02 into the function constant FC_mul_mm_id_map0_ne02 (FC_MUL_MM + 7), so that the threadgroup offsets of cnt/bas/sids and the prefix-sum bound fold at compile time. the pipeline name now encodes ne02 (same convention as the ssm_conv nc/silu constants) and the map0 pipeline is compiled with the constant value set. ne02 is no longer part of ggml_metal_kargs_mul_mm_id_map0 - the dispatch grid still comes from the tensor ne[2] on the host. no measurable perf change: 883/883 test-backend-ops MUL_MAT_ID tests pass, op-level and llama-bench results are within run variance. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : clean-up + TODO * metal : dispatch mul_mat_id entry tiles from the prepass follow the approach of the vulkan implementation (ggml-org#29998): instead of launching a grid over experts x ne21 entry tiles and letting most of the threadgroups early exit, kernel_mul_mm_id_map0 now also emits a row map - the total tile count plus one (expert, first entry, valid rows) descriptor per 32-entry tile that has work. kernel_mul_mm_id is launched with a static upper bound on the number of tiles (full tiles over all entries plus at most one partial tile per expert), reads its descriptor and processes a single tile - the tail early exits. This removes the round loop that was needed to cover duplicate ids past the grid (register pressure: th_max 768 -> back to the plain mul_mm level) and duplicates are now handled implicitly, since they only produce more descriptors within the bound. The grid shrinks by ~8-16x on typical shapes (e.g. 128 experts, 8 used, 512 tokens: 128 tiles instead of 2048 entry tiles, of which 94% only launched to exit) and ne21 is no longer needed by the mm kernel kargs. 883/883 test-backend-ops MUL_MAT_ID tests pass. llama-bench (Qwen3.6-35B-A3B Q8_0) is now faster than master on all pp sizes (pp2048/4096/8192 +2%, pp512 +0.5%), tg32 unchanged. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : add smem assert
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
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.
Overview
MUL_MAT_IDfor duplicate ids (the new tests from vulkan: handle mul_mat_id duplicates in the prepass rather than looping #29998): the oldkernel_mul_mm_id_map0assumed at most one match per token and a fixed capacity ofne21entries per expert, which produced wrong or dropped output rows and could overflow into the slice of the next expert.kernel_mul_mm_id_map0) builds a compact per-expert list of(token, slot)entries: pass 1 counts the entries of every expert, thread 0 computes the offsets of the expert slices and emits a row map - the total tile count plus one(expert, first entry, valid rows)descriptor per 32-entry tile that has work - and pass 2 fills the slices with one unconditional store per token (stores from divergent lanes are very slow; the dummy value of the tokens without a match lands on a padding slot that is never read), with duplicates handled by a rare branch.kernel_mul_mm_idis launched with a static upper bound on the tiles (min(n_rows, ceil(n_rows/32) + min(n_experts, n_rows) - 1)), reads its descriptor and processes a single tile - the tail early exits. This replaces the previous experts x ne21 grid (mostly threadgroups that launched only to exit, ~94% on typical shapes) and removes the need for any per-expert loop, so duplicate ids are covered implicitly and the kernel's register pressure is back to the plainmul_mat_idlevel.ne20*ne21 + ne02(entries + one padding slot per expert);ne02is passed tokernel_mul_mm_id_map0as a function constant and its id loops useFOR_UNROLL.ne02 <= 1024guard inggml_metal_op_mul_mat_id_use_mm,_ne20pipeline name suffix).Requirements