RMSNorm

class sdm.nn.RMSNorm(normalized_shape: int | list[int] | Size, eps: float | None = None, elementwise_affine: bool = True, device=None, dtype=None)

Bases: RMSNorm

A torch.nn.RMSNorm that dispatches to an efficient kernel.

Parameters:

• normalized_shape (tuple[int, ...])

• eps (float | None)

• elementwise_affine (bool)

forward(x: Tensor, rope: RotaryEmbedding | None = None) → Tensor

Apply optional rotary embeddings before normalization.

Parameters:

• x (Tensor) – Input with shape [..., C], or [..., S, H, C] with rotary embeddings, where S is the sequence length, H is the number of heads, and C is channels per head.

• rope (RotaryEmbedding | None) – Optional rotary embedding applied to the input.

Return type:

Tensor