Is this page helpful?

core.inference.moe.activations#

Padding-aware activation kernels for fused MoE.

These kernels skip padding rows (where permutation_map == -1) to avoid wasted computation on aligned-but-empty expert slots.

Module Contents#

Functions#

_ceil_div

_clamped_relu

ReLU, optionally tanh soft-clamped to (0, clamp_scale).

_squared_relu_kernel

Squared ReLU that skips rows beyond n_used and alignment-padding rows (perm_map == -1).

padded_squared_relu

Squared ReLU activation that skips rows beyond n_used and alignment-padding rows.

_swiglu_kernel

SwiGLU: SiLU(gate) * up, skipping rows beyond n_used and padding rows (perm_map == -1).

padded_swiglu

SwiGLU activation (SiLU(gate) * up); skips rows beyond n_used and alignment-padding rows.

_silu_mul_bounded_kernel

SiLU(gate) * up on the first n_rows rows; rows beyond are left untouched.

bounded_silu_mul

SwiGLU (SiLU(gate) * up) over the first n_rows rows of a flat [M, 2N] tensor.

_squared_relu_quantize_kernel

Fused squared ReLU + MXFP8 quantize + swizzle in one kernel.

squared_relu_and_quantize_mxfp8

Fused squared ReLU + MXFP8 quantize + swizzle.

API#

core.inference.moe.activations._ceil_div(a, b)#
core.inference.moe.activations._clamped_relu(x, clamp_scale, CLAMP: triton.language.constexpr)#

ReLU, optionally tanh soft-clamped to (0, clamp_scale).

Returns the pre-square value of the squared-ReLU activation. When CLAMP is set this is clamp_scale * tanh(ReLU(x) / clamp_scale), which equals training’s ReLU(clamp_scale * tanh(x / clamp_scale)): the soft clamp is non-decreasing and maps 0 to 0, so it commutes with the ReLU.

The clamped value stays in FP32, matching training’s fused weighted_clamped_squared_relu, which squares clamp_scale * tanh(...) directly with no intermediate downcast.

core.inference.moe.activations._squared_relu_kernel(
input_ptr,
output_ptr,
src_idx_ptr,
n_used_ptr,
N,
max_rows,
clamp_scale,
CLAMP: triton.language.constexpr,
BLOCK_N: triton.language.constexpr,
NUM_BLOCKS: triton.language.constexpr,
ZERO_PADDING: triton.language.constexpr,
)#

Squared ReLU that skips rows beyond n_used and alignment-padding rows (perm_map == -1).

Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple rows. n_used_ptr gates how many rows are processed — required for CUDA graph compatibility. MXFP8 quantization and grouped GEMM do not receive perm_map. ZERO_PADDING therefore materializes neutral values for dummy rows inside each expert segment.

core.inference.moe.activations.padded_squared_relu(
x: torch.Tensor,
permutation_map: torch.Tensor,
n_used: torch.Tensor,
clamp_scale: Optional[float] = None,
zero_padding: bool = False,
) → torch.Tensor#

Squared ReLU activation that skips rows beyond n_used and alignment-padding rows.

Parameters:

  • x – [output_size, ffn_hidden] BF16 FC1 output.

  • permutation_map – [output_size] int32, original token index or -1 for padding.

  • n_used – scalar int32 CUDA tensor = inclusive_expert_offsets[-1].

  • clamp_scale – config.activation_func_tanh_clamp_scale. If set, soft-clamp the pre-activation with s * tanh(x / s) first, bounding the output by s ** 2.

  • zero_padding – write zeros to alignment-padding rows instead of leaving them undefined. MXFP8 quantization and grouped GEMM cannot skip dummy rows through permutation_map. Rows beyond n_used remain undefined.

core.inference.moe.activations._swiglu_kernel(
input_ptr,
output_ptr,
src_idx_ptr,
n_used_ptr,
N,
max_rows,
BLOCK_N: triton.language.constexpr,
NUM_BLOCKS: triton.language.constexpr,
ZERO_PADDING: triton.language.constexpr,
)#

SwiGLU: SiLU(gate) * up, skipping rows beyond n_used and padding rows (perm_map == -1).

Input row width is 2N: gate = first N cols, up = last N cols (megatron chunk convention). Output row width is N. Fixed NUM_BLOCKS CTAs iterating rows -> CUDA-graph compatible. MXFP8 quantization and grouped GEMM do not receive perm_map. ZERO_PADDING therefore materializes neutral values for dummy rows inside each expert segment.

core.inference.moe.activations.padded_swiglu(
x: torch.Tensor,
permutation_map: torch.Tensor,
n_used: torch.Tensor,
zero_padding: bool = False,
) → torch.Tensor#

SwiGLU activation (SiLU(gate) * up); skips rows beyond n_used and alignment-padding rows.

Gated counterpart of padded_squared_relu: FC1 output is 2x wide (gate | up), so the output width is half the input.

Parameters:

  • x – [output_size, 2 * ffn_hidden] BF16 FC1 output.

  • permutation_map – [output_size] int32, original token index or -1 for padding.

  • n_used – scalar int32 CUDA tensor = inclusive_expert_offsets[-1].

  • zero_padding – write zeros to alignment-padding rows instead of leaving them undefined. MXFP8 quantization and grouped GEMM cannot skip dummy rows through permutation_map. Rows beyond n_used remain undefined.

Returns:

[output_size, ffn_hidden] BF16.

core.inference.moe.activations._silu_mul_bounded_kernel(
input_ptr,
output_ptr,
n_rows_ptr,
N,
max_rows,
BLOCK_N: triton.language.constexpr,
NUM_BLOCKS: triton.language.constexpr,
)#

SiLU(gate) * up on the first n_rows rows; rows beyond are left untouched.

For the flat token-major FC1 layout (no permutation map): rows [0, n_rows) are live. n_rows is a device scalar, so the fixed grid stays CUDA-graph compatible.

core.inference.moe.activations.bounded_silu_mul(
x: torch.Tensor,
n_rows: torch.Tensor,
) → torch.Tensor#

SwiGLU (SiLU(gate) * up) over the first n_rows rows of a flat [M, 2N] tensor.

Parameters:

  • x – [M, 2 * ffn_hidden] BF16 FC1 output (gate = first half, up = second half).

  • n_rows – scalar int32/int64 CUDA tensor with the number of live rows (e.g. valid_tokens * topk). Rows >= n_rows are skipped, not zeroed.

Returns:

[M, ffn_hidden] BF16 (rows >= n_rows undefined).

core.inference.moe.activations._squared_relu_quantize_kernel(
input_ptr,
out_fp8_ptr,
out_scale_ptr,
src_idx_ptr,
n_used_ptr,
K,
n_col_blocks,
max_rows,
clamp_scale,
CLAMP: triton.language.constexpr,
REAL_GROUPS: triton.language.constexpr,
BLOCK_K: triton.language.constexpr,
BLOCK_GROUPS: triton.language.constexpr,
NUM_BLOCKS: triton.language.constexpr,
)#

Fused squared ReLU + MXFP8 quantize + swizzle in one kernel.

Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple rows. Rows beyond n_used and alignment-padding rows (perm_map == -1) are skipped.

core.inference.moe.activations.squared_relu_and_quantize_mxfp8(
x: torch.Tensor,
permutation_map: torch.Tensor,
n_used: torch.Tensor,
clamp_scale: Optional[float] = None,
)#

Fused squared ReLU + MXFP8 quantize + swizzle.

Reads BF16 FC1 output, applies squared ReLU, quantizes to FP8 with swizzled scales. Single kernel replaces padded_squared_relu + mxfp8_quantize.

Parameters:

  • x – [output_size, K] BF16 FC1 output.

  • permutation_map – [output_size] int32, original token index or -1 for padding.

  • n_used – scalar int32 CUDA tensor = inclusive_expert_offsets[-1]. Rows beyond this are skipped before even checking the permutation_map.

  • clamp_scale – config.activation_func_tanh_clamp_scale. If set, soft-clamp the pre-activation with s * tanh(x / s) before the square.

Returns:

MXFP8Tensor with .data [output_size, K] float8_e4m3fn and .scale (swizzled e8m0).