[vla, diffusion] perf: fuse FastWAM DiT q/k RMSNorm into F.rms_norm

[XueSongTap]

Summary

Replace the hand-rolled RMS normalization in RMSNorm.forward in loongforge/embodied/model/fastwam/wan/dit.py with F.rms_norm, which dispatches to the fused ATen kernel. This module normalizes the q/k projections in every self- and cross-attention block of the FastWAM Wan video DiT, so it runs twice per attention op on every layer — a genuinely hot path.

The previous implementation expanded into a chain of separate kernels: an fp32 up-cast, pow, mean, rsqrt, a broadcast multiply, a down-cast, and a final weight multiply. Each step materializes an intermediate that the backward pass then has to keep alive. F.rms_norm collapses the sequence into one fused forward and one fused backward, cutting both launch overhead and activation memory at this spot.

Dtype behavior is preserved deliberately. F.rms_norm only fuses the weight multiply when weight and x share a dtype, so when they differ we normalize without the weight and apply the scale separately. That keeps the original promotion rule intact: an fp32 weight over a bf16 activation still yields an fp32 result.

Numerics

Normalization still accumulates in fp32 — the fused kernel uses an fp32 accumulator for reduced-precision inputs. The one behavioral difference is ordering in the same-dtype path: the previous code down-cast before multiplying by weight, whereas the fused kernel scales before the down-cast. Results are therefore very close to, but not bit-identical with, the previous implementation.

Impact