> 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

Kimi-K3 SiTU activation and attention-residual compute cores.

Memory- and dispatch-optimized fp32 chains extracted from `model.py`:
the chunked bf16-saving weighted-SiTU autograd Function, the attn-res
mixing chain, and the opt-in `torch.compile` wrapper
(`BackendConfig.compile_situ`) shared by all of them.

## Module Contents

### Classes

| Name                                                                                            | Description                                                          |
| ----------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`_WeightedSiTUFunction`](#nemo_automodel-components-models-kimi_k3-situ-_WeightedSiTUFunction) | Chunked fp32 weighted-SiTU that saves only the low-precision inputs. |

### Functions

| Name                                                                                                | Description                                                                 |
| --------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_apply_attn_res`](#nemo_automodel-components-models-kimi_k3-situ-_apply_attn_res)                 | Mix `[tokens, hidden]` with prior `[tokens, blocks, hidden]` residuals.     |
| [`_attn_res_core`](#nemo_automodel-components-models-kimi_k3-situ-_attn_res_core)                   | fp32 attention-residual mixing chain.                                       |
| [`_compile_situ_cores`](#nemo_automodel-components-models-kimi_k3-situ-_compile_situ_cores)         | Wrap the SiTU chunk cores and the attn-res core with `torch.compile`.       |
| [`_situ_bwd_core`](#nemo_automodel-components-models-kimi_k3-situ-_situ_bwd_core)                   | Compute analytic fp32 SiTU gradients for one chunk of rows.                 |
| [`_situ_fwd_core`](#nemo_automodel-components-models-kimi_k3-situ-_situ_fwd_core)                   | Compute the fp32 SiTU chain for one chunk of rows.                          |
| [`_situ_rw_is_row_aligned`](#nemo_automodel-components-models-kimi_k3-situ-_situ_rw_is_row_aligned) | Return True when `routing_weights` carries one entry per `gate_up` row.     |
| [`_weighted_situ`](#nemo_automodel-components-models-kimi_k3-situ-_weighted_situ)                   | Apply SiTU and routing weights to `[tokens, 2 * intermediate]` projections. |

### Data

[`_SITU_CHUNK_ROWS`](#nemo_automodel-components-models-kimi_k3-situ-_SITU_CHUNK_ROWS)

[`_SITU_CHUNK_THRESHOLD`](#nemo_automodel-components-models-kimi_k3-situ-_SITU_CHUNK_THRESHOLD)

[`_SITU_CORES_COMPILED`](#nemo_automodel-components-models-kimi_k3-situ-_SITU_CORES_COMPILED)

### API

```python
class nemo_automodel.components.models.kimi_k3.situ._WeightedSiTUFunction()
```

**Bases:** `Function`

Chunked fp32 weighted-SiTU that saves only the low-precision inputs.

The forward computes the same fp32 chain as the eager `_weighted_situ`
path in row chunks (bitwise-identical result); the backward recomputes
the fp32 intermediates per chunk with analytic gradients that match
autograd's fp32 chain, so autograd never stores full-size fp32 copies of
the \[tokens, 2 \* intermediate] projections.

```python
nemo_automodel.components.models.kimi_k3.situ._WeightedSiTUFunction.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor | None, None, None]
```

staticmethod

Recompute fp32 per chunk and return analytic gradients.

**Parameters:**

**`ctx`** `Any`

Autograd context holding the saved low-precision inputs.

---

**`grad_out`** `torch.Tensor`

Upstream gradient with the forward's output shape
\[..., intermediate].

---

**Returns:** `torch.Tensor`

Tuple `(d_gate_up, d_routing_weights, None, None)` where

```python
nemo_automodel.components.models.kimi_k3.situ._WeightedSiTUFunction.forward(
    ctx: typing.Any,
    gate_up: torch.Tensor,
    routing_weights: torch.Tensor,
    beta: float,
    linear_beta: float | None
) -> torch.Tensor
```

staticmethod

Apply SiTU and routing weights chunk by chunk.

**Parameters:**

**`ctx`** `Any`

Autograd context; saves `gate_up` and `routing_weights`
in their original (typically bf16 / fp32) dtypes.

---

**`gate_up`** `torch.Tensor`

Gate+up projections of shape \[..., 2 \* intermediate],
gate in the first half of the last axis, up in the second.

---

**`routing_weights`** `torch.Tensor`

Routing weights, either row-aligned with shape
\[..., k] matching `gate_up`'s leading dimensions (typically
\[tokens, 1]) or broadcastable against them.

---

**`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 \`\`broadcast(gate\_up.shape\[:-1] + \[intermediate],

```python
nemo_automodel.components.models.kimi_k3.situ._apply_attn_res(
    prefix_sum: torch.Tensor,
    block_residual: torch.Tensor,
    projection: torch.nn.Linear,
    norm: nemo_automodel.components.models.kimi_k3.model.KimiRMSNorm
) -> torch.Tensor
```

Mix `[tokens, hidden]` with prior `[tokens, blocks, hidden]` residuals.

```python
nemo_automodel.components.models.kimi_k3.situ._attn_res_core(
    values: torch.Tensor,
    norm_weight: torch.Tensor,
    proj_weight: torch.Tensor,
    variance_epsilon: float,
    out_dtype: torch.dtype
) -> torch.Tensor
```

fp32 attention-residual mixing chain.

The weighted combine is multiply+sum rather than
`torch.matmul(prob[T, 1, B], values[T, B, H])`: cuBLAS dispatches that
degenerate batched-GEMM shape to non-tensor-core fp32 kernels
(magma\_sgemmEx / gemv2 — 4.3% of busy GPU time on a 256×GB200 Kimi-K3
profile). multiply+sum runs the same fp32 math on the elementwise/reduce
path (identical up to fp32 accumulation order, which is below bf16
resolution for typical shapes) and fuses cleanly under `torch.compile`
when `BackendConfig.compile_situ` is set.

**Parameters:**

**`values`** `torch.Tensor`

Stacked residuals of shape \[tokens, blocks+1, hidden] (block
residuals concatenated with the current prefix sum along axis 1).

---

**`norm_weight`** `torch.Tensor`

RMSNorm weight of shape \[hidden].

---

**`proj_weight`** `torch.Tensor`

Squeezed attn-res projection weight of shape \[hidden].

---

**`variance_epsilon`** `float`

RMSNorm epsilon.

---

**`out_dtype`** `torch.dtype`

dtype of the returned mixed tensor.

---

**Returns:** `torch.Tensor`

Mixed residual of shape \[tokens, hidden] in `out_dtype`.

```python
nemo_automodel.components.models.kimi_k3.situ._compile_situ_cores() -> None
```

Wrap the SiTU chunk cores and the attn-res core with `torch.compile`.

Runs once per process: the compiled functions replace the module-level
eager cores, so every layer shares the same compiled kernels and repeated
model construction does not recompile. Compilation itself is lazy (at
first call). Compiled numerics are allclose to eager, not
bitwise-identical.

```python
nemo_automodel.components.models.kimi_k3.situ._situ_bwd_core(
    g: torch.Tensor,
    u0: torch.Tensor,
    w: torch.Tensor,
    go: torch.Tensor,
    beta: float,
    linear_beta: float | None,
    want_drw: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
```

Compute analytic fp32 SiTU gradients for one chunk of rows.

**Parameters:**

**`g`** `torch.Tensor`

fp32 gate projections of shape \[rows, intermediate].

---

**`u0`** `torch.Tensor`

fp32 up projections of shape \[rows, intermediate].

---

**`w`** `torch.Tensor`

fp32 routing weights broadcastable to \[rows, intermediate],
typically of shape \[rows, 1].

---

**`go`** `torch.Tensor`

fp32 upstream gradient of shape \[rows, intermediate].

---

**`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 the routing-weight gradient reduction is needed.

---

**Returns:** `torch.Tensor`

Tuple of fp32 tensors `(d_g, d_u, red)` where `d_g` and `d_u`

```python
nemo_automodel.components.models.kimi_k3.situ._situ_fwd_core(
    g: torch.Tensor,
    u0: torch.Tensor,
    w: torch.Tensor,
    beta: float,
    linear_beta: float | None
) -> torch.Tensor
```

Compute the fp32 SiTU chain for one chunk of rows.

**Parameters:**

**`g`** `torch.Tensor`

fp32 gate projections of shape \[rows, intermediate].

---

**`u0`** `torch.Tensor`

fp32 up projections of shape \[rows, intermediate].

---

**`w`** `torch.Tensor`

fp32 routing weights broadcastable to \[rows, intermediate],
typically of shape \[rows, 1].

---

**`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`

fp32 tensor of shape \[rows, intermediate]: `situ(g) * up(u0) * w`.

```python
nemo_automodel.components.models.kimi_k3.situ._situ_rw_is_row_aligned(
    gate_up: torch.Tensor,
    routing_weights: torch.Tensor
) -> bool
```

Return True when `routing_weights` carries one entry per `gate_up` row.

**Parameters:**

**`gate_up`** `torch.Tensor`

Gate+up projections of shape \[..., 2 \* intermediate].

---

**`routing_weights`** `torch.Tensor`

Routing weights; row-aligned when its shape is
\[..., k] with the same leading dimensions as `gate_up`.

---

```python
nemo_automodel.components.models.kimi_k3.situ._weighted_situ(
    gate_up: torch.Tensor,
    routing_weights: torch.Tensor,
    beta: float,
    linear_beta: float | None
) -> torch.Tensor
```

Apply SiTU and routing weights to `[tokens, 2 * intermediate]` projections.

```python
nemo_automodel.components.models.kimi_k3.situ._SITU_CHUNK_ROWS = 32768
```

```python
nemo_automodel.components.models.kimi_k3.situ._SITU_CHUNK_THRESHOLD = 12288
```

```python
nemo_automodel.components.models.kimi_k3.situ._SITU_CORES_COMPILED = False
```