refactor(norm): migrate normalization family into kernels/norm/ - #16
Conversation
96e50ae to
2285ea6
Compare
|
Increment over #15 is clean (~100 lines, mostly moves + doc sync). The fused One small thing: |
|
@ekryski clean, all 13 norm files moved + re-exported, block-scaled qgemv correctly left |
Moves the normalization kernels out of the legacy mlx/ + ffai/ split into
kernels/norm/ and drops the ffai_ / bare-name prefixes per the consolidation
plan (names describe the operation, not the source/model):
- mlx/rms_norm, mlx/layer_norm -> kernels/norm/ (mt_rms_norm*, mt_layer_norm)
- ffai/{rms_norm_residual,rms_norm_rope,gated_rmsnorm,rms_norm_qgemv,
gated_rms_norm_qgemv} -> mt_* (drop ffai_)
- ffai/adain1d: adain1d -> mt_adain1d
- the *_block_scaled_qgemv format matrices move as-is (already mt_<format>_*);
the format-axis fold (plan §7) is deferred to a later pass.
Generated MSL per kernel is unchanged; only paths + inventory names change.
Updated the two integration-test imports, cross-ref doc comments, the plan
(norm marked done) and KERNEL_AUDIT rows.
cargo build + clippy clean; all norm kernel tests + the rms_norm integration
tests pass.
The looser f32 tolerance is intentional: the one-pass E[x^2]-E[x]^2 variance over a length-300 reduce_sum is cancellation-prone and the GPU reduction order differs from the oracle's sequential sum. Comment-only; prevents the next person from tightening it and chasing a phantom fail.
c084b8b to
25b0608
Compare
What
Wave-1 kernel-family port: moves the normalization family out of the legacy
mlx/+ffai/split intokernels/norm/perdocs/specs/KERNEL_CONSOLIDATION_PLAN.md, and drops theffai_/ bare-nameprefixes so kernels are named for the operation, not the source.
Changes
mlx/{rms_norm,layer_norm}→kernels/norm/(mt_rms_norm*,mt_layer_norm).ffai/{rms_norm_residual,rms_norm_rope,gated_rmsnorm,rms_norm_qgemv,gated_rms_norm_qgemv}→
kernels/norm/with theffai_prefix dropped (mt_*).ffai/adain1d:adain1d→mt_adain1d(it had nomt_prefix at all).*_block_scaled_qgemvformat matrices move as-is — they already useclean
mt_<format>_*names; the format-axis fold (plan §7) is a later pass.batched_qkv_qgemv/rms_norm_block_scaled_qgemvcross-ref doc comments, the consolidation plan(norm marked done), and the
KERNEL_AUDITnorm rows.Generated MSL per kernel is unchanged — only module paths and inventory names
change. The FFAI emit consumer is regenerated from the new inventory.
Verification
cargo build -p metaltile-std+cargo clippy -p metaltile-std --testsclean.tile test(rms_norm, layer_norm,adain1d, gated, residual, rope, and the block-scaled format matrices).
rms_norm_per_head_gpuanddispatch_pinning_guardintegration tests pass.Note on stacking
This branch is based on the conv+rope reorg branch (#15), which in turn is based
on the docs branch (#14), so the diff here includes those until they merge. Each
family lands as its own PR; they consolidate cleanly once the earlier ones merge.