refactor(gemm): migrate quantized matmuls into kernels/gemm/ - #23
Merged
Conversation
Contributor
|
@ekryski clean, 28-in/28-out reconciled. One non- |
TheTom
force-pushed
the
ek/kernel-reorg-ssm
branch
from
June 22, 2026 21:26
040d89e to
4bc10bd
Compare
Move the 28 quantized matmul files from mlx/ + ffai/ into kernels/gemm/ alongside the dense GEMM/GEMV set: integer quantized_* (incl mpp/nax/int8/ dynamic_m), fp_quantized_*, block_scaled_* formats, dequant_gemv, the Q8/Q4 gemm_q8(_mpp)/gemm_q4_mpp, the batched QKV/4-way qgemv/qmm projections, and the block-scaled patch-embed. Drop ffai_ prefixes -> mt_ (gemm_q8 family, batched, dequant_gemv); mlx kernels already mt_-prefixed. Fix intra-set path refs (crate::mlx -> crate::kernels::gemm) and consumer-side test imports. gemv_q8 (grab-bag spanning gemm/kv_cache/moe/norm/ops) stays parked until moe lands; block_scaled_dequant -> quant, block_scaled_moe/moe_* and dequant_gemv_expert_indexed* -> moe (migrated later). Format-axis fold (plan §7) deferred. Completes wave 2.
TheTom
force-pushed
the
ek/kernel-reorg-gemm-quantized
branch
from
June 22, 2026 21:27
2b24cef to
742dcea
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 #22 (ssm). Completes wave 2 of the kernel consolidation (
docs/specs/KERNEL_CONSOLIDATION_PLAN.md).What
Moves the 28 quantized-matmul files from the legacy
mlx/+ffai/split intokernels/gemm/, alongside the dense GEMM/GEMV set landed in #20:quantized(+_mpp/_mpp_int8/_nax/_nax_int8/_mma_dynamic_m),dequant_gemvfp_quantized(+_mma/_nax)block_scaled_{matmul,mma,qmm,qmm_mpp,qmm_nax}batched_{4,qkv}(_block_scaled)_{qgemv,qmm}gemm_q8,gemm_q8_mpp,gemm_q4_mpppatch_embed(_mma)_block_scaledNaming
ffai_prefixes →mt_(gemm_q8/grouped_gemm_q8(_mpp),gemm_q4_mpp,batched_*,dequant_gemv*). MLX kernels were alreadymt_-prefixed.crate::mlx::*→crate::kernels::gemm::*) and all consumer-side test imports.Scope / parked
gemv_q8(1069-line grab-bag spanning gemm/kv_cache/moe/norm/ops) stays parked until moe lands.block_scaled_dequant→ quant;block_scaled_moe/moe_*anddequant_gemv_expert_indexed*→ moe.Verification
cargo build --workspace --bins✓cargo clippy --workspace --all-targets --all-features -- -D warnings✓ (clean)tile test: qmm (966), batched (390), dequant_gemv (114), patch_embed (195), fp_qmm (6), qmv (21) — all passgemm_q8(_mpp)/grouped_gemm_q8/qmm_mma_dynamic_m/qmm_mpp(_int8)/dequant_gemv_indirect_wrappers/steel_msl_snapshots— all pass (33 tests)Docs updated: plan §2/§6 mark
gemm/done; KERNEL_AUDIT paths/names updated.