[vla] perf: use fused F.rms_norm in FastWAM VAE RMSNorm

[XueSongTap]

Replace the hand-written RMSNorm in the Wan VAE with F.rms_norm.

The previous implementation expressed RMS normalization as F.normalize(x, dim) * dim**0.5 * gamma + bias, which launches four separate elementwise kernels over the full activation. F.rms_norm does the same work — mean(x^2) -> rsqrt -> scale by gamma — in a single fused kernel.

F.rms_norm normalizes over trailing dimensions only, while this module normalizes over the channel axis of (B, C, T, H, W) tensors. The channel-first path therefore moves channels to the last axis with movedim(1, -1) and moves them back afterwards; movedim only changes strides and does not copy.

Why it is equivalent

F.normalize(x, dim) is x / max(||x||_2, eps); multiplying by sqrt(dim) gives x / sqrt(mean(x^2)), which is the RMSNorm definition. The only difference is where eps is applied — clamped in the denominator before, added inside the square root now. eps is set explicitly to 1e-12 to match the previous F.normalize default.

Verification

• Training loss curves match the baseline.

• End-to-end throughput: ~1% improvement in the typical case.

Notes

• Requires PyTorch >= 2.4 for F.rms_norm. No fallback is kept.
• Removes the unused bias argument from RMSNorm.__init__. It defaulted to False and no call site ever set it, so no parameter was created and checkpoint compatibility is unaffected.

[nullnonenilNULL]

It is recommended to retain the bias support path rather than remove it entirely.

[XueSongTap]

The latest commit restores the bias support path while retaining the fused VAE RMSNorm implementation.