For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

gated_group_rmsnorm

`gated_group_rmsnorm()`

max.nn.legacy.state_space.gated_group_rmsnorm(y, gate, norm_weight, eps, group_size)

source

Fuses the gated group-RMSNorm into a single dispatch.

The norm is HF’s Zamba2RMSNormGated with norm_before_gate=False.

Collapses cast(y->f32) -> silu(gate)*y -> group rms_norm -> *norm_weight -> cast into one op. y and gate are [N, intermediate]; norm_weight is fp32 [intermediate]. Each contiguous group_size slice of the intermediate axis is normalized independently. Returns the model dtype (y.dtype), so the downstream out_proj cast is a no-op.

Parameters:

y (TensorValue) – The [N, intermediate] SSD scan output (model dtype).
gate (TensorValue) – The [N, intermediate] gate projection (any float dtype; may be a strided split view of the fused in-proj).
norm_weight (TensorValue) – The [intermediate] fp32 RMSNorm weight.
eps (float) – The epsilon inside rsqrt(mean_sq + eps).
group_size (int) – The width of each normalized group (intermediate // n_groups).

Returns:

The [N, intermediate] result in y.dtype.

Return type:

TensorValue

The following example normalizes a 256-wide intermediate in groups of 64:

import numpy as np
from max.driver import CPU, Accelerator, Buffer, accelerator_count
from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType
from max.nn.legacy.state_space import gated_group_rmsnorm

device = Accelerator() if accelerator_count() > 0 else CPU()
device_ref = DeviceRef.from_device(device)
num_tokens, intermediate, group_size = 4, 256, 64

with Graph(
    "gated_group_rmsnorm_example",
    input_types=(
        TensorType(
            DType.float32, [num_tokens, intermediate], device=device_ref
        ),
        TensorType(
            DType.float32, [num_tokens, intermediate], device=device_ref
        ),
        TensorType(DType.float32, [intermediate], device=device_ref),
    ),
) as graph:
    y, gate, norm_weight = (value.tensor for value in graph.inputs)
    output = gated_group_rmsnorm(
        y, gate, norm_weight, eps=1e-5, group_size=group_size
    )
    graph.output(output)

rng = np.random.default_rng(0)
y_np = rng.standard_normal((num_tokens, intermediate), dtype=np.float32)
gate_np = rng.standard_normal(
    (num_tokens, intermediate), dtype=np.float32
)
weight_np = rng.uniform(0.5, 1.5, intermediate).astype(np.float32)

model = InferenceSession(devices=[device]).load(graph)
inputs = [
    Buffer.from_numpy(x).to(device) for x in (y_np, gate_np, weight_np)
]
result = model.execute(*inputs)[0].to_numpy()

gated_group_rmsnorm()​

`gated_group_rmsnorm()`