nemo_automodel.components.models.kimi_k3.situ

Bases: Function

Attention-residual mix on the fused Triton kernels; saves only the inputs and 3 x (k+1) fp32 stats per token.

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

nemo_automodel.components.models.kimi_k3.situ._AttnResTritonFunction.forward(
    ctx: typing.Any,
    prefix_sum: torch.Tensor,
    block_residual: torch.Tensor,
    norm_weight: torch.Tensor,
    proj_weight: torch.Tensor,
    eps: float
) -> torch.Tensor

Bases: Function

Dense SiTU (no routing weights) on the Triton kernels, saving only the low-precision input.

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

nemo_automodel.components.models.kimi_k3.situ._DenseSiTUFunction.forward(
    ctx: typing.Any,
    x: torch.Tensor,
    beta: float,
    linear_beta: float | None
) -> torch.Tensor

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]

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

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

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:

Returns: torch.Tensor

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

Return True when the fused Triton kernels handle this attention-residual mix call.

Parameters:

Wrap the RMSNorm core with torch.compile (BackendConfig.compile_norm).

Runs once per process, same lazy pattern as _compile_situ_cores: the compiled function replaces the module-level eager core, so every KimiRMSNorm instance shares one compiled kernel. Compiled numerics are allclose to eager, not bitwise-identical.

Wrap the SiTU 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). On this path _WeightedSiTUFunction runs each pass as one fused whole-tensor kernel (the eager row-chunk loop, its casts and slice copies are skipped) and SituAndMul uses the compiled dense core. Compiled numerics are allclose to eager, not bitwise-identical.

SiTU gated activation for dense / shared-expert MLPs (SituAndMul).

Parameters:

Returns: torch.Tensor

Tensor of shape […, intermediate] in x’s dtype.

Route the attention-residual mix through the fused Triton kernels (KimiK3TextConfig.attn_res_triton).

Runs once per process. _apply_attn_res then launches attn_res_triton.attn_res_fwd_triton / attn_res_bwd_triton (no torch.cat, no fp32 copy of the stacked entries) for CUDA inputs with at most attn_res_triton.MAX_ENTRIES block entries; every other call keeps the eager or compile_situ chain. Without Triton the flag is a no-op. fp32 math matches _attn_res_core up to fp32 accumulation order.

Route SiTU through the hand-written Triton kernels (situ_backend = triton / triton_fast_math).

Runs once per process. _WeightedSiTUFunction (row-aligned [rows, 1] weights on a CUDA device) and SituAndMul then launch situ_triton.situ_fwd_triton / situ_bwd_triton instead of the eager chunk loop or the compile_situ inductor kernels; every other shape keeps its existing path. Without Triton the flag is a no-op (the eager or compiled path stays in place). fp32 math and operation order match the eager cores; only the routing-weight gradient’s fp32 accumulation order differs. fast_math (situ_backend = "triton_fast_math") selects the SFU variants of the kernels; once any layer of the process asks for them every SiTU call uses them.

Dispatch RMSNorm through the current (eager or compiled) core.

Indirection so KimiRMSNorm.forward picks up the compiled core even when _compile_norm_core runs after modules were constructed.

Parameters:

Returns: torch.Tensor

Normalized tensor with the same shape and dtype as hidden_states.

fp32 RMSNorm chain shared by every KimiRMSNorm instance.

Parameters:

Returns: torch.Tensor

Normalized tensor with the same shape and dtype as hidden_states.

Compute analytic fp32 SiTU gradients for one chunk of rows.

Parameters:

Returns: torch.Tensor

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

Compute the fp32 SiTU chain for one chunk of rows.

Parameters:

Returns: torch.Tensor

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

Return True when routing_weights carries one entry per gate_up row.

Parameters:

Return True when the Triton SiTU kernels handle this [rows, 2 * intermediate] call.

Parameters:

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

Whole-tensor weighted-SiTU backward for the compiled path.

Parameters:

Returns: tuple[torch.Tensor, torch.Tensor | None]

(d_gate_up2, d_rw) in the inputs’ dtypes; d_rw is None when not wanted.

Whole-tensor weighted-SiTU forward for the compiled path.

Low-precision in, low-precision out, fp32 math inside; under torch.compile the casts fuse into the elementwise kernel, so no fp32 intermediate exists and the eager row-chunk loop is unnecessary.

Parameters:

Returns: torch.Tensor

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

Apply the dense SiTU activation through the Triton kernels or the (possibly compiled) core.