refactor(kernels): migrate kv_cache + audio + vision families - #21
Merged
Conversation
Contributor
|
@ekryski clean. kv renames for downstream: |
- mlx/fft.rs + ffai/{kv_cache,kv_cache_update_many} -> kernels/kv_cache/
(kv_cache.rs -> cache.rs to avoid module_inception).
- mt_ prefixes on the bare-named kernels: mt_kv_cache_update(_many),
mt_quantize_kv(_int4/_int8/_fp8_e4m3/_fp8_e5m2), mt_bulk_dequant_kv(...).
fft already mt_*.
- update the kv_cache_quant_roundtrip integration test imports + the conv1d
doc-comment cross-ref.
cargo build clean; kv_cache + fft + quantize_kv tests pass, integration test passes.
ffai/{lstm,mel_spectrogram,snake1d,upsample_nearest1d,vocoder} -> kernels/audio/
with mt_ prefixes: mt_lstm, mt_mel_spectrogram(+_magnitude/_stft_window/_filterbank),
mt_snake1d, mt_upsample_nearest1d, mt_vocoder_istft.
cargo build + clippy clean; mel/lstm/snake/vocoder/upsample tests pass.
ffai/{avg_pool2d_nhwc,frame_diff_luma,im2col_patch,im2col_patch_interleaved,
patch_unfold_qwen,pos_emb_2d_add,resize_normalize,transpose_th} -> kernels/vision/
with mt_ prefixes. patch_unfold_qwen.rs -> patch_unfold.rs and
ffai_patch_unfold_qwen -> mt_patch_unfold (Qwen stays as a usage note in the
comment, not the name; test/bench fns de-modelled too).
cargo build + clippy clean; resize/im2col/patch_unfold/transpose/avg_pool/
frame_diff/pos_emb tests pass.
Doc-comment only; the kernels moved to kernels/kv_cache/ in this PR.
TheTom
force-pushed
the
ek/kernel-reorg-kv-audio-vision
branch
from
June 22, 2026 21:25
429cf8e to
0e99ee3
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 grouped port: three small, clean families migrated into
kernels/in onePR (you OK'd grouping when it makes sense). Each is its own commit.
mlx/fft.rs+ffai/{kv_cache,kv_cache_update_many}→kernels/kv_cache/(kv_cache.rs→cache.rsto avoidmodule_inception).mt_prefixes on the bare kernels (mt_kv_cache_update(_many),mt_quantize_kv*,mt_bulk_dequant_kv*); fft alreadymt_*.ffai/{lstm,mel_spectrogram,snake1d,upsample_nearest1d,vocoder}→kernels/audio/withmt_prefixes.ffai/{avg_pool2d_nhwc,frame_diff_luma,im2col_patch,im2col_patch_interleaved, patch_unfold_qwen,pos_emb_2d_add,resize_normalize,transpose_th}→kernels/vision/.patch_unfold_qwen→patch_unfold/mt_patch_unfold(Qwen stays a usage note in the comment; test/bench fns de-modelled too).
Plus the plan + KERNEL_AUDIT updated to mark these three done.
Generated MSL per kernel is unchanged — only module paths + inventory names.
Verification
cargo clippy --workspace --all-targets -- -D warningsclean.tile testgreen: kv_cache + quantize_kv + fft; mel/lstm/snake/vocoder/upsample;resize/im2col/patch_unfold/transpose/avg_pool/frame_diff/pos_emb. The
kv_cache_quant_roundtripintegration test passes.Note on stacking
Stacked on gemm-dense (#20) → ops (#18) → … →
dev. (This branch also carries afix that #20's force-push picked up:
kernels/gemm/gemm.rs→dense.rsto cleara
module_inceptionthat was failing #20's-D warningsclippy.)