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).
Mojo struct
GatedGroupRMSNorm
struct GatedGroupRMSNorm[group_size: Int]
Fused silu-gate + group RMSNorm + weight-scale for the Mamba-2 mixer.
Collapses cast -> silu(gate) * y -> group rms_norm -> * norm_weight -> cast
into one dispatch, matching HF Zamba2RMSNormGated with
norm_before_gate=False. The registration lives here in the built-in kernel
library (mirroring the causal_conv1d_varlen_fwd /
mamba2_ssd_chunk_scan_varlen_fwd_inplace precedents) so the graph compiler
/ serve path resolves the op with no out-of-tree custom_extensions. The
kernel math lives in state_space.gated_group_rmsnorm.
Tensor shapes:
- output: (n_rows, intermediate) - model dtype.
- y: (n_rows, intermediate) - SSD scan output, model dtype.
- gate: (n_rows, intermediate) - gate projection (any float dtype;
may be a strided split view of the fused in-proj).
- weight: (intermediate,) - fp32 RMSNorm weight.
- eps: Scalar epsilon (fp32) inside rsqrt(mean_sq + eps).
Parameters
- group_size (
Int): Width of each independently normalized group along the intermediate axis (intermediate // n_groups).
Implemented traits
Methods
execute
static def execute[dtype: DType, gate_dtype: DType, target: StringSlice[ImmStaticOrigin]](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], y: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=y.static_spec], gate: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate.static_spec], weight: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=weight.static_spec], eps: Float32, ctx: DeviceContext)