Vulkan: add hoisting support for row IDs and expert count in shaders - #26686
Vulkan: add hoisting support for row IDs and expert count in shaders#26686ravel7524 wants to merge 3 commits into
Conversation
|
FreeBSD 16, Mesa 26.2.0, 7900 xtx (RDNA3 gfx1100) Qwen3.6-35B-A3B (MoE)
Qwen3.8-27B (dense)
@ravel7524 your patch is working properly, it gives 100 tok/s for MoE... nice ! |
|
Please apply it to the coopmat2 path as well. I expect the gains will be small due tot he larger tile size, but probably still worth a bit and good for consistency. |
|
@jeffbolznv Thanks I agree it belongs in the coopmat2 path too. I have no hardware with coopmat2 support, so I can't validate that change or even smoke test it. Please feel free to build on this branch directly. I'd like to take this out of draft first and I'm waiting on my other open PRs to close before adding to the queue, so don't wait for me if you want to move sooner. One more thing: unless I'm misreading, this is currently a small regression on coopmat2 rather than just a missed optimisation. |
|
I had codex port to coopmat2 and I tested it, please cherry-pick this commit: fa24cd3. |
|
Cherry-picked as fa24cd3, thanks for writing and testing it. Reviewed the diff, the packing matches the existing cm2 load_row_ids, and the push constants were already plumbed through, so nothing else needed changing. This closes the coopmat2 regression I flagged above. No blocking questions from me. |
|
My results are noisier than usual today, but the trend is "it's a few percent faster": |
|
The latest results look very promising. Is there anything else you need from me to help move this PR toward merge? |
|
Just address the review feedback, thanks. |
- use vk_op_count_experts_push_constants instead of a raw uint vector - apply the fastdiv trick to the ne00 div/mod in count_experts - compute the per-expert offsets with subgroupExclusiveAdd when the device supports it, keeping the serial path as fallback - document the data_d layout and the hoisted_row_id_words bound - drop a leftover debug print in ggml_vk_matmul_id
github-actions
Bot
added
Vulkan
|
Addressed all six. Two things worth flagging rather than leaving in the diff: fastdiv is now a third local copy of the same helper; generic_unary_head.glsl and glu_head.glsl already carry it. Happy to hoist it into utils.glsl if you'd prefer that over another copy. The subgroup path is a count_experts_subgroup variant gated on subgroup_arithmetic && subgroup_require_full_support, with the serial version kept as fallback. Trip count is padded to a multiple of gl_SubgroupSize so the subgroup ops stay in uniform control flow. |
| (uint32_t)n_as, | ||
| uint32_t(hoist_row_ids), | ||
| 0, 0 }; | ||
| init_fastdiv_values(pc.ne00, pc.ne00mp, pc.ne00L); |
There was a problem hiding this comment.
I'd like this to be in a init_pushconst_fastdiv template specialization, to be consistent with other push constants.
|
I just had one nitpick remaining. I think these optimizations helped perf another 1% or so, but my numbers got noisier so I don't want to post them. |
|
I wrote the fastdiv cleanup here: #27526. Whichever change goes second will need to adapt, should be trivial. |
3 participants
Overview
This PR optimizes Vulkan
MUL_MAT_IDfor routed Mixture-of-Experts (MoE) prompt processing.Each expert’s assigned token-row IDs are collected once and reused by the related GPU workgroups. Previously, every workgroup independently searched the routing table for the same rows.
This removes redundant routing work around the GEMM. Dense-model paths and token-generation/decode paths are unchanged.
Additional information
Hardware/software:
gfx1100/ RDNA3)26.0.3-1ubuntu1; Vulkan device API1.4.335A-B-A-B-A-B, with-r 5per phaseGGML_VK_MUL_MAT_ID_HOIST_ROW_IDS=1Qwen3.6 35B-A3B Q4_K_M:
Gemma 4 26B A4B IT Q8_0
Qwen3.6 35B-A3B UD Q6_K
Requirements