refactor(moe): migrate the moe family into kernels/moe/ - #25
Merged
Conversation
Commit message hygiene checkAll commit messages and PR text are clean. ✅ |
ekryski
force-pushed
the
ek/kernel-reorg-moe
branch
3 times, most recently
from
June 14, 2026 05:08
5eaa8f6 to
559a85c
Compare
Contributor
|
@ekryski clean move + legit int4/int8 mpp fold (6→3 files via |
TheTom
force-pushed
the
ek/kernel-reorg-gemv-q8-split
branch
from
June 22, 2026 21:27
3255164 to
b86ce1e
Compare
Move all 35 moe-family files from ffai/ + mlx/ into kernels/moe/ (which #24 seeded with gather_q4 + sigmoid_bias): - orchestration (ex moe.rs): router_topk + permute/unpermute + 10 gather_qmm - routers: router_topk_biased (ex dsv4_router_topk), router_sigmoid_bias, router_sqrtsoftplus, sigmoid_bias - mpp grouped BGEMM: mpp(+int8/bm8/bm64/×int8/×block_scaled) + mpp_shared - gguf-format expert matmul: bgemm_{q2k,iq2xxs,q4}_*, gemv_{rows,ws}_*, gather_* - down combine: down_swiglu_accum, down_weighted_sum_f16 - expert-indexed + block-scaled: dequant_gemv_expert_indexed(_block_scaled), block_scaled (ex mlx/block_scaled_moe) Filenames drop the redundant moe_ prefix (folder provides it); kernel names keep mt_moe_*. Model-name purge: mt_dsv4_router_topk -> mt_moe_router_topk_biased (distinct from the generic mt_moe_router_topk: selects by the biased score, weights by the unbiased). Bare dequant_gemv_int4_expert_indexed -> mt_ prefix. Fixed mpp_shared intra-imports and ~27 consumer test files (grouped/mixed use-blocks included). Format-axis fold (§7) deferred. orchestration.rs (~4k lines) moves whole here; split follows next.
…ther_qmm
The former moe.rs (~4.3k lines, 13 kernels + a shared test/bench module) was
moved whole in the previous commit; split it now per the <1000-line guideline:
- router_topk.rs — mt_moe_router_topk (top-k expert selection)
- permute.rs — mt_moe_permute / mt_moe_unpermute
- gather_qmm.rs — the 10 grouped-gather quantized-matmul cells (int4/int8/
b{3,5,6,8}, m8/m16/m32, mma, bm16) + their CSR/MMA helpers
Each file carries its own kernel_tests + kernel_benches (the shared u32_bytes
helper is duplicated into router_topk/permute). Consumer test imports remapped
orchestration:: -> the new modules. Kernel set unchanged (1272 codegen).
Revert the prefix-drop from the moe migration: filenames carry the moe_ prefix again so they match their mt_moe_* kernel names (moe_mpp.rs, moe_gather_qmm.rs, moe_router_topk.rs, moe_gather_q4.rs, …). The two files whose kernels are format-prefixed instead keep names matching those (block_scaled_moe.rs -> mt_<fmt>_gather_qmm; dequant_gemv_expert_indexed*). Updated mod.rs + all consumer-test imports (grouped + mixed use-blocks).
…3 files)
The six integer MPP grouped-BGEMM files were three {int4,int8} pairs that share
an identical matmul path (same coop_tile descriptor / staging / write-back) and
differ only in the weight unpack. Fold each pair onto a compile-time BITS axis:
moe_mpp.rs = variants(BITS=[4,8]) bm16 (cooperative-tensor path)
moe_mpp_bm8.rs = variants(BITS=[4,8]) bm8 (direct-input path)
moe_mpp_bm64.rs = variants(BITS=[4,8]) bm64 (4-simdgroup 2x2 path)
The unpack parameterizes cleanly with no branch: vals_per_pack = 32/BITS codes
per u32, decoded by (packed >> j*BITS) & ((1<<BITS)-1). Deletes the three _int8
files; same six kernel names, same 1272-kernel codegen. The BM axis is NOT
foldable (bm8/bm16/bm64 use three different MPP descriptor paths).
Verified all 6 mpp correctness suites pass (25 tests); clippy -D warnings clean.
The kernel was renamed mt_dsv4_router_topk -> mt_moe_router_topk_biased in this PR; the test file, test-fn name, and a stale ffai:: doc path still carried the old name. Rename-only, no assertion changes.
TheTom
force-pushed
the
ek/kernel-reorg-moe
branch
from
June 22, 2026 21:28
559a85c to
392b5a0
Compare
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.
Stacked on #24. Migrates the moe family — the bulk of wave 3 — from
ffai/+mlx/intokernels/moe/(which #24 seeded withgather_q4+sigmoid_bias). Two commits, as agreed: move whole, then split the monster file.Commit 1 — migrate (35 files)
orchestration(exmoe.rs),router_topk_biased(exdsv4_router_topk),router_sigmoid_bias,router_sqrtsoftplusmpp(+_int8/_bm8/_bm64/×int8/×block_scaled) +mpp_sharedbgemm_{q2k,iq2xxs,q4}_*,gemv_{rows,ws,rows_view}_*,gather_{down_q2k,gemv_iq2xxs}down_swiglu_accum,down_weighted_sum_f16dequant_gemv_expert_indexed(_block_scaled),block_scaled(exmlx/block_scaled_moe)moe_prefix (folder provides it; matches refactor(gemm): split the gemv_q8 grab-bag by family #24); kernel names keepmt_moe_*.mt_dsv4_router_topk→mt_moe_router_topk_biased— it's a distinct kernel from the genericmt_moe_router_topk(selects top-k by the biased score, weights by the unbiased). Baredequant_gemv_int4_expert_indexed→mt_prefixed.mpp_sharedintra-imports and ~27 consumer test files (grouped + single-line-mixeduseblocks included).*_block_scaledmatrices move as-is.Commit 2 — split
orchestration.rsThe former
moe.rs(~4.3k lines, 13 kernels) moved whole in commit 1, then split per the <1000-line guideline:router_topk.rs—mt_moe_router_topkpermute.rs—mt_moe_permute/mt_moe_unpermutegather_qmm.rs— the 10 grouped-gather quantized-matmul cells + their CSR/MMA test/bench helpersEach file carries its own
kernel_tests+kernel_benches; consumer imports remappedorchestration::→ the new modules.Verification
cargo build --workspace --bins✓cargo clippy --workspace --all-targets --all-features -- -D warnings✓ (clean)tile build— 1272 kernels codegen to MSL (count unchanged across both commits); every renamed/relocated/split kerneltile inspect-able