> 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.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`](#nemo_automodel-components-models-kimi_k3-situ_triton-_check_2d_rows) | - | | [`_situ_bwd_kernel`](#nemo_automodel-components-models-kimi_k3-situ_triton-_situ_bwd_kernel) | Analytic SiTU gradients for BLOCK\_R full rows; the routing-weight gradient is reduced in-kernel. | | [`_situ_fwd_kernel`](#nemo_automodel-components-models-kimi_k3-situ_triton-_situ_fwd_kernel) | out\[r, c] = beta \* tanh(g / beta) \* sigmoid(g) \* up(u) \* w\[r] for one (rows x cols) tile. | | [`_tile_configs`](#nemo_automodel-components-models-kimi_k3-situ_triton-_tile_configs) | Autotune candidates: \~16 fp32 values per thread keeps the backward below the register cap. | | [`situ_bwd_triton`](#nemo_automodel-components-models-kimi_k3-situ_triton-situ_bwd_triton) | Weighted (or dense) SiTU backward. | | [`situ_fwd_triton`](#nemo_automodel-components-models-kimi_k3-situ_triton-situ_fwd_triton) | Weighted (or dense) SiTU forward on `[rows, 2 * intermediate]` projections. | ### Data [`HAVE_TRITON`](#nemo_automodel-components-models-kimi_k3-situ_triton-HAVE_TRITON) ### API ```python nemo_automodel.components.models.kimi_k3.situ_triton._check_2d_rows( name: str, t: torch.Tensor ) -> None ``` ```python 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. ```python 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. ```python 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. ```python 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 ```python 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. ```python nemo_automodel.components.models.kimi_k3.situ_triton.HAVE_TRITON = True ```