Skip to content

refactor(gemm): migrate quantized matmuls into kernels/gemm/ - #23

Merged
TheTom merged 2 commits into
devfrom
ek/kernel-reorg-gemm-quantized
Jun 22, 2026
Merged

refactor(gemm): migrate quantized matmuls into kernels/gemm/#23
TheTom merged 2 commits into
devfrom
ek/kernel-reorg-gemm-quantized

Conversation

@ekryski

@ekryski ekryski commented Jun 13, 2026

Copy link
Copy Markdown
Contributor

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 into kernels/gemm/, alongside the dense GEMM/GEMV set landed in #20:

  • Integer-quantizedquantized (+_mpp/_mpp_int8/_nax/_nax_int8/_mma_dynamic_m), dequant_gemv
  • FP-quantizedfp_quantized (+_mma/_nax)
  • Block-scaled / microscalingblock_scaled_{matmul,mma,qmm,qmm_mpp,qmm_nax}
  • Batched projectionsbatched_{4,qkv}(_block_scaled)_{qgemv,qmm}
  • Quantized GEMMgemm_q8, gemm_q8_mpp, gemm_q4_mpp
  • Quantized patch-embedpatch_embed(_mma)_block_scaled

Naming

  • Dropped ffai_ prefixes → mt_ (gemm_q8/grouped_gemm_q8(_mpp), gemm_q4_mpp, batched_*, dequant_gemv*). MLX kernels were already mt_-prefixed.
  • No model names in kernel names; files keep their format-distinguished names (the format-axis fold, plan §7, is deferred).
  • Fixed intra-set path refs (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.
  • Excluded (migrate with their home family): block_scaled_dequant → quant; block_scaled_moe/moe_* and dequant_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 pass
  • Touched integration tests: gemm_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.

@TheTom

TheTom commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

@ekryski clean, 28-in/28-out reconciled. One non-#[kernel] public rename to flag in the body: dequant_gemv_wants_indirectmt_dequant_gemv_wants_indirect (no alias), any runtime caller resolving indirect-dispatch by that symbol breaks. Approve.

@TheTom
TheTom force-pushed the ek/kernel-reorg-ssm branch from 040d89e to 4bc10bd Compare June 22, 2026 21:26
ekryski added 2 commits June 22, 2026 16:27
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
TheTom force-pushed the ek/kernel-reorg-gemm-quantized branch from 2b24cef to 742dcea Compare June 22, 2026 21:27
@TheTom
TheTom changed the base branch from ek/kernel-reorg-ssm to dev June 22, 2026 21:27
@TheTom
TheTom merged commit 05f038a into dev Jun 22, 2026
8 of 21 checks passed
@ekryski
ekryski deleted the ek/kernel-reorg-gemm-quantized 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