vllm.models.step5.nvidia.layernorm

Step5 normalization layers.

Classes:

• Step5IndexerLayerNorm –

  Per-head LayerNorm for the Step5 sparse-attention indexer.

Step5IndexerLayerNorm

Bases: Module

Per-head LayerNorm for the Step5 sparse-attention indexer.

Source code in vllm/models/step5/nvidia/layernorm.py

class Step5IndexerLayerNorm(nn.Module):
    """Per-head LayerNorm for the Step5 sparse-attention indexer."""

    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        normalized = F.layer_norm(
            x.unflatten(-1, (-1, self.hidden_size)),
            (self.hidden_size,),
            self.weight,
            self.bias,
            self.variance_epsilon,
        )
        return normalized.flatten(-2, -1)