Skip to content

refactor(moe): migrate the moe family into kernels/moe/ - #25

Merged
TheTom merged 7 commits into
devfrom
ek/kernel-reorg-moe
Jun 22, 2026
Merged

refactor(moe): migrate the moe family into kernels/moe/#25
TheTom merged 7 commits into
devfrom
ek/kernel-reorg-moe

Conversation

@ekryski

@ekryski ekryski commented Jun 13, 2026

Copy link
Copy Markdown
Contributor

Stacked on #24. Migrates the moe family — the bulk of wave 3 — from ffai/ + mlx/ into kernels/moe/ (which #24 seeded with gather_q4 + sigmoid_bias). Two commits, as agreed: move whole, then split the monster file.

Commit 1 — migrate (35 files)

Group Modules
Routing orchestration (ex moe.rs), router_topk_biased (ex dsv4_router_topk), router_sigmoid_bias, router_sqrtsoftplus
MPP grouped BGEMM mpp(+_int8/_bm8/_bm64/×int8/×block_scaled) + mpp_shared
GGUF-format expert matmul bgemm_{q2k,iq2xxs,q4}_*, gemv_{rows,ws,rows_view}_*, gather_{down_q2k,gemv_iq2xxs}
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; matches refactor(gemm): split the gemv_q8 grab-bag by family #24); kernel names keep mt_moe_*.
  • Model-name purge: mt_dsv4_router_topkmt_moe_router_topk_biased — it's a distinct kernel from the generic mt_moe_router_topk (selects top-k by the biased score, weights by the unbiased). Bare dequant_gemv_int4_expert_indexedmt_ prefixed.
  • Fixed the mpp_shared intra-imports and ~27 consumer test files (grouped + single-line-mixed use blocks included).
  • Format-axis fold (§7) deferred — the 13-kernel *_block_scaled matrices move as-is.

Commit 2 — split orchestration.rs

The former moe.rs (~4.3k lines, 13 kernels) moved whole in commit 1, then split per the <1000-line guideline:

  • router_topk.rsmt_moe_router_topk
  • permute.rsmt_moe_permute / mt_moe_unpermute
  • gather_qmm.rs — the 10 grouped-gather quantized-matmul cells + their CSR/MMA test/bench helpers

Each file carries its own kernel_tests + kernel_benches; consumer imports remapped orchestration:: → the new modules.

Verification

  • cargo build --workspace --bins
  • cargo clippy --workspace --all-targets --all-features -- -D warnings ✓ (clean)
  • tile build1272 kernels codegen to MSL (count unchanged across both commits); every renamed/relocated/split kernel tile inspect-able
  • 27 moe integration suites pass (59 tests, 0 failed); in-file kernel tests green (gather_qmm 315, router_topk 6, permute 6, expert_indexed 93)

@github-actions github-actions Bot added the ignore-for-release Hidden from release notes label Jun 13, 2026
@github-actions

github-actions Bot commented Jun 13, 2026

Copy link
Copy Markdown

Commit message hygiene check

All commit messages and PR text are clean. ✅

@TheTom

TheTom commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

@ekryski clean move + legit int4/int8 mpp fold (6→3 files via variants, expanded names byte-for-byte). mt_dsv4_router_topkmt_moe_router_topk_biased is fully propagated (stale test filename only). Approve.

@TheTom
TheTom force-pushed the ek/kernel-reorg-gemv-q8-split branch from 3255164 to b86ce1e Compare June 22, 2026 21:27
ekryski and others added 7 commits June 22, 2026 16:27
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
TheTom force-pushed the ek/kernel-reorg-moe branch from 559a85c to 392b5a0 Compare June 22, 2026 21:28
@TheTom
TheTom changed the base branch from ek/kernel-reorg-gemv-q8-split to dev June 22, 2026 21:29
@TheTom
TheTom merged commit c46c932 into dev Jun 22, 2026
2 checks passed
@ekryski
ekryski deleted the ek/kernel-reorg-moe branch July 13, 2026 20:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ignore-for-release Hidden from release notes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants