nemo_automodel.components.moe.swiglu_triton

Triton kernels for the clamped SwiGLU expert activation (DeepSeek-V4 style, swiglu_limit > 0).

The routed experts apply silu(clamp(gate, max=L)) * clamp(up, -L, L) * probs to the grouped-GEMM output x = [gate | up] (bf16, [rows, 2 * inter]) with fp32 routing probabilities probs ([rows, 1]). The compiled PyTorch version (experts._swiglu_clamped_deepep_compiled) is one inductor kernel forward, but its autograd-derived backward concatenates d_gate and d_up and ran at about a quarter of HBM bandwidth on GB200 (0.7 ms for a 1.1 GB pass at the 16-node shape). These kernels do one pass each:

forward: a 2-D tile per program, fp32 math, result rounded to x.dtype;
backward: one program per block of rows looping over the columns, writing d_gate / d_up straight into the two halves of grad_x and accumulating the per-row probability gradient in registers.

Clamp gradients follow torch.clamp (inclusive range masks on the input values). The dispatcher in experts.swiglu_clamped_deepep uses these kernels on contiguous CUDA inputs when Triton is importable and falls back to the compiled version otherwise.

Module Contents

Classes

Name	Description
`SwigluClampedTriton`	`silu(clamp(gate)) * clamp(up) * probs` with the Triton forward and backward kernels.

Functions

Name	Description
`_backward`	-
`_forward`	-
`_swiglu_clamped_bwd_kernel`	-
`_swiglu_clamped_fwd_kernel`	-
`swiglu_clamped_triton`	Apply the clamped SwiGLU with the Triton kernels; callers check `uses_triton_kernel` first.
`uses_triton_kernel`	Whether the Triton kernels can take these inputs (contiguous 2-D CUDA `x`, one probability per row).

Data

API

class nemo_automodel.components.moe.swiglu_triton.SwigluClampedTriton()

Bases: Function

silu(clamp(gate)) * clamp(up) * probs with the Triton forward and backward kernels.

nemo_automodel.components.moe.swiglu_triton.SwigluClampedTriton.backward(
    ctx,
    grad_y: torch.Tensor
)

staticmethod

nemo_automodel.components.moe.swiglu_triton.SwigluClampedTriton.forward(
    ctx,
    x: torch.Tensor,
    probs: torch.Tensor,
    limit: float
) -> torch.Tensor

staticmethod

nemo_automodel.components.moe.swiglu_triton._backward(
    x: torch.Tensor,
    probs: torch.Tensor,
    grad_y: torch.Tensor,
    limit: float
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.moe.swiglu_triton._forward(
    x: torch.Tensor,
    probs: torch.Tensor,
    limit: float
) -> torch.Tensor

nemo_automodel.components.moe.swiglu_triton._swiglu_clamped_bwd_kernel(
    x_ptr,
    probs_ptr,
    grad_y_ptr,
    grad_x_ptr,
    grad_probs_ptr,
    rows,
    inter,
    stride_x,
    stride_gy,
    stride_gx,
    limit,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr
)

nemo_automodel.components.moe.swiglu_triton._swiglu_clamped_fwd_kernel(
    x_ptr,
    probs_ptr,
    y_ptr,
    rows,
    inter,
    stride_x,
    stride_y,
    limit,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr
)

nemo_automodel.components.moe.swiglu_triton.swiglu_clamped_triton(
    x: torch.Tensor,
    probs: torch.Tensor,
    limit: float
) -> torch.Tensor

Apply the clamped SwiGLU with the Triton kernels; callers check uses_triton_kernel first.

nemo_automodel.components.moe.swiglu_triton.uses_triton_kernel(
    x: torch.Tensor,
    probs: torch.Tensor
) -> bool

Whether the Triton kernels can take these inputs (contiguous 2-D CUDA x, one probability per row).

nemo_automodel.components.moe.swiglu_triton.HAVE_TRITON: bool = bool(_HAVE_TRITON and _HAVE_TRITON_LANGUAGE)

nemo_automodel.components.moe.swiglu_triton._BWD_BLOCK_M = 4

nemo_automodel.components.moe.swiglu_triton._BWD_BLOCK_N = 512

nemo_automodel.components.moe.swiglu_triton._FWD_BLOCK_M = 8

nemo_automodel.components.moe.swiglu_triton._FWD_BLOCK_N = 256

nemo_automodel.components.moe.swiglu_triton._SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)