Implement codegen for optimized RMSNorm kernel in LLaMA4

In the context of RMSNorm in LLaMA4, I observed that mapping independent tokens (rows) to individual waves yields a significant speedup, around 4× faster than the current naive version emitted by Wave.

The idea is to assign one wave per token, so that each wave performs the norm computation and scaling independently. This avoids cross-wave synchronization and the need for shared memory, leading to better performance.
However the benefit is dependent on the available parallelism. In low-token scenarios (e.g., 1 token) the speedup is not better.

We could maybe generalize this mapping strategy to other kernels by automatically performing wave-level partitioning of independent work, especially when parallelism is high enough to benefit from it.

You can find an example of such a handwritten kernels and performance comparison in the following commit: https://github.com/adedespirlet/wave/commit/f7427b24a1a62afc8adef645c8d098bfdc1b151b