refactor: migrate the final ffai/ families — sdpa, quant, stragglers - #27
Merged
Conversation
Commit message hygiene checkAll commit messages and PR text are clean. ✅ |
ekryski
force-pushed
the
ek/kernel-reorg-sdpa-quant
branch
from
June 14, 2026 04:51
99f3a5d to
b3f6f2c
Compare
ekryski
force-pushed
the
ek/kernel-format-axis-fold
branch
from
June 14, 2026 04:59
9a849fd to
3c4bd4d
Compare
ekryski
force-pushed
the
ek/kernel-reorg-sdpa-quant
branch
from
June 14, 2026 05:02
b3f6f2c to
c521ca7
Compare
ekryski
force-pushed
the
ek/kernel-format-axis-fold
branch
from
June 14, 2026 05:08
3c4bd4d to
29d92dd
Compare
ekryski
force-pushed
the
ek/kernel-reorg-sdpa-quant
branch
from
June 14, 2026 05:08
c521ca7 to
0be5a16
Compare
Contributor
TheTom
force-pushed
the
ek/kernel-format-axis-fold
branch
from
June 22, 2026 21:29
29d92dd to
0f871c8
Compare
…s family Move the leftover activation kernels into kernels/ops/ (leaky_relu, gelu_erf, gate_up_swiglu_fused, + dsv4_swiglu_limit -> swiglu_limit) and the manifold- constrained hyper-connection kernels into a new kernels/hyper_connections/ family (mhc, mhc_sinkhorn_split). Drop ffai_/dsv4_ prefixes -> mt_ (dsv4 is a model name; mhc/swiglu_limit are the operations). Verified 1272 kernels codegen; swiglu_limit (3) + mhc (18) tests pass; clippy clean.
Move all 26 attention kernels from mlx/ + ffai/ into a new kernels/sdpa/ family:
the bidirectional / decode / multi / prefill paths (head-dims d32..d512, relpos /
windowed / conformer / 2pass / batched / sink variants), the quantized +
block-scaled flash forms, the AURA compressed-domain attention (flash + score +
value), the DSv4 compressed-sparse-attention (CSA) decode + Lightning Indexer +
compressor, the head-gate, and the mlx scaled_dot_product_attention + sdpa_vector.
Naming: ffai_ -> mt_; bare aura_*/flash_quantized_sdpa/sdpa_decode_{2pass,batched}_*
-> mt_; dsv4_ model prefix dropped (csa/indexer/compressor are the operations).
Fixed consumer-test imports + the AURA MSL snapshots (renamed kernel).
Verified 1272 kernels codegen; sdpa in-file tests (160) + external suites pass;
clippy -D warnings clean.
… ffai/
Move src/quant/ (codec, format, gguf host SSOT) into kernels/quant/ and migrate
the dequant GPU kernels there too: aura_{dequant_rotated,encode}, dequant_gather
(+block_scaled), gguf_dequant_* (+iq2xxs extract), dsv4_{fp8_block,mxfp4}_dequant
-> {fp8_block,mxfp4}_dequant, ffai_dequant_q4 -> dequant_q4, mlx block_scaled_dequant.
Repointed all 228 crate::quant:: refs (incl the relative quant:: inside the §7
test setups' use crate::{} blocks) -> crate::kernels::quant::. Dropped ffai_/dsv4_
prefixes -> mt_. Pruned the dead mlx conv.rs/shared.rs stubs and the now-empty
ffai/ module. Updated consumer-test imports + AURA MSL snapshots.
ffai/ is fully migrated (folder removed); only mlx/steel/ (steel attn/conv
templates) remains. Verified: 1272 kernels codegen; quant tests (gguf/dequant/
aura/...) pass; clippy -D warnings clean.
TheTom
force-pushed
the
ek/kernel-reorg-sdpa-quant
branch
from
June 22, 2026 21:29
0be5a16 to
e043987
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
2 participants
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.
The final consolidation: empties
ffai/entirely and folds the last families intokernels/.What moved
kernels/sdpa/(new, 26 files) — all attention: bidirectional / decode / multi / prefill (head-dims d32–d512, relpos/windowed/conformer/2pass/batched/sink), the quantized + block-scaled flash forms, AURA compressed-domain attention (flash/score/value), DSv4 compressed-sparse-attention (CSA) decode + Lightning Indexer + compressor, head-gate, and the mlxscaled_dot_product_attention+sdpa_vector.kernels/quant/— unifies the host SSOT (codec/format/gguf, moved fromsrc/quant/) with the dequant GPU kernels (gguf_dequant_*, aura encode/dequant, dequant_gather, fp8_block/mxfp4 dequant, dequant_q4, block_scaled_dequant). All 228crate::quant::refs repointed.kernels/ops/— activation stragglers (leaky_relu, gelu_erf, gate_up_swiglu_fused, swiglu_limit).kernels/hyper_connections/(new) — manifold-constrained hyper-connections (mhc, mhc_sinkhorn_split).Naming
ffai_→mt_; bare kernels (aura_, flash_quantized_sdpa, sdpa_decode_{2pass,batched}_, dequant_gather variants) →mt_;dsv4_model prefix dropped (csa/indexer/compressor/mhc/swiglu_limit are the operations).aurakept (codec/technique, per plan §2).Cleanup
Pruned the dead
mlx/conv.rs+mlx/shared.rsstubs and the now-emptyffai/module.ffai/is gone; onlymlx/steel/(steel attn/conv templates) remains as a separate template-infra concern.Verification
cargo build --workspace,clippy --all-targets -D warningsclean, fulltile test= 1394 suites / 4125 tests / 0 failed, 1272 kernels codegen. Consumer-test imports + AURA MSL snapshots updated.Stacked on #26 (the §7 format-axis fold-out).