[CUDA][Improvement] RMSNorm forward speed up

[nastya236]

We have 2 variants of RNSNorm and for majority of workloads we were using res_norm_small:

• each thread does 1 vectorised load of width 8 (for bf16)
• 1 block per row, so block size 1024 for N = 8192

Now block size is 128, we do multiple vectorised loads per thread.
The issue is the register pressure when using safe load_vector with loop fall back. Therefore in a new version there is an alignment check. But in majority of cases x, w and output should be 16 bytes aligned.

b200, pytorch 2.13, cuda 13.3

N	M	Torch (µs)	MLX (µs, before)	MLX (µs, after)
1024	32768	51.7	44.9	39.2
2048	32768	57.7	77.9	53.0
3072	32768	83.9	142.5	70.2
4096	32768	92.7	146.4	89.5
7168	32768	160.3	307.0	146.1
8192	32768	180.7	318.2	170.6

[angeloskath]

🚀