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 class

GatedActivation

GatedActivation

class max.nn.moe.GatedActivation(fused_activation='silu', alpha=0.0, limit=0.0, gated_activation_fn=None)

source

Bases: object

The gated activation between an MoE’s gate-up and down projections.

Kernels that fuse the activation compute fused_activation, named like the HF config hidden_act:

• "silu": silu(gate) * up. A non-zero limit first clamps the inputs to min(gate, limit) and clamp(up, -limit, limit).
• "swigluoai": g * sigmoid(alpha * g) * (u + 1) with g = min(gate, limit) and u = clamp(up, -limit, limit).
• "gelu_tanh": gelu(gate) * up with the tanh approximation gelu(x) = 0.5 * x * (1 + tanh(0.7978845608028654 * (x + 0.044715 * x**3))).

gated_activation_fn replaces the activation with graph ops, for activations no kernel fuses. It takes the [gate | up] projection and moe_dim; see make_concatenated_gated_activation_fn(). A layer that has one runs it instead of any fused kernel, which costs performance: the activation runs as separate graph ops over the whole gate-up output, and the fused grouped-matmul and quantize paths are skipped. Prefer a fused_activation whenever one matches.

Parameters:

• fused_activation (str)
• alpha (float)
• limit (float)
• gated_activation_fn (Callable[[TensorValue, int], TensorValue] | None)

alpha

alpha: float = 0.0

source

The sigmoid scale of "swigluoai"; must be 0 for the others.

fused_activation

fused_activation: str = 'silu'

source

The activation fused kernels compute.

gated_activation_fn

gated_activation_fn: Callable[[TensorValue, int], TensorValue] | None = None

source

Graph-op replacement for the activation, or None to fuse it.

kernel_parameters()

kernel_parameters()

source

Returns fused_activation encoded as kernel parameters.

Raises:

ValueError – If gated_activation_fn is set; a kernel can only fuse fused_activation.

Return type:

dict[str, float | str]

limit

limit: float = 0.0

source

The clamp bound of "swigluoai" and "silu" (0 leaves "silu" unclamped); must be 0 for "gelu_tanh".