align RMSNorm API with hipDNN

[AaronStGeorge]

Why

RMSNorm over an example 1d vector is defined as

$$y_i = \frac{x_i}{\texttt{RMS}(x)} \cdot \gamma_i$$ $$\texttt{RMS}(x) = \sqrt{\epsilon + \frac{1}{n}\sum_{j=1}^{n} x_j^2}$$

$\texttt{RMS}$ is computed over what PyTorch calls normalized_shape.

Fusilli currently always uses all non-batch dimensions for normalized_shape. For example with a [N, C, H, W] fusilli computes $\texttt{RMS}$ over [C, H, W] regardless of what scale is set and expects inverse rms output to be [N, 1, 1, 1] for all input shapes.

The semantics diverge from cuDNN/hipDNN which derive normalized_shape from the maximal trailing suffix of input where scale[i] == input[i].

cuDNN's rules for determining normalized shape from scale:

• scale rank == input rank
• scale[0] == 1 (batch broadcast)
• normalized shape = maximal trailing suffix where scale[i] == input[i]; must be non-empty
• in the leading region (dims before the trailing match), scale[i] == 1

Inverse RMS shape is always the "kept dimensions" aka inputs dimensions where scale[i] == 1, exactly the dimensions that are not normalized over.

Examples:

• input [N, C, H, W], scale [1, C, H, W] -> normalized_shape = [C, H, W], inv_rms = [N, 1, 1, 1]
• input [N, C, H, W], scale [1, 1, H, W] ->normalized_shape = [H, W], inv_rms = [N, C, 1, 1]
• input [N, C, H, W], scale [1, 1, 1, W] -> normalized_shape = [W], inv_rms = [N, C, H, 1]
• (degenerate but valid) input [N, 1, 1, 1], scale [1, 1, 1, 1] -> normalized_shape = [1,1,1], inv_rms = [N, 1, 1, 1]

Note: cuDNN and PyTorch frontends require normalized_shape to match input dims without broadcasts. For example, input [N, C, H, W], scale [1, 1, H, 1] would be rejected rather than assuming that scale is broadcast over W dimension.

What

• Updates scale to required parameter for RMSNorm
• Updates normalized_shape calculation and emission to match cuDNN/hipDNN semantics
• Updates expected inv_rms shape to be the kept dimensions

[sjain-stanford]

Could you populate the PR description with the "what" and "why" for this change? The title says align the API with hipdnn but I don't know what was misaligned and how we're fixing that. The PR description also gets rolled into the final commit message so good for record-keeping (especially for agents to look back at the history).

[sjain-stanford]

I think we should make normalized_shape an input for rmsnorm and layernorm and not derive normalized_shape from the maximal trailing suffix of input. Here's why:

For it to work, we need either scale to be required (to infer norm_shape) or we require a norm shape to be specified.

Both layernorm and RMSnorm in PT don't require a scale -> only if elementwise_affine is true. Making it required in fusilli is a bit jarring if users trained without scaling.

So the least common divisor is to expect a norm_shape on both, and then let scale / bias be same shape as norm_shape.

We should do this at fusilli layer. hipdnn can do what it does today which is using trailing suffix of input for norm shape, and feed that "inferred" normalized shape into fusilli's op builder.