nemo_automodel.components.models.kimi_k3.situ_triton

Hand-written Triton kernels for the Kimi-K3 SiTU activation.

The SiTU chain (beta * tanh(g / beta) * sigmoid(g) * linear_beta * tanh(u / linear_beta) * w) is a pure elementwise function of the gate/up halves of one [rows, 2 * intermediate] projection plus one routing weight per row. torch.compile (BackendConfig.compile_situ) fuses it into a single kernel, but the generated 1-D pointwise kernels pay for two things on every element: 64-bit div/mod to recover the row index for the row-broadcast routing weight (the pointwise index space is flattened), and, in the backward, both cat branches evaluated under tl.where (twelve masked loads and four tanh per output element). On GB200 that leaves the backward ~10x above its bandwidth bound (17 ms per 262k x 6144 bf16 call on a 2-node K3 profile).

These kernels tile the problem in 2-D (rows x columns) instead: the row index is a per-tile 64-bit multiply, the column index stays 32-bit, each transcendental is evaluated once, d_gate and d_up are stored straight into the two halves of the output, and the routing-weight gradient (“sum_cols(go * situ(g)

up(u))) is reduced inside the backward kernel while the tile is in registers. The fp32 math and its operation order are identical to situ._situ_fwd_core/situ._situ_bwd_core“; only the fp32 accumulation order of the routing-weight reduction differs.

Module Contents

Functions

Name	Description
`_check_2d_rows`	-
`_situ_bwd_kernel`	Analytic SiTU gradients for BLOCK_R full rows; the routing-weight gradient is reduced in-kernel.
`_situ_fwd_kernel`	out[r, c] = beta * tanh(g / beta) * sigmoid(g) * up(u) * w[r] for one (rows x cols) tile.
`_tile_configs`	Autotune candidates: ~16 fp32 values per thread keeps the backward below the register cap.
`situ_bwd_triton`	Weighted (or dense) SiTU backward.
`situ_fwd_triton`	Weighted (or dense) SiTU forward on `[rows, 2 * intermediate]` projections.

Data

HAVE_TRITON

API

nemo_automodel.components.models.kimi_k3.situ_triton._check_2d_rows(
    name: str,
    t: torch.Tensor
) -> None

nemo_automodel.components.models.kimi_k3.situ_triton._situ_bwd_kernel(
    gu_ptr,
    rw_ptr,
    go_ptr,
    dgu_ptr,
    drw_ptr,
    n_rows,
    half,
    stride_gu,
    stride_go,
    stride_dgu,
    beta,
    linear_beta,
    HAS_RW: triton.language.constexpr,
    HAS_LINEAR: triton.language.constexpr,
    WANT_DRW: triton.language.constexpr,
    BLOCK_R: triton.language.constexpr,
    BLOCK_C: triton.language.constexpr
)

Analytic SiTU gradients for BLOCK_R full rows; the routing-weight gradient is reduced in-kernel.

nemo_automodel.components.models.kimi_k3.situ_triton._situ_fwd_kernel(
    gu_ptr,
    rw_ptr,
    out_ptr,
    n_rows,
    half,
    stride_gu,
    stride_out,
    beta,
    linear_beta,
    HAS_RW: triton.language.constexpr,
    HAS_LINEAR: triton.language.constexpr,
    BLOCK_R: triton.language.constexpr,
    BLOCK_C: triton.language.constexpr
)

out[r, c] = beta * tanh(g / beta) * sigmoid(g) * up(u) * w[r] for one (rows x cols) tile.

nemo_automodel.components.models.kimi_k3.situ_triton._tile_configs() -> list

Autotune candidates: ~16 fp32 values per thread keeps the backward below the register cap.

nemo_automodel.components.models.kimi_k3.situ_triton.situ_bwd_triton(
    gate_up2: torch.Tensor,
    routing_weights2: torch.Tensor | None,
    grad_out2: torch.Tensor,
    beta: float,
    linear_beta: float | None,
    want_drw: bool
) -> tuple[torch.Tensor, torch.Tensor | None]

Weighted (or dense) SiTU backward.

Parameters:

gate_up2

torch.Tensor

Saved gate+up projections of shape [rows, 2 * intermediate].

routing_weights2

torch.Tensor | None

Saved routing weights of shape [rows, 1], or None (dense).

grad_out2

torch.Tensor

Upstream gradient of shape [rows, intermediate], unit last stride.

beta

float

SiTU beta applied to the gate branch.

linear_beta

float | None

Optional bounded-linear beta applied to the up branch.

want_drw

bool

Whether to reduce the routing-weight gradient (requires weights).

Returns: torch.Tensor

(d_gate_up2, d_routing_weights2) in the inputs’ dtypes; the second

nemo_automodel.components.models.kimi_k3.situ_triton.situ_fwd_triton(
    gate_up2: torch.Tensor,
    routing_weights2: torch.Tensor | None,
    beta: float,
    linear_beta: float | None
) -> torch.Tensor

Weighted (or dense) SiTU forward on [rows, 2 * intermediate] projections.

Parameters:

gate_up2

torch.Tensor

Gate+up projections of shape [rows, 2 * intermediate] on a CUDA device, unit stride along the last axis; gate in the first half.

routing_weights2

torch.Tensor | None

Optional routing weights of shape [rows, 1] (any float dtype, contiguous), or None for the dense activation.

beta

float

SiTU beta applied to the gate branch.

linear_beta

float | None

Optional bounded-linear beta applied to the up branch.

Returns: torch.Tensor

Tensor of shape [rows, intermediate] in gate_up2’s dtype.

nemo_automodel.components.models.kimi_k3.situ_triton.HAVE_TRITON = True