> This page is for version 0.4.0 · 26.04.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-moe-swiglu_triton-SwigluClampedTriton) | `silu(clamp(gate)) * clamp(up) * probs` with the Triton forward and backward kernels. |

### Functions

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

### Data

[`HAVE_TRITON`](#nemo_automodel-components-moe-swiglu_triton-HAVE_TRITON)

[`_BWD_BLOCK_M`](#nemo_automodel-components-moe-swiglu_triton-_BWD_BLOCK_M)

[`_BWD_BLOCK_N`](#nemo_automodel-components-moe-swiglu_triton-_BWD_BLOCK_N)

[`_FWD_BLOCK_M`](#nemo_automodel-components-moe-swiglu_triton-_FWD_BLOCK_M)

[`_FWD_BLOCK_N`](#nemo_automodel-components-moe-swiglu_triton-_FWD_BLOCK_N)

[`_SUPPORTED_DTYPES`](#nemo_automodel-components-moe-swiglu_triton-_SUPPORTED_DTYPES)

### API

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

**Bases:** `Function`

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

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

staticmethod

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

staticmethod

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

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

```python
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
)
```

```python
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
)
```

```python
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.

```python
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).

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

```python
nemo_automodel.components.moe.swiglu_triton._BWD_BLOCK_M = 4
```

```python
nemo_automodel.components.moe.swiglu_triton._BWD_BLOCK_N = 512
```

```python
nemo_automodel.components.moe.swiglu_triton._FWD_BLOCK_M = 8
```

```python
nemo_automodel.components.moe.swiglu_triton._FWD_BLOCK_N = 256
```

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