Skip to content

refactor(gemm): migrate the dense matmul family into kernels/gemm/ - #20

Merged
TheTom merged 1 commit into
devfrom
ek/kernel-reorg-gemm
Jun 22, 2026
Merged

refactor(gemm): migrate the dense matmul family into kernels/gemm/#20
TheTom merged 1 commit into
devfrom
ek/kernel-reorg-gemm

Conversation

@ekryski

@ekryski ekryski commented Jun 13, 2026

Copy link
Copy Markdown
Contributor

What

Wave-2 family port (part 1): migrates the dense matmul kernels into
kernels/gemm/ per docs/specs/KERNEL_CONSOLIDATION_PLAN.md. gemm is the
largest family and is deeply quant-entangled, so it's split: this PR moves the
non-quantized ops; the quantized op×format matrix follows.

Changes

  • mlx/{gemv,gemv_masked} + mlx/steel/gemm/*kernels/gemm/ (with a
    steel/ subdir). Already mt_*-named.
  • ffai/{gemm,patch_embed,patch_embed_mma,gemv_axpy_inplace}kernels/gemm/
    with mt_ prefixes: mt_gemm/mt_gemm_batched, mt_patch_embed/mt_patch_embed_mma,
    mt_gemv_axpy_inplace.
  • Fixed the cross-family ffai/ssm.rs test ref and the steel_msl_snapshots
    import; refreshed ffai_gemm doc-comment references; updated the KERNEL_AUDIT
    paths/names and marked gemm dense done in the plan.

Generated MSL per kernel is unchanged — only module paths + inventory names.

Deferred (follow-up)

The quantized matmuls — quantized_* / fp_quantized_* / block_scaled_qmm*,
gemm_q4/q8, gemv_q8, and the batched *_qgemv/*_qmm (+ block-scaled) forms —
are the quantized form of these same ops. They land in kernels/gemm/ in a
follow-up PR; the format-axis fold (plan §7) stays deferred, same as norm/.

Verification

  • cargo build + cargo clippy --workspace --all-targets clean.
  • Kernel tests pass via tile test: gemm (12), gemv (591 incl. quantized still
    in ffai), gemv_axpy (12), patch_embed (195), steel fused/gather/masked/segmented/splitk (60).
  • steel_msl_snapshots (7) and the ssm correctness test pass.

Note on stacking

Stacked on ops (#18) → sampling (#17) → norm (#16) → conv+rope (#15) → dev; the
diff includes those until they merge.

@TheTom

TheTom commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

@ekryski clean mass-move, signatures + MSL identical. Heads-up: patch_embed/patch_embed_mma were previously un-prefixed, so the registered name changes (mt_patch_embed*), easy to miss in a ffai_→mt_ sweep. The pub use … as conv back-compat shim in mlx/steel/mod.rs is a good follow-up delete target once mlx/ fully dissolves. Approve.

First (dense) half of the gemm family — the matmul ops that aren't quantized:

- mlx/{gemv,gemv_masked} + mlx/steel/gemm/* -> kernels/gemm/ (+ steel/ subdir);
  already mt_* named.
- ffai/{gemm,patch_embed,patch_embed_mma,gemv_axpy_inplace} -> kernels/gemm/
  with mt_ prefixes (mt_gemm(_batched) in dense.rs — named dense.rs to avoid
  module_inception; mt_patch_embed(_mma), mt_gemv_axpy_inplace).
- update the cross-family ssm test ref + steel_msl_snapshots import; refresh the
  ffai_gemm doc-comment references and the KERNEL_AUDIT paths/names; mark gemm
  dense done in the plan.

The quantized matmul matrix lands in a follow-up to the same folder; the
format-axis fold (§7) stays deferred.

cargo build + workspace clippy (-D warnings) clean; gemm/gemv/patch_embed/steel/
axpy kernel tests pass, steel_msl_snapshots + ssm correctness pass.
@TheTom
TheTom force-pushed the ek/kernel-reorg-gemm branch from a028af6 to 4557c60 Compare June 22, 2026 21:24
@TheTom
TheTom merged commit 88ff027 into dev Jun 22, 2026
8 checks passed
@ekryski
ekryski deleted the ek/kernel-reorg-gemm 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