Skip to content

refactor(kernels): migrate kv_cache + audio + vision families - #21

Merged
TheTom merged 5 commits into
devfrom
ek/kernel-reorg-kv-audio-vision
Jun 22, 2026
Merged

refactor(kernels): migrate kv_cache + audio + vision families#21
TheTom merged 5 commits into
devfrom
ek/kernel-reorg-kv-audio-vision

Conversation

@ekryski

@ekryski ekryski commented Jun 13, 2026

Copy link
Copy Markdown
Contributor

What

Wave-2 grouped port: three small, clean families migrated into kernels/ in one
PR (you OK'd grouping when it makes sense). Each is its own commit.

  • kv_cachemlx/fft.rs + ffai/{kv_cache,kv_cache_update_many}
    kernels/kv_cache/ (kv_cache.rscache.rs to avoid module_inception).
    mt_ prefixes on the bare kernels (mt_kv_cache_update(_many),
    mt_quantize_kv*, mt_bulk_dequant_kv*); fft already mt_*.
  • audioffai/{lstm,mel_spectrogram,snake1d,upsample_nearest1d,vocoder}
    kernels/audio/ with mt_ prefixes.
  • visionffai/{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_qwenpatch_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 warnings clean.
  • tile test green: 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_roundtrip integration test passes.

Note on stacking

Stacked on gemm-dense (#20) → ops (#18) → … → dev. (This branch also carries a
fix that #20's force-push picked up: kernels/gemm/gemm.rsdense.rs to clear
a module_inception that was failing #20's -D warnings clippy.)

@TheTom

TheTom commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

@ekryski clean. kv renames for downstream: kv_cache_update[_many]mt_kv_cache_update[_many], quantize_kv/bulk_dequant_kvmt_*. Nit: tests/kv_cache_quant_roundtrip_gpu.rs:6 doc comment still says "ship in ffai::kv_cache". Approve.

ekryski and others added 5 commits June 22, 2026 16:25
- 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
TheTom force-pushed the ek/kernel-reorg-kv-audio-vision branch from 429cf8e to 0e99ee3 Compare June 22, 2026 21:25
@TheTom
TheTom merged commit 72f443e into dev Jun 22, 2026
8 checks passed
@ekryski
ekryski deleted the ek/kernel-reorg-kv-audio-vision 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