> 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.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`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-_attn_res_bwd_coef_kernel) | One token per program: \ per entry, then the softmax-backward coefficients. | | [`_attn_res_bwd_grad_kernel`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-_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`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-_attn_res_fwd_kernel) | One token per program: scores, softmax and weighted sum over NB block entries + the prefix row. | | [`_check_inputs`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-_check_inputs) | - | | [`_chunk_configs`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-_chunk_configs) | Autotune candidates: chunk width x warps (register footprint vs loads in flight). | | [`_entries_pad`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-_entries_pad) | - | | [`attn_res_bwd_triton`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-attn_res_bwd_triton) | Fused attention-residual mix backward. | | [`attn_res_fwd_triton`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-attn_res_fwd_triton) | Fused attention-residual mix forward. | ### Data [`BWD_ROWS`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-BWD_ROWS) [`HAVE_TRITON`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-HAVE_TRITON) [`MAX_ENTRIES`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-MAX_ENTRIES) [`MAX_HIDDEN`](#nemo_automodel-components-models-kimi_k3-attn_res_triton-MAX_HIDDEN) ### API ```python 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: \ 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 = \: dp\_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. ```python 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). ```python 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. ```python 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 ) ``` ```python nemo_automodel.components.models.kimi_k3.attn_res_triton._chunk_configs() -> list ``` Autotune candidates: chunk width x warps (register footprint vs loads in flight). ```python nemo_automodel.components.models.kimi_k3.attn_res_triton._entries_pad( entries: int ) -> int ``` ```python 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 ```python 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 ```python nemo_automodel.components.models.kimi_k3.attn_res_triton.BWD_ROWS = 8 ``` ```python nemo_automodel.components.models.kimi_k3.attn_res_triton.HAVE_TRITON = True ``` ```python nemo_automodel.components.models.kimi_k3.attn_res_triton.MAX_ENTRIES = 16 ``` ```python nemo_automodel.components.models.kimi_k3.attn_res_triton.MAX_HIDDEN = 16384 ```