Skip to content

refactor(norm): migrate normalization family into kernels/norm/ - #16

Merged
TheTom merged 2 commits into
devfrom
ek/kernel-reorg-norm
Jun 22, 2026
Merged

refactor(norm): migrate normalization family into kernels/norm/#16
TheTom merged 2 commits into
devfrom
ek/kernel-reorg-norm

Conversation

@ekryski

@ekryski ekryski commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

What

Wave-1 kernel-family port: moves the normalization family out of the legacy
mlx/ + ffai/ split into kernels/norm/ per
docs/specs/KERNEL_CONSOLIDATION_PLAN.md, and drops the ffai_ / bare-name
prefixes 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 the ffai_ prefix dropped (mt_*).
  • ffai/adain1d: adain1dmt_adain1d (it had no mt_ prefix at all).
  • The two *_block_scaled_qgemv format matrices move as-is — they already use
    clean mt_<format>_* names; the format-axis fold (plan §7) is a later pass.
  • Updated the two integration-test imports, the batched_qkv_qgemv /
    rms_norm_block_scaled_qgemv cross-ref doc comments, the consolidation plan
    (norm marked done), and the KERNEL_AUDIT norm 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 --tests clean.
  • All norm-family kernel tests pass via tile test (rms_norm, layer_norm,
    adain1d, gated, residual, rope, and the block-scaled format matrices).
  • rms_norm_per_head_gpu and dispatch_pinning_guard integration 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.

@TheTom

TheTom commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

Increment over #15 is clean (~100 lines, mostly moves + doc sync). The fused mt_mxfp4/nvfp4/mxfp8_e4m3_gated_rms_norm_qgemv family keeps one kernel per format for now, which is the right interim call: folding the format onto a variants(FMT = ...) axis is wave-3 quant work per the consolidation plan, and doing it here would bloat a move-PR with body rewrites.

One small thing: mt_adain1d ships with f32 tol 1e-3 where the house default for f32 is 1e-4. If the looser tolerance is needed (long reduction or transcendental in the body), a one-line comment on the #[test_kernel] saying why prevents the next person from tightening it and chasing a phantom failure.

@TheTom

TheTom commented Jun 16, 2026

Copy link
Copy Markdown
Contributor

@ekryski clean, all 13 norm files moved + re-exported, block-scaled qgemv correctly left mt_-named, non-move edits are doc/import only. Approve.

ekryski and others added 2 commits June 22, 2026 16:22
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.
@TheTom
TheTom force-pushed the ek/kernel-reorg-norm branch from c084b8b to 25b0608 Compare June 22, 2026 21:22
@TheTom
TheTom merged commit 59ca3c7 into dev Jun 22, 2026
8 checks passed
@ekryski
ekryski deleted the ek/kernel-reorg-norm 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