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`	Chunked fp32 weighted-SiTU that saves only the low-precision inputs.

Functions

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

Data

_SITU_CHUNK_ROWS

_SITU_CHUNK_THRESHOLD

_SITU_CORES_COMPILED

API

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.

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

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],

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.

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.

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.

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:

torch.Tensor

fp32 gate projections of shape [rows, intermediate].

torch.Tensor

fp32 up projections of shape [rows, intermediate].

torch.Tensor

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

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

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:

torch.Tensor

fp32 gate projections of shape [rows, intermediate].

torch.Tensor

fp32 up projections of shape [rows, intermediate].

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.

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.

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.

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

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

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