refactor(gemm): migrate the dense matmul family into kernels/gemm/ - #20
Merged
Conversation
ekryski
force-pushed
the
ek/kernel-reorg-gemm
branch
from
June 13, 2026 16:19
8c1ee8f to
a028af6
Compare
Contributor
|
@ekryski clean mass-move, signatures + MSL identical. Heads-up: |
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
force-pushed
the
ek/kernel-reorg-gemm
branch
from
June 22, 2026 21:24
a028af6 to
4557c60
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.
What
Wave-2 family port (part 1): migrates the dense matmul kernels into
kernels/gemm/perdocs/specs/KERNEL_CONSOLIDATION_PLAN.md. gemm is thelargest 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 asteel/subdir). Alreadymt_*-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.ffai/ssm.rstest ref and thesteel_msl_snapshotsimport; refreshed
ffai_gemmdoc-comment references; updated theKERNEL_AUDITpaths/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 afollow-up PR; the format-axis fold (plan §7) stays deferred, same as
norm/.Verification
cargo build+cargo clippy --workspace --all-targetsclean.tile test: gemm (12), gemv (591 incl. quantized stillin ffai), gemv_axpy (12), patch_embed (195), steel fused/gather/masked/segmented/splitk (60).
steel_msl_snapshots(7) and thessmcorrectness test pass.Note on stacking
Stacked on ops (#18) → sampling (#17) → norm (#16) → conv+rope (#15) →
dev; thediff includes those until they merge.