Add mosaic normalization kernel

[samanklesaria]

Adds a Mosaic version of the normalization kernel. Importantly, this implementation does not use any modern GPU features, and works with Ampere GPUs (e.g. A100) too!

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[samanklesaria]

Normalization benchmarks: Mosaic GPU vs Pallas-Triton

NVIDIA A100-SXM4-40GB, bf16, from _results/bench_output.log (2026-08-11).
Times in µs. Ratio = mosaic / pallas; <1 means mosaic is faster.

All lines logged an autotuning cache miss, so both backends ran heuristics
configs rather than tuned ones.

case	pallas	mosaic	xla	ratio
forward
384res_64chan	29.8	30.7	30.8	1.03
384res_128chan	59.7	58.3	59.7	0.98
768res_128chan	222	221	223	1.00
384res_128chan_axis0	67.5	63.3	144	0.94
768res_128chan_axis0	240	230	495	0.96
384res_64chan_vmap	221	220	227	1.00
384res_128chan_vmap	438	438	439	1.00
768res_128chan_vmap	1733	1739	1738	1.00
384res_128chan_axis0_vmap	473	455	1086	0.96
768res_128chan_axis0_vmap	1875	1815	4222	0.97
forward + vjp
384res_64chan	84.7	95.1	188	1.12
384res_128chan	175	167	367	0.95
768res_128chan	617	602	1348	0.98
384res_128chan_axis0	185	197	411	1.06
768res_128chan_axis0	671	674	1477	1.00
384res_64chan_vmap	623	624	1349	1.00
384res_128chan_vmap	1205	1181	2637	0.98
768res_128chan_vmap	4680	4635	10397	0.99
384res_128chan_axis0_vmap	1300	1356	3519	1.04
768res_128chan_axis0_vmap	5138	5069	17301	0.99

[samanklesaria]

If the input is of shape (1, A, 1) and we write back the residuals, this kernel will currently fail. That's because we try to write the result directly to GMEM, which currently fails for scalars: store_untiled always use the optimized=True flag, which isn't available for writing to GMEM. This is ideally fixed upstream: the current branch used for scalars should actually be conditional on the layout rather than the size: keep the scalar path for WGSplatFragLayout (and rank-0), and let a tiled value of any rank ≥ 1 fall through to the general case which is correct for it already.

[samanklesaria]

I've edited the code to make use of the warpgroup semantics and emit_pipeline functionality now available for Ampere.

[samanklesaria]

I'll remove these eventually. Currently trying out jj, which means explicitly accounting for all untracked files.