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

fused_silu

fused_silu()

max.experimental.nn.common_layers.functional_kernels.fused_silu(input, row_offsets)

source

Performs the fused SILU operation for all the MLPs in the EP MoE module.

After the EP dispatch phase, the host does not know how many tokens were actually received, so this op reads the row offsets to determine the received token count and applies the SILU operation only to the received tokens.

Parameters:

• input (TensorValue) – The input tokens to apply the SILU operation to. Shape: [max_recv_tokens, hidden_size].
• row_offsets (TensorValue) – The row offsets that determine the actual number of received tokens in the input tensor. Shape: [n_local_experts + 1].

Returns:

[max_recv_tokens, hidden_size].

Return type:

The output tokens after the SILU operation. Shape