Ggml moe gathermatmul gpu#36932
Open
cavusmustafa wants to merge 3 commits into
Open
Conversation
The ggml-openvino frontend lowers the gemma-style top-k expert matmul
(MUL_MAT_ID) using public ops only, as a rank-2 carrier:
Reshape( [Convert]( Gather( CompressedWeightsBlock[n_expert, m*k],
ids[n_tokens, n_used], axis=0 ) ) )
-> MatMul(activations, ., transpose_b=true)
The frontend keeps the expert Gather on the rank-2 [n_expert, m*k] compressed
block on purpose: it folds to a plain GatherCompressed on any OpenVINO build
(weights stay compressed, no OOM), so the frontend is correct and memory-safe
with or without this pass.
When this pass is present (GPU systolic + oneDNN MoE pipeline) it does better:
it rewrites the carrier to the internal ov::op::internal::GatherMatmul so the
existing ConvertGatherMatmulToGatherMatmulCompressed folds it into a single
GatherMatmulCompressed (gather + dequantize-selected-experts + matmul in one op).
GatherMatmul needs a rank-3 [n_expert, N, K] weights block, so the pass rebuilds
the block's low-bit / scale / zero-point Constants at the grouped rank-4
[n_expert, N, K/group, group] shape (a bit-identical reinterpret of the same
data) and reshapes to rank-3.
Matching a genuine CompressedWeightsBlock as the Gather's source both guarantees
the downstream fold can fire and structurally excludes any non-compressed expert
matmul. Registered after ConvertTiledMoeBlockToGatherMatmuls and before
ConvertGatherMatmulToGatherMatmulCompressed.
Cover the ggml-openvino MoE expert-matmul fusion pass: - fires on the rank-2 compressed carrier for u4/u8 (asymmetric, with zero point) and symmetric i4 (no zero point), rewriting the expert MatMul to GatherMatmul; - end-to-end, the produced GatherMatmul folds to GatherMatmulCompressed and the low-bit weight Constant survives (weights stay compressed, no f32 materialization); - negative: does NOT fire on a non-compressed f16 expert Gather, nor when the MatMul is not transpose_b.
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.
Details:
MatcherPass(ov::intel_gpu::ConvertGgmlMoeGatherToGatherMatmul) that captures the ggml-openvino frontend's public-op MoE expert-matmul carrier —MatMul(activations, Reshape([Convert(]Gather(CompressedWeightsBlock[n_expert, m*k], ids)[)]), transpose_b=true)— and rewrites it to the internalov::op::internal::GatherMatmul.[n_expert, m*k], butGatherMatmulneeds rank-3[n_expert, N, K]; the pass rebuilds the block's low-bit weight / scale / (optional) zero-pointConstants at the grouped rank-4[n_expert, N, K/group, group]shape over the same data (a bit-identical reinterpret), reshapes to rank-3, and feeds it toGatherMatmul. The downstream stockConvertGatherMatmulToGatherMatmulCompressedthen folds it into a singleGatherMatmulCompressed(gather + dequantize-selected-experts + matmul), keeping the low-bit weights compressed.supports_immad && use_onednnMoE-compressed pipeline, afterConvertTiledMoeBlockToGatherMatmulsand beforeConvertGatherMatmulToGatherMatmulCompressed. Matching a genuineCompressedWeightsBlock(u4/i4/u8/i8) both guarantees the downstream fold can fire and structurally excludes any non-compressed expert matmul.GatherCompressed; with this pass it is upgraded toGatherMatmulCompressed. No env flag — the fusion is automatic when the plugin is present.transformations_pipeline.cpp(0 deletions to existing code).convert_ggml_moe_gather_to_gathermatmul_test.cpp, 6 cases): fires on u4/u8 (with zero point) and symmetric i4 (no zero point); end-to-end folds toGatherMatmulCompressedwith the u4 weightConstantsurviving (weights stay compressed); negative cases confirm it does not fire on a non-compressed f16 expertGathernor whentranspose_bis false.llama-bench, warmup on): prefill ~1.8x (7.7 -> 13.6 t/s), decode ~1.2x (3.96 -> 4.81 t/s) vs. the non-fused rank-2 path.Tickets:
AI Assistance:
ov_gpu_unit_teststarget; the 6 new unit tests all pass; ran the GPU transformation test subset (*Transform*:*Convert*:*Fuse*:*Compress*:*MoE*:*GatherMatmul*:*FullyConnected*:*MatMul*-> 260 passed, 1 pre-existing skip, 0 failed) to confirm no regression; verified end-to-end on the real 26B model (dumped IR shows all 60 expert matmuls rewritten and folded toGatherMatmulCompressed, a before/after constant census confirms no expert weights are materialized), a live 26B GPU run completes (EXIT 0), and the throughput numbers above were measured directly.cc: @wine99