nemo_automodel.components.models.kimi_k3.attn_res_triton

Hand-written Triton kernels for the Kimi-K3 attention-residual mix.

The mix (situ._attn_res_core) scores every entry of [tokens, k+1, hidden] (the k block residuals plus the current prefix sum) with an RMS-normalised dot product against norm_weight * proj_weight, softmaxes the k+1 scores per token and returns the probability-weighted sum of the entries. The eager chain materialises torch.cat((block_residual, prefix_sum)) and an fp32 upcast of it, and the torch.compile version (BackendConfig.compile_situ) lowers the row-broadcast multiply and the entry-axis reduction into 1-D pointwise / reduction kernels that run far below HBM bandwidth (on a 256-GPU K3 profile the two named kernels cost 761 us and 390 us per call at stage 0, where k is at most 1; k grows to 8 on the last pipeline stage).

These kernels handle one token per program and stream the hidden row in C-wide chunks (autotuned 512-4096 columns x 4-8 warps), so the register footprint stays small and many programs share an SM. The forward streams every entry once for the two row reductions (sum of squares, dot with the score weight) with elementwise accumulators reduced once per entry, softmaxes the scores in registers and streams the entries a second time for the weighted sum; no concatenation, no fp32 copy, bf16 in / bf16 out with fp32 math. The per-token probabilities, inverse RMS and dot products (3 x (k+1) fp32 values) are saved for the backward, which runs as two kernels: a per-token coefficient kernel (<grad, entry> reductions and the softmax backward) and a [tokens / BWD_ROWS, chunks] gradient kernel that writes the analytic input gradients and accumulates the score-weight gradient per (row block, chunk) into a partial buffer that is summed in torch (deterministic, no atomics).

The fp32 math matches _attn_res_core up to fp32 accumulation order (the reference multiplies each element by the inverse RMS before its dot-product reduction; the kernel reduces first and scales the sum), which is below bf16 resolution for the shapes this is used at.

Module Contents

Functions

Name	Description
`_attn_res_bwd_coef_kernel`	One token per program: <g, x_j> per entry, then the softmax-backward coefficients.
`_attn_res_bwd_grad_kernel`	Program (row block, chunk): dx_j = p_j g + c_sw_j sw + c_x_j x_j for R tokens on one C-wide chunk.
`_attn_res_fwd_kernel`	One token per program: scores, softmax and weighted sum over NB block entries + the prefix row.
`_check_inputs`	-
`_chunk_configs`	Autotune candidates: chunk width x warps (register footprint vs loads in flight).
`_entries_pad`	-
`attn_res_bwd_triton`	Fused attention-residual mix backward.
`attn_res_fwd_triton`	Fused attention-residual mix forward.

Data

API

nemo_automodel.components.models.kimi_k3.attn_res_triton._attn_res_bwd_coef_kernel(
    br_ptr,
    ps_ptr,
    g_ptr,
    p_ptr,
    r_ptr,
    a_ptr,
    csw_ptr,
    cx_ptr,
    T,
    s_br_t,
    s_br_k,
    s_ps_t,
    s_g_t,
    s_st_t,
    H: triton.language.constexpr,
    NB: triton.language.constexpr,
    NB_PAD: triton.language.constexpr,
    C: triton.language.constexpr
)

One token per program: <g, x_j> per entry, then the softmax-backward coefficients.

With p = softmax(s), s_j = r_j * a_j, r_j = rsqrt(mean(x_j^2) + eps), a_j = <x_j, sw>: dp_j = <g, x_j>, ds_j = p_j * (dp_j - sum_i p_i dp_i). Writes c_sw_j = ds_j * r_j (the coefficient on the score weight) and c_x_j = -ds_j * a_j * r_j^3 / H (the coefficient on x_j), so that dx_j = p_j * g + c_sw_j * sw + c_x_j * x_j.

nemo_automodel.components.models.kimi_k3.attn_res_triton._attn_res_bwd_grad_kernel(
    br_ptr,
    ps_ptr,
    nw_ptr,
    pw_ptr,
    g_ptr,
    p_ptr,
    csw_ptr,
    cx_ptr,
    dbr_ptr,
    dps_ptr,
    dsw_ptr,
    T,
    s_br_t,
    s_br_k,
    s_ps_t,
    s_g_t,
    s_st_t,
    s_dbr_t,
    s_dbr_k,
    s_dps_t,
    H: triton.language.constexpr,
    NB: triton.language.constexpr,
    NB_PAD: triton.language.constexpr,
    R: triton.language.constexpr,
    C: triton.language.constexpr
)

Program (row block, chunk): dx_j = p_j g + c_sw_j sw + c_x_j x_j for R tokens on one C-wide chunk.

Accumulates the score-weight gradient of its chunk over its R tokens and writes it to dsw[row block, chunk] (deterministic partials, summed by the caller).

nemo_automodel.components.models.kimi_k3.attn_res_triton._attn_res_fwd_kernel(
    br_ptr,
    ps_ptr,
    nw_ptr,
    pw_ptr,
    out_ptr,
    p_ptr,
    r_ptr,
    a_ptr,
    T,
    eps,
    s_br_t,
    s_br_k,
    s_ps_t,
    s_out_t,
    s_st_t,
    H: triton.language.constexpr,
    NB: triton.language.constexpr,
    NB_PAD: triton.language.constexpr,
    C: triton.language.constexpr
)

One token per program: scores, softmax and weighted sum over NB block entries + the prefix row.

nemo_automodel.components.models.kimi_k3.attn_res_triton._check_inputs(
    prefix_sum: torch.Tensor,
    block_residual: torch.Tensor,
    norm_weight: torch.Tensor,
    proj_weight: torch.Tensor
)

nemo_automodel.components.models.kimi_k3.attn_res_triton._chunk_configs() -> list

Autotune candidates: chunk width x warps (register footprint vs loads in flight).

nemo_automodel.components.models.kimi_k3.attn_res_triton._entries_pad(
    entries: int
) -> int

nemo_automodel.components.models.kimi_k3.attn_res_triton.attn_res_bwd_triton(
    prefix_sum: torch.Tensor,
    block_residual: torch.Tensor,
    norm_weight: torch.Tensor,
    proj_weight: torch.Tensor,
    grad_out: torch.Tensor,
    stats: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Fused attention-residual mix backward.

Parameters:

prefix_sum, block_residual, norm_weight, proj_weight

the forward inputs.

grad_out

torch.Tensor

[tokens, hidden] gradient of the mixed output (contiguous last dim).

stats

torch.Tensor

the forward’s [3, tokens, k+1] fp32 statistics.

Returns: torch.Tensor

(d_prefix_sum, d_block_residual, d_score_weight): the first two in their inputs’ dtypes, the

nemo_automodel.components.models.kimi_k3.attn_res_triton.attn_res_fwd_triton(
    prefix_sum: torch.Tensor,
    block_residual: torch.Tensor,
    norm_weight: torch.Tensor,
    proj_weight: torch.Tensor,
    eps: float
) -> tuple[torch.Tensor, torch.Tensor]

Fused attention-residual mix forward.

Parameters:

prefix_sum

torch.Tensor

[tokens, hidden] current residual stream (any float dtype; sets the output dtype).

block_residual

torch.Tensor

[tokens, k, hidden] prior block starts, k <= MAX_ENTRIES (may be 0).

norm_weight

torch.Tensor

[hidden] RMSNorm weight of the score norm.

proj_weight

torch.Tensor

[hidden] squeezed [1, hidden] projection weight.

eps

float

RMSNorm epsilon.

Returns: torch.Tensor

(mixed [tokens, hidden] in prefix_sum's dtype, stats [3, tokens, k+1] fp32) where the stats

nemo_automodel.components.models.kimi_k3.attn_res_triton.BWD_ROWS = 8

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

nemo_automodel.components.models.kimi_k3.attn_res_triton.MAX_ENTRIES = 16

nemo_automodel.components.models.kimi_k3.attn_res_triton.MAX_HIDDEN = 16384