rms_norm: run f16 rows natively instead of widening the whole tensor

[czoli1976]

RmsNorm's f16 fast path converted the whole tensor to f32, ran the f32 kernel, and converted back — three times the tensor's size through memory and two allocations, for an op that is otherwise memory-bandwidth bound. This adds rms_norm_f16 to the linalg dispatch table with an aarch64 kernel, and normalises the f16 tensor in place.

The kernel

rms_norm_f16_inner is the existing rms_norm_f32_inner with the loads and stores changed and nothing else. FCVTL/FCVTL2 widen each 16-element group into the same v4-v7 the f32 kernel accumulates from, so the FMLA chain, the horizontal reduce, the scalar mean/eps/rsqrt and the scalar tail are untouched. On the way out FCVTN/FCVTN2 round to nearest-even under the default FPCR, which is what f16::from_f32 does.

Since f16 -> f32 is exact, that makes the kernel bit-identical to widening the row in the caller, which is what the path did before. f16_matches_widening_the_row_to_f32 asserts exactly that at ten lengths spanning the 16-element body and every tail shape.

Not regressing hosts without an f16 kernel

Worth flagging, because the naive version of this change would quietly hurt x86: with only an aarch64 kernel registered, x86 would drop from "widen + AVX-512 f32 kernel" to a generic scalar f16 loop. So the generic rms_norm_f16 widens one row at a time into a reused thread-local buffer and calls whichever rms_norm_f32 the host registered — AVX-512 keeps its arithmetic and its speed, and still loses the whole-tensor allocation. Bit-identical there too.

Numbers

f16 RmsNorm on the trailing axis, through a real plan, criterion against a saved baseline on the merge-base:

shape	before	after
128x768	79.0 us	19.3 us	-75.5%
512x4096	1.763 ms	424 us	-76.1%
2048x4096	6.672 ms	1.615 ms	-75.8%

p = 0.00 throughout.

What I tried first

Doing this core-side — keeping the f32 kernel and converting a row at a time into a scratch — is a regression (+7.7% at 128x768, +6.5% at 2048x4096). Row-local scratch buys locality, but the conversion loops are then scalar to_f32() per element, and that costs more than the bulk cast_to it replaces. The conversion has to happen in-register for this to pay, which is why the kernel is asm rather than core Rust.

Tests

tract-linalg 4415, tract-core 270, test-f16 2378, test-unit-core 816 — green. cargo fmt --all and cargo clippy clean.

🍍