From f33ce6b3e58b213d8b55800cfa70a76faaab7f00 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:10:07 +0800 Subject: [PATCH 1/8] Add NVFP4 Kimi K3 model and agentx support (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../__init__.py | 3 + .../decode_comm.py | 402 +++ .../decode_gemv.py | 467 +++ .../modeling.py | 3012 +++++++++++++++++ .../spec_worker.py | 1343 ++++++++ .../weights.py | 42 + .../modeling_v2/models/kimi_k3_vl/routing.py | 46 +- .../_torch/models/modeling_kimi_linear.py | 9 +- .../modeling/test_kimi_linear_checkpoint.py | 33 +- .../test_modeling_v2_k3_situ_mul.py | 4 +- ...test_modeling_v2_k3_drafter_attn_qknorm.py | 6 +- .../gemm/test_modeling_v2_k3_ctm_gemv.py | 4 +- .../gemm/test_modeling_v2_k3_ctm_gemv_long.py | 4 +- .../test_modeling_v2_k3_ctm_gemv_swiglu.py | 4 +- .../gemm/test_modeling_v2_k3_ctm_gemv_wide.py | 4 +- .../gemm/test_modeling_v2_k3_decode_gemv.py | 4 +- .../gemm/test_modeling_v2_k3_head_gemv.py | 4 +- .../norm/test_modeling_v2_k3_embed_norm.py | 4 +- .../modeling_v2/test_modeling_v2_routing.py | 205 ++ .../_torch/moe/test_kimi_k3_situ_moe.py | 87 + 20 files changed, 5664 insertions(+), 23 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py new file mode 100644 index 000000000000..c43a02997371 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 (NVFP4) / sm_107 / tp16 attention, routed experts moe_tp 4 x moe_ep 4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py new file mode 100644 index 000000000000..f5327daf9474 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's residual-update collectives on the catalog's Kimi K3 MNNVL and sandwich entries. + +A layer's post-attention step adds the attention output to the running prefix sum, selects the attention residual +and applies the post-attention RMSNorm. Under TP the attention output is a sum over the TP group, so the step is that +all-reduce followed by the residual update. On a step of at most `AR_ATTN_RES_MAX_TOKENS` tokens (wide decode steps +aside) one collective runs both: + +* `K3DecodeComm.sandwich_oproj`, `comm/k3_sandwich_oproj`: o_proj, its all-reduce and the residual update in one + kernel, at most `SANDWICH_MAX_TOKENS` tokens of an o_proj of the TP16 per-rank shape [7168, 768]. The attention + hands over its gated o_proj input. +* `K3DecodeComm.allreduce_attn_res`, `comm/mnnvl_allreduce_attn_res`, everywhere else: the one-shot MNNVL all-reduce + of the attention's unreduced o_proj partial, with the residual update as its epilogue. + +The pre-attention step after a MoE layer whose row-parallel tail was handed on (a `PendingTail`, `decode_moe.py`) is +the same kind of collective: `K3DecodeComm.sandwich_tail`, `comm/k3_sandwich_tail`, runs the tail GEMV, its +all-reduce and the next layer's residual update (the final norm's, after the last layer) in one kernel. + +The DSpark drafter (`K3DSparkDrafter`) runs its collectives with a plain residual add and RMSNorm on the same state: + +* `K3DecodeComm.sandwich_plain`, `comm/k3_sandwich_plain`: a row-parallel projection of at most `SANDWICH_MAX_TOKENS` + tokens (a layer's attention output projection, or its MLP's SiLU-and-mul and down projection), its all-reduce, the + residual add and the RMSNorm in one kernel. +* `K3DecodeComm.allreduce_norm`, `comm/mnnvl_fusion_allreduce`: the all-reduce of an unreduced projection output of + any token count the MNNVL workspace holds, with the residual add and the RMSNorm in its epilogue (the context + projection's with a zero residual). + +The state is one `MnnvlWorkspace` and one `K3SandwichWorkspace` of the TP group (`K3DecodeComm.create`): collective +over the group and eager, built by the target in `post_load_weights` before any CUDA-graph capture. Every rank must +make the same calls on each in the same order. Which call a step takes is decided from its token count and kind and +from the load-time layout alone, which every rank of the group shares. + +The plain MNNVL all-reduces the decode path keeps (the stock modules') send one-shot up to +`DECODE_AR_ONE_SHOT_MAX_BYTES` (`use_decode_one_shot`), except a wide decode step's (`wide_all_reduce`). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_oproj import ( + K3SandwichWorkspace, + k3_sandwich_oproj, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_plain import ( + k3_sandwich_plain, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_tail import ( + k3_sandwich_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allreduce_attn_res import ( + mnnvl_allreduce_attn_res, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_fusion_allreduce import ( + mnnvl_fusion_allreduce, + required_buffer_bytes, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) + +# The kernel's support predicate: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_sandwich import op as _sandwich_op +from tensorrt_llm._torch.distributed import AllReduceParams + +# The most tokens of a step whose post-attention all-reduce runs one-shot with the residual update in its epilogue. +AR_ATTN_RES_MAX_TOKENS = 16 + +# The most tokens the sandwich kernel takes (one token tile). +SANDWICH_MAX_TOKENS = _sandwich_op.MAX_TOKENS + +# The MNNVL workspace's buffer size: a one-shot call pushes T x 7168 bf16 from each of 16 ranks, 3.5 MiB at T = 16. +MNNVL_BUFFER_BYTES = 4 << 20 + +# The one-shot ceiling of the stock MNNVL all-reduces on the decode path: 8 tokens x 7168 x 16 ranks x 2 B is +# 1.75 MiB, which the stock 1 MiB ceiling would send two-shot. +DECODE_AR_ONE_SHOT_MAX_BYTES = 4 << 20 + +# A wide decode step's ceiling: the stock 1 MiB. At 16 ranks two-shot is faster for every wide step's rows. +WIDE_AR_ONE_SHOT_MAX_BYTES = 1 << 20 + + +def use_decode_one_shot(model: nn.Module) -> None: + """Every stock MNNVL all-reduce of ``model`` sends one-shot up to `DECODE_AR_ONE_SHOT_MAX_BYTES`. Each grows its + workspace on its first eager call of a larger size, before the capture of that size.""" + for module in model.modules(): + mnnvl = getattr(module, "mnnvl_allreduce", None) + if mnnvl is not None: + mnnvl.one_shot_max_bytes = DECODE_AR_ONE_SHOT_MAX_BYTES + mnnvl.early_trigger = True + + +def skip_all_reduce() -> AllReduceParams: + """All-reduce parameters under which a row-parallel module returns this rank's unreduced output (its projection + without its all-reduce): a ``Linear``'s ``all_reduce_params``, a ``GatedMLP``'s ``final_all_reduce_params``.""" + return AllReduceParams(enable_allreduce=False) + + +def wide_all_reduce(all_reduce: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``all_reduce(x)`` of a wide decode step (a stock ``AllReduce`` module, no fusion): its MNNVL all-reduce with + the `WIDE_AR_ONE_SHOT_MAX_BYTES` ceiling, else the module itself.""" + mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None: + out = mnnvl( + x.contiguous(), AllReduceParams(), one_shot_max_bytes=WIDE_AR_ONE_SHOT_MAX_BYTES + ) + if out is not None: + return out + return all_reduce(x) + + +class PendingTail(NamedTuple): + """A MoE layer's row-parallel tail left to its consumer's fused pre-attention step (``K3DecodeComm.sandwich_tail``): + the reduced latent ``[T, 3584]``, the shared experts' activation, the tail weight ``[latent up columns | padding | + shared down]``, this rank's first latent column and the latent norm's epsilon.""" + + latent: torch.Tensor + act: torch.Tensor + weight: torch.Tensor + lo: int + lat_eps: float + + +def _eps(norm: nn.Module) -> float: + """The epsilon of a KimiK3RMSNorm or ``torch.nn.RMSNorm`` (``eps``; for torch's None, its default: the machine + epsilon of the weight's dtype) or a stock RMSNorm (``variance_epsilon``).""" + eps = norm.eps if hasattr(norm, "eps") else norm.variance_epsilon + return float(torch.finfo(norm.weight.dtype).eps if eps is None else eps) + + +def _res_args(res_proj: nn.Module, res_norm: nn.Module, out_norm: nn.Module) -> tuple: + """The residual update's weights and epsilons in the entries' order.""" + return ( + res_proj.weight.reshape(-1), + res_norm.weight, + out_norm.weight, + _eps(res_norm), + _eps(out_norm), + ) + + +@dataclass(eq=False) +class K3DecodeComm: + """The decode path's collective state for the TP group: the `MnnvlWorkspace` and the `K3SandwichWorkspace` the + post-attention steps run on. Built by `create`; owned by the target and shared by every layer.""" + + mnnvl: MnnvlWorkspace + sandwich: K3SandwichWorkspace + + @classmethod + def create(cls, mapping, oproj_weight: Optional[torch.Tensor] = None) -> "K3DecodeComm": + """The state for ``mapping``'s TP group. Collective: every rank of the group calls it at the same point, + eagerly, before any CUDA-graph capture; each workspace fails on every rank or on none. + + ``oproj_weight``: an o_proj weight ``takes_oproj`` holds for. The sandwich kernel compiles here for its + shape, with one call on a zero row of a zero weight, so no capture compiles it; the call advances the + sandwich workspace on every rank alike.""" + state = cls( + MnnvlWorkspace.create(mapping, MNNVL_BUFFER_BYTES), + K3SandwichWorkspace.create(mapping), + ) + if oproj_weight is not None: + weight = torch.zeros_like(oproj_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_oproj( + weight.new_zeros(1, weight.shape[1]), + weight, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + state.sandwich, + ) + torch.cuda.synchronize(weight.device) + return state + + def compile_tail(self, latent_size: int, act_size: int, tail_weight: torch.Tensor) -> None: + """Compile the sandwich tail kernel for a MoE tail of ``latent_size`` latent and ``act_size`` shared columns + and ``tail_weight``'s shape, with one call on a zero row of a zero weight, before any capture. Collective: every + rank of the group makes the call; it advances the sandwich workspace on every rank alike.""" + weight = torch.zeros_like(tail_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_tail( + weight.new_zeros(1, latent_size), + weight.new_zeros(1, act_size), + weight, + 0, + 1e-6, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + self.sandwich, + ) + torch.cuda.synchronize(weight.device) + + def takes_post_attention(self, hidden_states: torch.Tensor, step) -> bool: + """Whether one of this state's collectives runs the post-attention step of a layer whose attention input is + ``hidden_states``: at most `AR_ATTN_RES_MAX_TOKENS` bf16 rows of a hidden size the MNNVL entry takes, on any + step but a wide decode step (``step.wide``), whose post-attention update stays the all-reduce and the fused + add + attn_res + RMSNorm.""" + if step is not None and step.wide: + return False + if hidden_states.dim() != 2 or hidden_states.dtype != torch.bfloat16: + return False + rows, hidden = hidden_states.shape + return ( + hidden % 1024 == 0 + and hidden <= 8192 + and 0 < rows <= min(AR_ATTN_RES_MAX_TOKENS, self.mnnvl.max_one_shot_tokens(hidden)) + ) + + @staticmethod + def takes_oproj(o_proj: nn.Module, max_snapshots: int) -> bool: + """Whether ``sandwich_oproj`` takes the layer of output projection ``o_proj`` on a step of at most + `SANDWICH_MAX_TOKENS` tokens, decided once the weights are final: a bias-free, contiguous bf16 weight of the + kernel's shape (the TP16 per-rank [7168, 768]) and a snapshot bank of at most ``max_snapshots`` rows the + kernel's candidate count holds.""" + weight = getattr(o_proj, "weight", None) + return ( + getattr(o_proj, "bias", None) is None + and isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.is_cuda + and _sandwich_op.supports(weight.new_empty((1, weight.shape[1])), weight, max_snapshots) + ) + + @staticmethod + def takes_tail( + latent: torch.Tensor, act: torch.Tensor, tail_weight: torch.Tensor, max_snapshots: int + ) -> bool: + """Whether ``sandwich_tail`` takes a MoE tail of these tensors' shapes (rows of ``latent`` and ``act``, at most + `SANDWICH_MAX_TOKENS`; ``tail_weight``) with a snapshot bank of at most ``max_snapshots`` rows: the TP16 + per-rank shapes.""" + return _sandwich_op.supports_tail(latent, act, tail_weight, max_snapshots) + + def allreduce_attn_res( + self, + partial: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_allreduce_attn_res`` call: ``updated = prefix_sum + + allreduce(partial)`` (the sum alone without ``prefix_sum``), ``normed = out_norm(attn_res(block_residual..., + updated))``, the attention residual selected with ``res_proj`` and ``res_norm``. ``block_residual`` holds the + valid snapshots, ``[S, T, H]``.""" + return mnnvl_allreduce_attn_res( + partial.contiguous(), + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.mnnvl, + ) + + def sandwich_tail( + self, + pending: PendingTail, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + updated_out: Optional[torch.Tensor] = None, + tap: Optional[torch.Tensor] = None, + tap_updated: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of a MoE layer's row-parallel tail ``pending`` + (``[RMSNorm(latent)[:, lo:lo + 224] | act] @ weight.T``), in one ``comm/k3_sandwich_tail`` call. + ``updated_out``: a bf16 ``[T, H]`` tensor the call stores ``updated`` into (the consumer's snapshot bank row), + returned as ``updated``. ``tap``: a bf16 ``[T, H]`` view (a speculative capture slot) the call also stores + the pre-norm attention-residual mixture into, or ``updated`` with ``tap_updated``.""" + return k3_sandwich_tail( + pending.latent, + pending.act, + pending.weight, + pending.lo, + pending.lat_eps, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + tap=tap, + tap_updated=tap_updated, + updated_out=updated_out, + ) + + def sandwich_oproj( + self, + core: torch.Tensor, + o_weight: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of ``core @ o_weight.T``, in one ``comm/k3_sandwich_oproj`` + call: ``core`` is this rank's gated o_proj input ``[T, 768]``, ``o_weight`` its o_proj slice. Bit for bit + o_proj followed by ``allreduce_attn_res``.""" + return k3_sandwich_oproj( + core.contiguous(), + o_weight, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + ) + + def compile_plain(self, weight: torch.Tensor, swiglu: bool = False) -> bool: + """Compile the plain sandwich (``sandwich_plain``) for a projection of ``weight``'s shape (with ``swiglu``, a + down projection after the SiLU-and-mul) with one call on a zero row of a zero weight, before any capture; + False, with nothing launched, where the kernel does not take that shape. Collective: every rank of the group + makes the call; it advances the sandwich workspace on every rank alike.""" + zero = torch.zeros_like(weight) + hidden, k_in = zero.shape + x = zero.new_zeros(1, 2 * k_in if swiglu else k_in) + residual = zero.new_zeros(1, hidden) + ones = zero.new_ones(hidden) + if not _sandwich_op.supports_plain(x, zero, residual, ones, swiglu): + return False + k3_sandwich_plain(x, zero, residual, ones, 1e-6, self.sandwich, swiglu=swiglu) + torch.cuda.synchronize(zero.device) + return True + + @staticmethod + def takes_plain( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> bool: + """Whether ``sandwich_plain`` takes the call: at most `SANDWICH_MAX_TOKENS` contiguous bf16 rows ``x`` of a + row-parallel slice ``weight`` [7168, K] (K a multiple of 128 up to 896; with ``swiglu`` K 896 and ``x`` the + ``[gate | up]`` rows, 2 K wide), ``residual`` [rows, 7168], and ``norm``'s bf16 [7168] weight.""" + return _sandwich_op.supports_plain(x, weight, residual, norm.weight, swiglu) + + def sandwich_plain( + self, + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/k3_sandwich_plain`` call: ``updated = residual + allreduce(x @ + weight.T)`` (with ``swiglu``, of ``silu_and_mul(x) @ weight.T``), ``normed = norm(updated)`` for a plain + RMSNorm ``norm``. Bit for bit, by the kernel's statement, the projection on ``k3_ctm_gemv`` at split 1 (with + ``swiglu``, ``k3_ctm_gemv_swiglu`` at split 2) followed by ``allreduce_norm``'s one-shot call.""" + return k3_sandwich_plain( + x.contiguous(), + weight, + residual.contiguous(), + norm.weight, + _eps(norm), + self.sandwich, + swiglu=swiglu, + ) + + def takes_allreduce_norm(self, rows: int, hidden: int) -> bool: + """Whether ``allreduce_norm`` takes ``rows`` bf16 rows of ``hidden`` columns: the MNNVL workspace holds the + call at the decode path's one-shot ceiling (`DECODE_AR_ONE_SHOT_MAX_BYTES`, two-shot above it).""" + if rows <= 0 or hidden <= 0 or hidden % 8: + return False + world, buffer_bytes = self.mnnvl.world_size, self.mnnvl.buffer_bytes + need = required_buffer_bytes( + rows, hidden, world, torch.bfloat16, DECODE_AR_ONE_SHOT_MAX_BYTES + ) + two_shot = rows * hidden * world * 2 > DECODE_AR_ONE_SHOT_MAX_BYTES + return need <= buffer_bytes and not (two_shot and buffer_bytes % 32) + + def allreduce_norm( + self, partial: torch.Tensor, residual: torch.Tensor, norm: nn.Module + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_fusion_allreduce`` call: ``updated = residual + + allreduce(partial)``, ``normed = norm(updated)`` for a plain RMSNorm ``norm``, sent one-shot up to + `DECODE_AR_ONE_SHOT_MAX_BYTES`. ``partial`` is this rank's unreduced ``[rows, hidden]`` bf16 output of a + row-parallel projection, of a shape ``takes_allreduce_norm`` holds for.""" + return mnnvl_fusion_allreduce( + partial.contiguous(), + self.mnnvl, + DECODE_AR_ONE_SHOT_MAX_BYTES, + residual.contiguous(), + norm.weight, + _eps(norm), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py new file mode 100644 index 000000000000..0729fe05b91e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py @@ -0,0 +1,467 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. + +* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest + at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, + `gemm/k3_ctm_gemv_wide`, `gemm/k3_ctm_gemv_long`, `gemm/k3_ctm_gemv` or `gemm/k3_ctm_gemv_swiglu`, and, where the + site lists it, more rows (up to `WIDE_ROWS`) on `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused + projections (MLA's [W_a; W_g] with the gate rows through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention + output projection, the built-in MLA path's q_a / kv_a, q_b and gate projections, the MoE decode path's projections + (`decode_moe.py`), two of them with fp32 outputs, and the DSpark drafter's q / k / v, output, gate / up and down + projections (`K3DSparkDrafter`). +* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on + `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the + stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the + drafter's logits on the same head take it too. +* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual + bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. +* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: + gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then + runs the down projection's all-reduce. + +Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. + +A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. +`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The +embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call +whose kernel has not run eagerly is refused. + +The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one +stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, Set + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( + k3_ctm_gemv_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( + K3HeadGemvWorkspace, + k3_head_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.concat import concat +from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.split import split + +# The kernels' support predicates: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._utils import mpi_disabled + +# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 +# requests of 8 tokens). +MAX_ROWS = 8 +WIDE_ROWS = 64 + + +@dataclass(frozen=True) +class Site: + """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows + ("decode", "wide", "long", "ctm" for k3_ctm_gemv, or "swiglu" for k3_ctm_gemv_swiglu, whose rows are the + ``[gate | up]`` input of the SiLU-and-mul, 2 ``k`` wide), and k3_ctm_gemv_wide at MAX_ROWS+1..``wide_rows`` rows + where ``wide``. Output columns from ``sig_col0`` on are stored through a sigmoid; ``out_fp32``: an fp32 output + (k3_ctm_gemv_wide only). ``split`` / ``ring`` / ``push``: the CTM kernels' CTAs per 128-row weight tile, + k3_ctm_gemv_long's weight-ring stages, and whether the partial sums are pushed to each row's owner.""" + + n: int + k: int + small: str + wide: bool = False + sig_col0: int = -1 + split: int = 0 + ring: int = 0 + push: bool = False + out_fp32: bool = False + wide_rows: int = WIDE_ROWS + + +SITES: Dict[str, Site] = { + # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), + # the gate rows through a sigmoid. + "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), + # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). + "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), + # The attention output projection (row parallel; 6 heads x 128 in). + "o_proj": Site(7168, 768, "decode", wide=True), + # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. + "kv_a": Site(2112, 7168, "decode"), + "q_b": Site(1152, 1536, "wide"), + "g_proj": Site(768, 7168, "wide"), + # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. + "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), + "dense_down": Site(7168, 2112, "long", split=2, ring=6), + # The MoE decode path (decode_moe.py): this rank's head slice [latent down 224 | router 56] with an fp32 output, + # the shared experts' gate_up, the row-parallel tail [latent up 224 | padding 32 | shared down 384] and the + # replicated tail's latent up projection with an fp32 output. The wide kernel takes the head slice and the tail up + # to 32 rows, where it is faster than the stock GEMM. + "moe_head": Site(280, 7168, "wide", wide=True, out_fp32=True, wide_rows=32), + "moe_shared_gate_up": Site(768, 7168, "wide", wide=True), + "moe_tail": Site(7168, 640, "wide", wide=True, wide_rows=32), + "moe_up": Site(7168, 3584, "wide", out_fp32=True), + # The DSpark drafter's projections (K3DSparkDrafter; 6 query heads, 1 KV head of 64, MLP width 896): the fused + # q / k / v, the output projection (row parallel), gate_up [gate 896 | up 896], and the down projection with the + # SiLU-and-mul (row parallel). + "drafter_qkv": Site(512, 7168, "long", split=8, ring=6, push=True), + "drafter_o": Site(7168, 384, "ctm", split=1, push=True), + "drafter_gate_up": Site(1792, 7168, "long", split=8, ring=6, push=True), + "drafter_down": Site(7168, 896, "swiglu", split=2, push=True), +} + + +def _width(spec: Site) -> int: + """The row width a site's kernel reads: ``k``, or the ``[gate | up]`` input of a SiLU-and-mul site, 2 ``k``.""" + return 2 * spec.k if spec.small == "swiglu" else spec.k + + +def _wide_tile(rows: int) -> int: + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel + + return k3_ctm_gemv_kernel.wide_tile(rows) + + +def _run( + spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor +) -> Optional[torch.Tensor]: + """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" + if kernel == "decode": + if spec.sig_col0 >= 0 or spec.out_fp32 or not _decode_op.supports(x2d, weight): + return None + return k3_decode_gemv(x2d, weight) + if kernel == "wide": + if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, spec.out_fp32): + return None + return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0, out_fp32=spec.out_fp32) + if spec.out_fp32: + return None + if kernel == "ctm": + if spec.sig_col0 >= 0 or not _ctm_op.supports(x2d, weight, spec.split): + return None + return k3_ctm_gemv(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + if kernel == "swiglu": + if spec.sig_col0 >= 0 or not _ctm_op.supports_swiglu(x2d, weight, spec.split): + return None + return k3_ctm_gemv_swiglu(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. + sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count + if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( + x2d, weight, spec.split, spec.ring + ): + return None + return k3_ctm_gemv_long( + x2d, + weight, + sig_col0=spec.sig_col0, + split=spec.split, + ring=spec.ring, + trigger_early=True, + push=spec.push, + ) + + +def _capturing() -> bool: + return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() + + +def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: + """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" + if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: + return x2d.clone(memory_format=torch.contiguous_format) + return x2d + + +def _head_takes_module(lm_head: nn.Module) -> bool: + """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along + the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" + weight = getattr(lm_head, "weight", None) + mapping = getattr(lm_head, "mapping", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and weight.is_contiguous() + and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" + and getattr(lm_head, "gather_output", False) + and getattr(lm_head, "gather_output_sizes", None) is None + and getattr(lm_head, "padding_size", None) == 0 + and getattr(lm_head, "bias", None) is None + and not getattr(lm_head, "has_any_quant", True) + and mapping is not None + and not mapping.enable_attention_dp + ) + + +def _plain_rmsnorm(norm: nn.Module) -> bool: + """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for + bit. An unknown module fails closed.""" + weight = getattr(norm, "weight", None) + return ( + IS_FLASHINFER_AVAILABLE + and isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "use_cuda_tile", True) + and not getattr(norm, "return_hp_output", True) + and getattr(norm, "nvfp4_scale", None) is None + and hasattr(norm, "variance_epsilon") + ) + + +class K3DecodeGemvs: + """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran + eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are + final; owned by the target.""" + + def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: + self.head_workspace = head_workspace + self._ran: Set[tuple] = set() + + @classmethod + def create( + cls, + lm_head: Optional[nn.Module] = None, + sites: Iterable[str] = tuple(SITES), + device: Optional[torch.device] = None, + ) -> "K3DecodeGemvs": + """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: + it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on + zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or + head whose kernel does not take its shape keeps the generic path.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" + ) + head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + if device is None: + device = ( + head_weight.device + if isinstance(head_weight, torch.Tensor) + else torch.device("cuda", torch.cuda.current_device()) + ) + state = cls() + for site in sites: + spec = SITES[site] + weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) + for rows in (1, 16, 32, 64) if spec.wide else (1,): + if rows > spec.wide_rows: + continue + state._project(site, weight.new_zeros(rows, _width(spec)), weight, warm=True) + del weight + if "dense_gate_up" in sites: + gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) + if _ctm_op.supports_situ_mul(gu): + for linear_beta in (None, 1.0): + k3_situ_mul(gu, 1.0, linear_beta) + state._ran.add(("situ_mul", linear_beta is not None)) + if lm_head is not None and _head_takes_module(lm_head): + x = head_weight.new_zeros(1, head_weight.shape[1]) + if _head_op.supports(x, head_weight): + workspace = K3HeadGemvWorkspace.create( + head_weight.shape[0], head_weight.shape[1], head_weight.device + ) + k3_head_gemv(x, head_weight, workspace) + state.head_workspace = workspace + state._ran.add(("lm_head",)) + torch.cuda.synchronize(device) + return state + + def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: + """``x @ weight.T`` (``[..., N]``: bf16, the site's sigmoid columns through the sigmoid; fp32 at an + ``out_fp32`` site; on a SiLU-and-mul site, ``silu(gate) * up`` of ``x = [gate | up]`` first) for ``site``'s + weight on its decode kernel, or None where none takes the call: more rows than the site's kernels take, + another shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its + module.""" + return self._project(site, x, weight, warm=False) + + def _project( + self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool + ) -> Optional[torch.Tensor]: + spec = SITES[site] + width = _width(spec) + if ( + weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (spec.n, spec.k) + or not weight.is_contiguous() + or x.dtype != torch.bfloat16 + or x.dim() < 1 + or x.shape[-1] != width + ): + return None + rows = x.numel() // width + if 0 < rows <= MAX_ROWS: + kernel = spec.small + elif spec.wide and MAX_ROWS < rows <= spec.wide_rows: + kernel = "wide" + else: + return None + key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) + capturing = _capturing() + if capturing and not warm and key not in self._ran: + return None + y = _run(spec, kernel, _dense_rows(x.reshape(rows, width)), weight) + if y is None: + return None + if not capturing: + self._ran.add(key) + return y.view(*x.shape[:-1], spec.n) + + def lm_head_logits( + self, rows: torch.Tensor, lm_head: nn.Module, gather: bool = True + ) -> Optional[torch.Tensor]: + """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on + ``gemm/k3_head_gemv`` (``gather`` False: the shard alone, ``[M, vocab / TP]``); None where it does not take + the call (more than `MAX_ROWS` rows, another head, ...).""" + workspace = self.head_workspace + if workspace is None or not _head_takes_module(lm_head): + return None + weight = lm_head.weight + if ( + tuple(weight.shape) != (workspace.n_out, workspace.k_in) + or weight.device != workspace.partials.device + or rows.dim() != 2 + or rows.dtype != torch.bfloat16 + or not 0 < rows.shape[0] <= MAX_ROWS + or rows.shape[1] != workspace.k_in + ): + return None + if _capturing() and ("lm_head",) not in self._ran: + return None + group = lm_head.mapping.tp_group + if gather and len(group) > 1 and mpi_disabled(): + return None + x = _dense_rows(rows) + if not _head_op.supports(x, weight): + return None + local = k3_head_gemv(x, weight, workspace) + if not gather or len(group) == 1: + return local + gathered = allgather(local, None, group) + return concat(list(split(gathered, rows.shape[0], dim=0)), dim=-1) + + def dense_mlp( + self, + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + beta: float, + linear_beta: Optional[float], + ) -> Optional[torch.Tensor]: + """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on + ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on + ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then + runs the module.""" + gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] + rows = x.shape[0] if x.dim() == 2 else 0 + if ( + not 0 < rows <= MAX_ROWS + or linear_beta == 0.0 + or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) + or tuple(down_weight.shape) != (down.n, down.k) + or gate_up.n != 2 * down.k + ): + return None + if ( + _capturing() + and not { + ("dense_gate_up", "long", 0), + ("dense_down", "long", 0), + ("situ_mul", linear_beta is not None), + } + <= self._ran + ): + return None + gu = self.project("dense_gate_up", x, gate_up_weight) + if gu is None or not _ctm_op.supports_situ_mul(gu): + return None + return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) + + def embed_norm( + self, + input_ids: torch.Tensor, + table: torch.Tensor, + norm: nn.Module, + bank: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one + ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not + apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose + kernel has not run eagerly).""" + if not _plain_rmsnorm(norm): + return None + ids = input_ids.reshape(-1) + if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): + return None + key = ("embed_norm", ids.numel(), ids.dtype) + capturing = _capturing() + if capturing and key not in self._ran: + return None + normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) + if not capturing: + self._ran.add(key) + return normed + + +class _K3Head: + """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" + + __slots__ = ("_gemvs", "_lm_head") + + def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: + self._gemvs = gemvs + self._lm_head = lm_head + + def __call__(self, rows: torch.Tensor) -> torch.Tensor: + logits = self._gemvs.lm_head_logits(rows, self._lm_head) + return self._lm_head(rows) if logits is None else logits + + +class K3LogitsProcessor(nn.Module): + """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. + + It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock + processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. + """ + + def __init__(self, stock: nn.Module) -> None: + super().__init__() + self.stock = stock + self.gemvs: Optional[K3DecodeGemvs] = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) + return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) + + def lm_head_shard(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: + """This rank's bf16 vocabulary shard of ``lm_head(rows)`` on ``gemm/k3_head_gemv``, without the gather, or + None where ``gemvs`` does not take the rows. A speculative worker that keeps its logits vocabulary-sharded + computes them here, so they hold the values of the logits this processor gathers.""" + if self.gemvs is None: + return None + return self.gemvs.lm_head_logits(rows, lm_head, gather=False) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py new file mode 100644 index 000000000000..883f1a0ea418 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py @@ -0,0 +1,3012 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 on SM107, TP16 attention and MoE TP4 x EP4. + +The target reads the ModelOpt mixed-precision checkpoint without converting its +NVFP4 experts or FP8 attention weights to another format. Quantized attention and +routed experts use their stock implementations. The BF16 embedding, dense MLP, +LM head, attention-residual collectives and standalone GQA DSpark drafter retain +the PR's decode specializations. MLA DSpark uses its stock drafter and worker. +This target has no SM107 accuracy or performance certification until its +corresponding GPU tests and model gates have run. + +Unlike the MXFP4 target, this path permits the stock attention implementation's +quantized KV cache and BF16 recurrent states. It does not allocate the MXFP4 +fused MoE state or the BF16 fused attention projection workspaces. +""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( + k3_drafter_attn_qknorm, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata + +# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. +from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.models.modeling_dspark import ( + GQADSparkForCausalLM, + draft_is_embedded_in_target, +) +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.moe.fused_moe import ( + ConfigurableMoE, + SiTuActivation, + TRTLLMGenFusedMoE, + create_moe, +) +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.config_utils import is_mla +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from . import decode_comm as _decode_comm +from . import decode_gemv as _decode_gemv +from . import spec_worker as _spec_worker +from . import weights as _weights + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, +# and the certification is per arch. +_SM = (10, 7) + +#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in +#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router +#: and fused-A GEMMs), then the decode kernels'. +REQUIRED_TRTLLM_OPS = ( + "attn_res_fwd", + "attn_res_rmsnorm_fwd", + "attn_res_add_rmsnorm_fwd", + "attn_res_add_rmsnorm_persistent_fwd", + "kda_prefill", + "kda_decode", + "kda_mtp_decode", + "dsv3_router_gemm_op", + "dsv3_fused_a_gemm_op", + # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). + "k3_decode_gemv", + "k3_ctm_gemv_wide", + "k3_ctm_gemv_long", + "k3_situ_mul", + "k3_head_gemv", + "k3_embed_norm", + "allgather", + # The decode path's collectives (decode_comm.py). + "mnnvl_allreduce_attn_res", + "k3_sandwich_oproj", + "k3_sandwich_tail", + "mnnvl_allgather_split", + # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces + # with the residual add and RMSNorm (the split context projection's with hidden_norm). + "k3_ctm_gemv", + "k3_ctm_gemv_swiglu", + "k3_drafter_attn_qknorm", + "k3_sandwich_plain", + "mnnvl_fusion_allreduce", +) + +#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. +#: A renamed field upstream fails here, loudly, instead of reading as a default. +REQUIRED_ENGINE_FIELDS = { + "attn_metadata": ( + "num_contexts", + "num_generations", + "num_seqs", + "num_tokens", + "seq_lens", + "tokens_per_block", + "kv_cache_manager", + "mamba_metadata", + ), + "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), +} + +#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every +#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry +#: leaves the list when a catalog entry replaces it. +UNCERTIFIED_GENERIC_CALLS = ( + # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. + "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", + "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", + "tensorrt_llm._torch.models.modeling_utils.DecoderModel", + # The text model's stock modules. + "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", + "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", + "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", + "tensorrt_llm._torch.moe.fused_moe.create_moe", + "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", + "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", + "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", + "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", + "tensorrt_llm._torch.modules.situ.SituAndMul", + "tensorrt_llm._torch.modules.rms_norm.RMSNorm", + "tensorrt_llm._torch.distributed.AllReduce", + "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", + # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries + # do not take a block, its context projection without the TP group's collective state, its context k / v, its + # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. + "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", + "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", + "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", + "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", +) + +# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware +# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). +DECODE_MAX_TOKENS = 8 +MAX_REQUESTS = 8 +MAX_TOKENS_PER_REQUEST = 8 + +_LANG_PREFIX = "language_model." + + +# ---------------------------------------------------------------------------------------------------------------------- +# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and +# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, +# `K3DecodeMLA` below). +# ---------------------------------------------------------------------------------------------------------------------- + +# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections +# instead of the min-latency fused GEMM op (read once at import). +_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( + os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" +) + + +# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table +# holds cos=1/sin=0), but the chunked-context path indexes the table by +# absolute position, so it must cover max_position_embeddings (~512MB per +# backend for the 1M-position checkpoint); a smaller table is read out of +# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context +# deployments. +_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" + + +class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): + """Compute expert choices separately from the quantized MoE GEMMs. + + Preserve the PR's routing arithmetic while the backend applies the + checkpoint's NVFP4 activation and weight scales. + """ + + @property + def requires_separated_routing(self) -> bool: + return True + + +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + @property + def routing_method(self) -> KimiK3MoeRoutingMethod: + """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return KimiK3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + +def _get_text_config(pretrained_config: "PretrainedConfig"): + """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" + if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( + not hasattr(pretrained_config, "linear_attn_config") + and hasattr(pretrained_config, "text_config") + ): + return pretrained_config.text_config + return pretrained_config + + +def _is_kda_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] + + +def _is_mla_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" + + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" + + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" + + +KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" + + +_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") + + +# Turning the port on should not require knowing which of the three topologies +# is the right one: "1" selects the measured policy (``split``), so enabling the +# feature and choosing the policy are one action through one variable. +_ATTN_RES_TOPOLOGY_ON = "1" + + +def _read_attn_res_topology() -> str: + """Default ``per_token``: the persistent kernel is opt-in. + + ``1`` is the only accepted on-value and resolves to ``split``. The named + topologies stay available for measurement: ``persistent`` uses the + persistent kernel at every shape it implements, ``per_token`` at none. + """ + raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") + if raw == _ATTN_RES_TOPOLOGY_ON: + return "split" + if raw not in _ATTN_RES_TOPOLOGIES: + # Loudly, for the same reason as the token ceiling below: a mistyped A/B + # arm that silently fell back to the default would measure one side + # twice and report no difference. + raise ValueError( + f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " + f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " + f"(which means 'split'), got {raw!r}" + ) + return raw + + +def _read_fused_attn_res_max_tokens() -> int: + """Resolved after the topology, because its default follows it. + + With the persistent port off -- the default -- the ceiling is 1, the + pre-existing gate: the fused epilogue is taken at the single-token decode + shape and nowhere else. Enabling the port raises it to 32, the top of the + measured range, so that "off" keeps meaning "unchanged". + """ + default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" + raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) + try: + value = int(raw) + except ValueError: + # Failing loudly matters more than usual here: a mistyped A/B arm that + # silently fell back to the default would measure the candidate twice + # and report no difference. + raise ValueError( + f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value < 1: + raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") + return value + + +_ATTN_RES_TOPOLOGY = _read_attn_res_topology() + + +_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() + + +def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: + """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. + + No token ceiling: the persistent grid is sized by the SM count, not by the + token count, so prefill is the case it exists for. + """ + del M # deliberately unused; see above + return H == 7168 and 2 <= N <= 9 + + +def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: + """Pick between the two fused kernels for this call site. + + ``persistent`` takes the persistent kernel at every shape it implements; + ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in + for the prefill/decode boundary. Shapes the persistent kernel does not + implement fall through to the caller's existing gate and land on the + unfused path. + """ + if not _persistent_attn_res_applicable(M, H, N): + return False + if _ATTN_RES_TOPOLOGY == "persistent": + return True + return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> Optional[torch.Tensor]: + """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. + + Returns ``None`` when the call falls outside the fused kernel's + contract (dtype/shape/arch) so the caller can use the exact fp32 reference + instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` + layout. Candidate order matches the reference: snapshots first, the + running prefix sum last. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + try: + attn_res_op = torch.ops.trtllm.attn_res_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output, _rsigma, _probs, _logits = attn_res_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + ) + return output.reshape(M, H) + + +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that + the shape gate let the call through, and a rejected call looks exactly like + a disabled one in the logs. Emitted at debug level: it is once per distinct + shape, not once per process, so it is a diagnostic rather than a summary. + """ + logger.debug_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, + default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, + # which is the measured window; larger token counts have not been measured + # and fall back to the unfused add + attn_res_fwd + RMSNorm path. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + _, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + None, + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) + return output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in + ``_apply_attn_res_rmsnorm_fused``). + + The production residual add produces a BF16 tensor that remains live across + the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual output, + while avoiding a separate add launch and a re-read of the intermediate by + attention-residual selection. + """ + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # Same measured window as _apply_attn_res_rmsnorm_fused above. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + updated_prefix_sum, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + addend.reshape(M, 1, H).contiguous(), + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = attn_res_add_rmsnorm_op( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + +def _apply_attn_res( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> torch.Tensor: + """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). + + prefix_sum: ``[num_tokens, hidden_size]`` + block_residual: ``[num_snapshots, num_tokens, hidden_size]`` + + Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's + contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. + Only the fallback boundary restores the HF ``[M, K, H]`` layout. + """ + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + block_residual_hf = block_residual.transpose(0, 1) + v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.eps) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the + fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_rmsnorm_fused( + prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in + ``_apply_attn_res_and_rmsnorm``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + + +# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where +# they are faster than the add -> attn_res -> RMSNorm chain. +_WIDE_ATTN_RES_MAX_TOKENS = 32 + + +def _attn_res_max_tokens(step: DecodeStep) -> int: + """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step + ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" + return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS + + +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + + +# The routed experts' module path prefixes. ``exclude_modules`` matches with +# wildcards and walks ancestor prefixes, so an empty prefix would widen what +# matches instead of just missing. +_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") + + +def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_streaming_nvfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + **{ + f"{w}_{kind}": get_tensor(f"{base}.{expert_idx}.{w}.{kind}") + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.NVFP4: _K3ExpertCkptSpec( + kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"), + loader=_load_nvfp4_expert, + loaded_slots_attr="_streamed_expert_slots", + needs_layer_finalize=True, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + self.num_experts = cfg.num_experts + self.top_k = cfg.num_experts_per_token + self.moe_hidden_size = cfg.routed_expert_hidden_size + # ValueError (not assert): these guard unsupported checkpoint + # configurations and must stay active under ``python -O``. + if self.moe_hidden_size is None: + raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") + if not getattr(cfg, "latent_moe_use_norm", False): + raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") + + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) + self._situ_betas = (situ_beta, situ_linear_beta) + dtype = torch.bfloat16 + + # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 + # accumulate/output (checkpoint stores the gate weight in bf16; saves a + # per-layer input cast + fp32 splitK pair on the bs1 decode path). + # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True + self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) + + routed_moe_model_config = self._routed_moe_model_config(model_config) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) + routed_moe_kwargs = dict( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_intermediate_size, + dtype=dtype, + # Kimi owns the latent reduction so it can order that collective + # after the shared expert's auxiliary-stream reduction. + reduce_results=False, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + # Let CommunicationFactory select the best available strategy. + communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. As measured once: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver + # instead, which reports the rejection trail. + # + # CUTLASS is absent on purpose: it is the fallback target, so + # "degraded to CUTLASS" is not a thing that can happen to it. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + ) + self._check_trtllm_situ_quant( + routed_moe_model_config.moe_backend, routed_quant_config.quant_algo + ) + + self.routed_experts = create_moe(**routed_moe_kwargs) + if not isinstance(self.routed_experts, ConfigurableMoE): + raise RuntimeError( + "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." + ) + if self.routed_experts.layer_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "dynamic EPLB or replicated expert slots." + ) + local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) + if local_expert_ids != list( + range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) + ): + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming currently requires a " + "contiguous static expert partition." + ) + self.local_expert_ids = tuple(local_expert_ids) + self.experts_per_rank = len(local_expert_ids) + self.expert_lo = local_expert_ids[0] + self.expert_hi = self.expert_lo + self.experts_per_rank + + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = model_config.mapping.tp_size > 1 + self._reduce_routed_output = ( + use_shared_tp + and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + ) + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self.shared_experts = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=shared_model_config, + reduce_output=use_shared_tp, + layer_idx=layer_idx, + is_shared_expert=True, + ) + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.moe_main_event = torch.cuda.Event() + self.moe_shared_event = torch.cuda.Event() + self.routed_expert_down_proj = nn.Linear( + cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype + ) + # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer + # fallback is the same fp32-variance eager math as KimiK3RMSNorm). + self.routed_expert_norm = RMSNorm( + hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. + self.decode_moe = None + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, projection.weight.t(), None, None + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit + ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (4 x 4, asserted at + construction).""" + return mapping.moe_tp_size, mapping.moe_ep_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Read the NVFP4 expert format, preserving all per-layer scale metadata.""" + global_quant = model_config.quant_config + if global_quant is None: + raise ValueError("The NVFP4 target requires checkpoint quantization metadata") + if any( + global_quant.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + raise ValueError(f"NVFP4 experts at layer {layer_idx} are excluded from quantization") + declarations = model_config.quant_config_dict or {} + matches = [ + declarations[name] + for prefix in ("language_model.model.", "model.", "") + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + if (name := f"{prefix}layers.{layer_idx}.{suffix}") in declarations + ] + if matches: + if any( + quant.quant_algo != matches[0].quant_algo + or quant.group_size != matches[0].group_size + for quant in matches[1:] + ): + raise ValueError(f"Conflicting expert quantization aliases at layer {layer_idx}") + return matches[0] + if global_quant.quant_algo == QuantAlgo.NVFP4: + return global_quant + raise ValueError(f"Missing NVFP4 expert declaration at layer {layer_idx}") + + @staticmethod + def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: + """Reject formats without a fused SiTU cubin in the resolved TRTLLM backend. + + Read the supported formats from the backend so model construction and + kernel selection agree, including when AUTO resolves to TRTLLM. + """ + situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() + if moe_backend != "TRTLLM" or quant_algo in situ_supported: + return + supported = ", ".join(sorted(algo.name for algo in situ_supported)) + raise ValueError( + f"Kimi K3 routed experts are quantized as {quant_algo}, which the " + "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " + f"exist only for {supported}. Set moe_config.backend to CUTLASS " + "or MEGAMOE_CUTEDSL." + ) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config, with the split of ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. + supported_backends = { + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + f"got {model_config.moe_backend!r}." + ) + if model_config.moe_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "EPLB or replicated expert slots." + ) + mapping = model_config.mapping + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + logger.info_once( + f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " + f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", + key="kimi_k3_moe_tp_ep_split", + ) + + mapping_dict = mapping.to_dict() + mapping_dict["moe_cluster_size"] = 1 + mapping_dict["moe_tp_size"] = moe_tp + mapping_dict["moe_ep_size"] = moe_ep + routed_mapping = Mapping.from_dict(mapping_dict) + + routed_model_config = copy.copy(model_config) + routed_model_config._frozen = False + routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) + routed_model_config.mapping = routed_mapping + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend in { + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + }: + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) + if configured_moe_max_num_tokens < default_moe_max_num_tokens: + logger.info_once( + "Kimi K3 MegaMoE raises moe_max_num_tokens from " + f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " + "because the global DP SymmBuffer requires capacity for " + "max_num_tokens * dp_size.", + key=( + "kimi_k3_megamoe_capacity_override_" + f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" + ), + ) + routed_model_config.moe_max_num_tokens = max( + configured_moe_max_num_tokens, + default_moe_max_num_tokens, + ) + routed_model_config._frozen = True + return routed_model_config + + def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: + """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel + share (``partial_tail``): the decode path takes the step with that tail.""" + return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + partial_tail: bool = False, + *, + step: DecodeStep, + ) -> Union[torch.Tensor, _decode_comm.PendingTail]: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification + (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where + ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at + most 8 tokens, a tensor on a wide decode step.""" + decode = self.decode_moe + if decode is not None and decode.takes(hidden_states, step, partial_tail): + return decode.forward( + self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push + ) + if partial_tail: + raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") + return self.forward_generic(hidden_states, all_rank_num_tokens) + + def forward_generic( + self, hidden_states: torch.Tensor, all_rank_num_tokens=None + ) -> torch.Tensor: + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None + + def _routed_output(): + # Latent down/up projections via the min-latency fused GEMM op: + # at <=16 tokens (decode graphs) it runs a single pipelined + # bf16 kernel per projection instead of cuBLAS's split-K GEMV + + # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); + # for larger token counts the op falls back to cuBLAS internally. + # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear + # (A/B escape hatch). When the FP8 weight-read conversion has + # replaced the projection module, call it directly: its weight is + # an e4m3 buffer the bf16 dsv3 op must not read, and its forward + # is already a single fused GEMM (fp8_swap_ab_gemm). + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + if self._reduce_routed_output: + return y + # Communication-backed paths return a complete routed result. + y = self.routed_expert_norm(y) + return self._routed_projection(y, self.routed_expert_up_proj) + + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. The shared GatedMLP includes its output all-reduce on + # the auxiliary stream. The join below must precede the routed + # all-reduce: concurrent collectives on different streams can corrupt + # SYMM_MEM all-reduce state. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + disable_on_compile=True, + ) + if self._reduce_routed_output: + routed_latent = moe_all_reduce(routed_out) + routed_latent = self.routed_expert_norm(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + return routed_out + shared_out + + +def resolve_attention_quant_config( + config: ModelConfig | None, layer_idx: int, projection: str +) -> QuantConfig: + """Resolve a checkpoint projection, including mixed-precision exclusions.""" + if config is None: + return QuantConfig() + global_config = config.quant_config or QuantConfig() + names = [ + f"{prefix}layers.{layer_idx}.self_attn.{projection}" + for prefix in ("language_model.model.", "model.", "") + ] + if any(global_config.is_module_excluded_from_quantization(name) for name in names): + return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) + declarations = config.quant_config_dict or {} + matches = [declarations[name] for name in names if name in declarations] + if matches: + selected = matches[0] + if any( + (match.quant_algo, match.group_size) != (selected.quant_algo, selected.group_size) + for match in matches[1:] + ): + raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") + elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: + selected = QuantConfig() + else: + selected = global_config + if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " + f"quantization {selected.quant_algo}" + ) + if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): + raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") + result = copy.copy(selected) + result.kv_cache_quant_algo = global_config.kv_cache_quant_algo + return result + + +class KimiMLARuntime(nn.Module): + """Wraps K3 MLA and applies its external TP output reduction.""" + + def __init__( + self, + cfg: "PretrainedConfig", + layer_idx: int, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the + # final output reduction in this wrapper so the output gate remains + # between attention and the row-parallel o_proj. + mapping = model_config.mapping + reduce_output = mapping.tp_size > 1 + self._o_allreduce = ( + AllReduce( + mapping=mapping, + strategy=model_config.allreduce_strategy, + dtype=torch.bfloat16, + ) + if reduce_output + else None + ) + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in ( + "q_a_proj", + "kv_a_proj_with_mqa", + "q_b_proj", + "kv_b_proj", + "g_proj", + "o_proj", + ) + } + attention_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + attention_config._frozen = model_config._frozen + self.mixer = K3DecodeMLA( + hidden_size=cfg.hidden_size, + num_heads=cfg.num_attention_heads, + q_lora_rank=cfg.q_lora_rank, + kv_lora_rank=cfg.kv_lora_rank, + qk_nope_head_dim=cfg.qk_nope_head_dim, + qk_rope_head_dim=cfg.qk_rope_head_dim, + v_head_dim=cfg.v_head_dim, + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + use_output_gate=cfg.mla_use_output_gate, + max_position_embeddings=max_positions, + model_config=attention_config, + aux_stream_dict=aux_stream_dict, + ) + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" + return self.mixer.will_run_decode_branch(attn_metadata, step) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the + mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer( + None, hidden_states, attn_metadata, step=step, project_output=project_output + ) + return self._reduce(out) if project_output and reduce_output else out + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + out = self.mixer.forward_generic(None, hidden_states, attn_metadata) + return self._reduce(out) if reduce_output else out + + def _reduce(self, out: torch.Tensor) -> torch.Tensor: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + return self._o_allreduce(out) if self._o_allreduce is not None else out + + +class KimiLinearDecoderLayer(nn.Module): + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + dtype = torch.bfloat16 + + self.is_kda = _is_kda_layer(cfg, layer_idx) + is_mla = _is_mla_layer(cfg, layer_idx) + if self.is_kda == is_mla: + raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") + + if self.is_kda: + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in projection_names + } + attention_config._frozen = model_config._frozen + self.linear_attn = K3DecodeKDA( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + ) + + self.is_moe = ( + cfg.num_experts is not None + and layer_idx >= cfg.first_k_dense_replace + and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 + ) + if self.is_moe: + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's + # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take + # (decode_gemv.SITES); otherwise it stays within one node. + spans_nodes = self._mnnvl_allreduce() is not None + if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. + self.mlp = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, + ) + self._situ = (situ_beta, situ_linear_beta) + # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics + # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied + # after downcast, use_gemma=False). + self.input_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3). The res norms stay + # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by + # _apply_attn_res and the fused attn_res op, never called as + # modules. + self.attn_res_block_size = cfg.attn_res_block_size + assert self.attn_res_block_size is not None, ( + "Kimi K3 runtime expects attn_res_block_size to be set" + ) + self.self_attention_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set + # by the target's post_load_weights. + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + self.sandwich_oproj = False + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + step: DecodeStep, + capture: Optional[Tuple[Any, int]] = None, + prenormed: bool = False, + pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, + defer_moe_tail: bool = False, + ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + ``block_residual`` is a preallocated snapshot bank in kernel-native + ``[K_max, M, H]`` layout. Returns the running prefix sum and the + number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. + + ``step`` is the step's classification (``decode_step``), handed to the + attention module. + + ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is + this layer's input norm, and the layer's input, the step's embedding, + already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). + + The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens + runs the attention's all-reduce and the residual update in one + collective once the target has built its ``decode_comm``: with o_proj + (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, + the step and the attention's decode branch, else on the attention's + unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). + + ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the + previous layer's MoE output, unreduced (``defer_moe_tail``); then + ``hidden_states`` is the prefix sum without it, and this layer's + pre-attention step reduces and adds it: a ``PendingTail`` in one + ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its + all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer + the ``sandwich_tail`` call also stores the tap into the capture slot + where the speculative metadata exposes it (``capture_view``); other + tapped layers take the tap after the step. + + ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` + instead, ``partial`` this layer's MoE output unreduced, for the next + consumer to reduce. + """ + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + attn_res_max_tokens = _attn_res_max_tokens(step) + tail = ( + pending_moe_partial + if isinstance(pending_moe_partial, _decode_comm.PendingTail) + else None + ) + # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum + # straight into the bank row it snapshots. + snapshot_row = None + if tail is not None and self.layer_idx % self.attn_res_block_size == 0: + snapshot_row = block_residual[num_snapshots] + # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the + # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). + tap_view = None + if tail is not None and capture is not None: + tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) + + if prenormed: + assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 + assert pending_moe_partial is None and capture is None + elif tail is not None: + hidden_states, prefix_sum = self.decode_comm.sandwich_tail( + tail, + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + updated_out=snapshot_row, + tap=tap_view, + tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, + ) + elif pending_moe_partial is not None: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + elif capture is not None: + # The mixture tap needs the PRE-norm value, which the fused + # attn-res + RMSNorm kernel does not expose. Keep the two steps + # split on captured layers only and fuse everywhere else. + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if capture is not None and pending_moe_partial is not None and tap_view is None: + # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output + # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. + tapped = ( + _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + if _AUX_ATTN_RES_STREAM_ENABLED + else prefix_sum + ) + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + + if self.layer_idx % self.attn_res_block_size == 0: + if not prenormed and snapshot_row is None: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, step): + if ( + self.sandwich_oproj + and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS + and attention.will_run_decode_branch(attn_metadata, step) + ): + core = attention(hidden_states, attn_metadata, step=step, project_output=False) + hidden_states, prefix_sum = comm.sandwich_oproj( + core, + self._o_proj().weight, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + if comm is not None and step.wide: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) + else: + attention_output = attention(hidden_states, attn_metadata, step=step) + if prefix_sum is None: + prefix_sum = attention_output + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention_output, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + if self.is_moe and defer_moe_tail: + partial = self.block_sparse_moe( + hidden_states, all_rank_num_tokens, partial_tail=True, step=step + ) + return prefix_sum, num_snapshots, partial + if self.is_moe: + hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) + else: + hidden_states = self._dense_mlp(hidden_states, step) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def forward_generic( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + ) -> Tuple[torch.Tensor, int]: + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + if capture is not None: + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, None): + partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + elif prefix_sum is None: + prefix_sum = attention.forward_generic(hidden_states, attn_metadata) + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention.forward_generic(hidden_states, attn_metadata), + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + if self.is_moe: + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + hidden_states = self.block_sparse_moe.forward_generic( + hidden_states, all_rank_num_tokens + ) + else: + hidden_states = self.mlp(hidden_states) + return prefix_sum + hidden_states, num_snapshots + + def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: + """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode + kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" + gemvs = self.decode_gemvs + if gemvs is not None and step.small: + out = gemvs.dense_mlp( + hidden_states, + self.mlp.gate_up_proj.weight, + self.mlp.down_proj.weight, + *self._situ, + ) + if out is not None: + return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out + return self.mlp(hidden_states) + + def _mnnvl_allreduce(self): + """The MNNVL all-reduce of this layer's attention output, or None.""" + attention = self.linear_attn if self.is_kda else self.self_attn + return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) + + def _o_proj(self) -> nn.Module: + """This layer's attention output projection (row parallel).""" + return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj + + def _o_allreduce(self) -> AllReduce: + """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" + return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce + + def accepts_moe_partial(self, num_snapshots: int) -> bool: + """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target + built its collective state, and the snapshot bank is not empty.""" + return num_snapshots > 0 and self.decode_comm is not None + + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + + +class KimiLinearModel(DecoderModel): + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = _get_text_config(model_config.pretrained_config) + self._text_cfg = cfg + dtype = torch.bfloat16 + + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream() for _ in range(4)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + } + + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + + # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) + # by _apply_attn_res and the fused attn_res op. + self.output_attn_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.num_attn_res_snapshots = ( + cfg.num_hidden_layers + cfg.attn_res_block_size - 1 + ) // cfg.attn_res_block_size + # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the + # weights are final. None: every step embeds on the generic path. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + + @property + def kda_draft_records(self) -> bool: + """Whether the hybrid cache manager keeps the last verify's draft records (``kda_state_tok``), from which + ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn`` replay the accepted drafts. The engine reads it once the weights + are loaded, to build the manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer + taking the K3 kernels. Otherwise the KDA verify replays the accepted drafts from the replay caches (the + built-in verify).""" + spec_config = getattr(self.model_config, "spec_config", None) + return bool( + spec_config is not None + and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) + and spec_config.tokens_per_gen_step in (2, 4, 6, 8) + and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) + ) + + def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, + from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from + load-time state, so every rank and every replay of the graph decides alike.""" + if not torch.cuda.is_current_stream_capturing(): + return False + kda = [layer.linear_attn for layer in self.layers if layer.is_kda] + mla = [layer.self_attn for layer in self.layers if not layer.is_kda] + return latent_push( + step, + capturing=True, + breakable=is_in_breakable_cuda_graph(), + kda_draft_records=self.kda_draft_records, + kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), + mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), + ) + + def _defer_moe_tail( + self, + i: int, + hidden_states: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + step: DecodeStep, + ) -> bool: + """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE + takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after + the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps + its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated + tail.""" + layer = self.layers[i] + if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): + return False + if spec_metadata is not None and ( + capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) + ): + return False + return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] + step = decode_step(attn_metadata, num_tokens) + if step is None: + return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) + if self._latent_push(attn_metadata, step): + step = replace(step, latent_push=True) + # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. + prenormed = None + if ( + inputs_embeds is None + and self.decode_gemvs is not None + and len(self.layers) > 0 + and self.num_attn_res_snapshots > 0 + ): + table = self.embed_tokens.weight + block_residual = table.new_empty( + self.num_attn_res_snapshots, num_tokens, table.shape[1] + ) + prenormed = self.decode_gemvs.embed_norm( + input_ids, table, self.layers[0].input_layernorm, block_residual + ) + if prenormed is not None: + hidden_states = prenormed + else: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) + # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. + pending_moe_partial = None + for i, layer in enumerate(self.layers): + defer_moe_tail = self._defer_moe_tail( + i, hidden_states, num_snapshots, spec_metadata, capture_set, step + ) + outputs = layer( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + step=step, + capture=self._capture(spec_metadata, capture_set, i), + prenormed=i == 0 and prenormed is not None, + pending_moe_partial=pending_moe_partial, + defer_moe_tail=defer_moe_tail, + ) + if defer_moe_tail: + hidden_states, num_snapshots, pending_moe_partial = outputs + else: + (hidden_states, num_snapshots), pending_moe_partial = outputs, None + + if isinstance(pending_moe_partial, _decode_comm.PendingTail): + normed, _ = self.layers[-1].decode_comm.sandwich_tail( + pending_moe_partial, + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + return normed + if pending_moe_partial is not None: + _, normed = _apply_attn_res_add_and_rmsnorm( + hidden_states, + _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + _attn_res_max_tokens(step), + ) + return normed + return self._final_norm( + hidden_states, + block_residual, + num_snapshots, + spec_metadata, + capture_set, + _attn_res_max_tokens(step), + ) + + def _generic_forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor], + inputs_embeds: Optional[torch.FloatTensor], + spec_metadata, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + for i, layer in enumerate(self.layers): + hidden_states, num_snapshots = layer.forward_generic( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + capture=self._capture(spec_metadata, capture_set, i), + ) + return self._final_norm( + hidden_states, block_residual, num_snapshots, spec_metadata, capture_set + ) + + def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + return (spec_metadata, self.layers[i - 1].layer_idx) + return None + + def _final_norm( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + max_fused_tokens: Optional[int] = None, + ) -> torch.Tensor: + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) + + return _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + max_fused_tokens, + ) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. +# ---------------------------------------------------------------------------------------------------------------------- + + +class K3DecodeKDA(KimiKDALinearAttention): + """Quantized KDA with the target's output-reduction boundary.""" + + takes_k3_kernels = False + k3_buffers = None + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + return False + + def forward(self, hidden_states, attn_metadata, step, reduce_output=True, project_output=True): + if not project_output: + raise ValueError("Quantized KDA requires its own output projection") + return self.forward_generic(hidden_states, attn_metadata, reduce_output) + + def forward_generic(self, hidden_states, attn_metadata, reduce_output=True): + if reduce_output: + return super().forward(hidden_states, attn_metadata) + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + else: + core = self._forward_impl(hidden_states, attn_metadata) + return self.o_proj(core.reshape(-1, self.proj_size)) + + +class K3DecodeMLA(KimiK3MLAAttention): + """Stock quantized MLA, including its KV-cache format handling.""" + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + return False + + def forward(self, position_ids, hidden_states, attn_metadata, step, project_output=True): + if not project_output: + raise ValueError("Quantized MLA requires its own output projection") + return self.forward_generic(position_ids, hidden_states, attn_metadata) + + def forward_generic(self, position_ids, hidden_states, attn_metadata): + return super().forward(position_ids, hidden_states, attn_metadata) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. +# ---------------------------------------------------------------------------------------------------------------------- + +#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at +#: max_draft_len 7. +DRAFTER_ATTN_SPLITS = frozenset( + {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} +) +# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k +# RMSNorm's epsilon and the NeoX RoPE's base. +_DRAFTER_HEADS = (6, 1) +_DRAFTER_HEAD_DIM = 64 +_DRAFTER_PAGE = 64 +_DRAFTER_EPS = 1e-5 +_DRAFTER_ROPE_BASE = 10000.0 + + +def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: + """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input + feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" + if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: + return None + width = in_features // tp_size + return tp_rank * width, (tp_rank + 1) * width + + +class K3FcSlice(nn.Module): + """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is + ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's + ranks is ``fc``'s output.""" + + def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: + super().__init__() + self.weight = nn.Parameter(weight, requires_grad=False) + self.start = start + self.end = end + + @classmethod + def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: + """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns + (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" + columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) + if columns is None: + return None + start, end = columns + return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the + GEMM reads the column block in place, through the rows' stride.""" + return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) + + +class K3DSparkDrafter(GQADSparkForCausalLM): + """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. + + A block the entries take runs, per layer: + + * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; + * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to + its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the + worker writes the accepted rows' context k / v before a block reads them; + * the output projection on ``drafter_o``, then the module's all-reduce; + * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the + module's all-reduce. + + The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state + (below); a projection whose site does not take its rows runs its module. Every other block runs the stock + ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or + normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under + CUDA-graph capture. + + Where the target hands the drafter the TP group's collective state (``use_decode_comm``): + + * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and + ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; + * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the + post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the + block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the + projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection + after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is + followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers + whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. + + Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the + stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); + the context k / v and the Markov head stay upstream code. + """ + + def __init__( + self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" + ) -> None: + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. + self._k3_take: Optional[bool] = None + # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes + # only these. + self._k3_attn_ran: set = set() + # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + # The zero residual of the split context projection's all-reduce, up to a decode step's rows. + self._k3_zero_rows: Optional[torch.Tensor] = None + # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms + # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. + self._k3_norms_fuse = False + self._k3_sandwich_forms: frozenset = frozenset() + + def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: + """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the + target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are + loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. + + * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they + split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` + sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. + * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and + RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose + shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not + compile is not used.""" + self.decode_comm = comm + self._k3_split_fc() + self._k3_norms_fuse = self._k3_norms_take_comm() + forms = set() + if self._k3_norms_fuse: + layers = self.model.layers + for form, swiglu, weights in ( + ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), + ("down", True, [layer.mlp.down_proj.weight for layer in layers]), + ): + if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( + weights[0], swiglu=swiglu + ): + forms.add(form) + self._k3_sandwich_forms = frozenset(forms) + logger.info( + "Kimi K3 DSpark drafter: residual adds and RMSNorms " + + ( + f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " + "mnnvl_fusion_allreduce)" + if self._k3_norms_fuse + else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" + ) + ) + + def _k3_norms_take_comm(self) -> bool: + """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and + post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no + quantized output), and every output and down projection is row parallel with its all-reduce.""" + hidden = self.config.hidden_size + + def plain(norm: nn.Module) -> bool: + weight = getattr(norm, "weight", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (hidden,) + and hasattr(norm, "variance_epsilon") + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "return_hp_output", True) + ) + + def reduces(linear: nn.Module) -> bool: + return ( + getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" + and getattr(linear, "reduce_output", False) + and getattr(linear, "all_reduce", None) is not None + ) + + return plain(self.model.norm) and all( + plain(layer.input_layernorm) + and plain(layer.post_attention_layernorm) + and reduces(layer.self_attn.o_proj) + and reduces(layer.mlp.down_proj) + for layer in self.model.layers + ) + + def _k3_tp_all_reduce(self) -> Optional[nn.Module]: + """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" + o_proj = self.model.layers[0].self_attn.o_proj + if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": + return None + return getattr(o_proj, "all_reduce", None) + + def _k3_split_fc(self) -> None: + """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero + residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is + a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP + all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" + fc = getattr(self, "fc", None) + if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): + return + mapping = self.model_config.mapping + weight = getattr(fc, "weight", None) + sliced = None + if ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and getattr(fc, "bias", None) is None + and self._k3_tp_all_reduce() is not None + ): + sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) + if sliced is None: + logger.info( + "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " + f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" + ) + return + self.fc = sliced + self._k3_zero_rows = weight.new_zeros( + MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] + ) + logger.info( + f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " + f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" + ) + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" + result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + self._k3_split_fc() + return result + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. + + With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group + with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode + step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then + ``hidden_norm``. Otherwise the stock projection.""" + fc = self.fc + if not isinstance(fc, K3FcSlice): + return super().project_target_hidden(hidden_states) + partial = fc(hidden_states.to(fc.weight.dtype)) + rows = partial.shape[0] + zeros = self._k3_zero_rows + if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): + normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) + return normed + if rows == 0: + return self.hidden_norm(partial) + return self.hidden_norm(self._k3_tp_all_reduce()(partial)) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class + docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows + ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows + as a view instead of gathering them.""" + keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) + if keys is None: + return super().dflash_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ) + out = self._k3_block_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ctx_rows_start, + ) + if not torch.cuda.is_current_stream_capturing(): + self._k3_attn_ran |= keys + return out + + def _k3_layers_take(self) -> bool: + """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head + layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, + non-causal unwindowed block attention, and no block convolution or post-attention gate.""" + if self._k3_take is None: + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + take = ( + self._use_fused_qk_norm_rope + and not self.has_block_conv + and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate + ) + for layer_idx, layer in enumerate(self.model.layers): + attn, mlp = layer.self_attn, layer.mlp + pos = getattr(attn, "pos_embd_params", None) + rope = getattr(pos, "rope", None) + norms = (attn.q_norm, attn.k_norm) + take = take and ( + (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS + and attn.head_dim == _DRAFTER_HEAD_DIM + and getattr(attn, "is_qk_norm", False) + and not getattr(attn, "use_gemma_rms_norm", True) + and rope is not None + and pos.is_neox + and getattr(pos, "mrope_section", None) is None + and rope.theta == _DRAFTER_ROPE_BASE + and all( + norm.variance_epsilon == _DRAFTER_EPS + and norm.weight.dtype == torch.bfloat16 + and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) + for norm in norms + ) + and attn.qkv_proj.bias is None + and attn.o_proj.bias is None + and mlp.activation is torch.nn.functional.silu + and mlp.swiglu_limit is None + and mlp.swiglu_alpha in (None, 1.0) + and mlp.swiglu_beta in (None, 0.0) + and mlp.gate_up_proj.bias is None + and mlp.down_proj.bias is None + and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) + ) + self._k3_take = bool(take) + logger.info( + f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " + + ( + "on k3_drafter_attn_qknorm and the drafter GEMV sites" + if self._k3_take + else "on the stock block forward (a layer the drafter entries do not take)" + ) + ) + return self._k3_take + + def _k3_block_keys( + self, + noise_embedding: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor], + ctx_page_table: Optional[torch.Tensor], + ) -> Optional[frozenset]: + """The attention's compile keys for this block when the drafter entries take it, else None.""" + if ( + self.dflash_attention_backend != "TRTLLM" + or ctx_kv_cache is None + or ctx_page_table is None + or noise_embedding.dtype != torch.bfloat16 + or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS + or is_in_breakable_cuda_graph() + or not self._k3_layers_take() + ): + return None + num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads + page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) + dense = ( + num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + ) + strides = set() + for layer_idx in range(len(self.model.layers)): + cache = ctx_kv_cache[layer_idx] + if ( + cache.dtype != torch.bfloat16 + or cache.dim() != 5 + or tuple(cache.shape[1:]) != page + or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) + or cache.stride(0) % 8 + ): + return None + strides.add(cache.stride(0)) + keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) + if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: + return None + return keys + + def _k3_block_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: torch.Tensor, + ctx_page_table: torch.Tensor, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + batch, block = noise_embedding.shape[:2] + rows = batch * block + ctx_len = num_ctx_per_req[:batch].to(torch.int32) + if ctx_rows_start is None: + page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) + else: + # The batch's page-table rows are one contiguous run: a view, no gather. + page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] + positions = query_positions.reshape(-1).contiguous() + hidden = noise_embedding.reshape(rows, -1) + layers = self.model.layers + comm = self._k3_fused_comm(hidden) + residual = None + if comm is not None: + # The fused all-reduces read the residual and return the updated one: the block's input serves as the + # first residual, uncopied. + residual = hidden + normed = layers[0].input_layernorm(hidden) + for layer_idx, layer in enumerate(layers): + attn = layer.self_attn + if comm is None: + if residual is None: + residual = hidden.clone() + normed = layer.input_layernorm(hidden) + else: + normed, residual = layer.input_layernorm(hidden, residual) + qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) + out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) + k3_drafter_attn_qknorm( + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + positions, + attn.q_norm.variance_epsilon, + attn.pos_embd_params.rope.theta, + ctx_kv_cache[layer_idx], + page_table, + ctx_len, + attn.num_heads, + attn.num_key_value_heads, + out, + ) + if comm is not None: + # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after + # the last layer the final norm. + next_norm = ( + layers[layer_idx + 1].input_layernorm + if layer_idx + 1 < len(layers) + else self.model.norm + ) + normed, residual = self._k3_project_norm( + comm, out, attn.o_proj, residual, layer.post_attention_layernorm + ) + normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) + continue + hidden = self._k3_project("drafter_o", out, attn.o_proj) + hidden, residual = layer.post_attention_layernorm(hidden, residual) + hidden = self._k3_mlp(layer.mlp, hidden) + if comm is not None: + return normed + out, _ = self.model.norm(hidden, residual) + return out + + def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: + """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its + all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace + holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" + comm = self.decode_comm + if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): + return None + return comm + + def _k3_project_norm( + self, + comm: _decode_comm.K3DecodeComm, + x: torch.Tensor, + linear: nn.Module, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, + the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the + projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` + where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" + if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( + x, linear.weight, residual, norm + ): + return comm.sandwich_plain(x, linear.weight, residual, norm) + gemvs = self.decode_gemvs + partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) + if partial is None: + partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_mlp_norm( + self, + comm: _decode_comm.K3DecodeComm, + mlp: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down + projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul + form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down + projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; + then ``comm/mnnvl_fusion_allreduce``.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + partial = None + if gate_up is not None: + if "down" in self._k3_sandwich_forms and comm.takes_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ): + return comm.sandwich_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ) + partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + if partial is None: + partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: + """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel + projection's all-reduce.""" + y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) + if y is None: + return linear(x) + return _row_parallel_reduce(linear, y) + + def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on + ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + down = ( + None + if gate_up is None + else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + ) + if down is None: + return mlp(x) + return _row_parallel_reduce(mlp.down_proj, down) + + +def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: + """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" + all_reduce = getattr(linear, "all_reduce", None) + if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: + return all_reduce(partial) + return partial + + +# ---------------------------------------------------------------------------------------------------------------------- +# The target: step classification, the construction checks and the class route returns. +# ---------------------------------------------------------------------------------------------------------------------- + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. + + The checkpoint names its language-model modules `language_model.`; the text model's are ``, with + `layers.*` under `model.`. + """ + config = model_config.pretrained_config + text = copy.copy(model_config) + text._frozen = False + text.pretrained_config = config.text_config + excluded = text.quant_config.exclude_modules + if excluded: + text.quant_config = copy.copy(text.quant_config) + renamed = [] + for name in excluded: + if name.startswith(_LANG_PREFIX): + name = name[len(_LANG_PREFIX) :] + if name.startswith("layers."): + name = "model." + name + renamed.append(name) + text.quant_config.exclude_modules = renamed + text.skip_create_weights_in_init = True + text._frozen = True + return text + + +@dataclass(frozen=True) +class DecodeStep: + """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` + generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" + + num_tokens: int + num_requests: Optional[int] = None + tokens_per_request: Optional[int] = None + #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` + #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. + latent_push: bool = False + + @property + def small(self) -> bool: + """Whether the step fits one token tile of the token-count kernels.""" + return self.num_tokens <= DECODE_MAX_TOKENS + + @property + def decode(self) -> bool: + """Whether the step is a pure decode step the request-aware kernels take.""" + return self.num_requests is not None + + @property + def wide(self) -> bool: + """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode + layout's MoE head and tail, on M-general ops.""" + return self.decode and not self.small + + +def latent_push( + step: DecodeStep, + *, + capturing: bool, + breakable: bool, + kda_draft_records: bool, + kda_decode_kernels: bool, + mla_decode_branch: bool, +) -> bool: + """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at + most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a + breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels + (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA + layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). + Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel + between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a + graph launch orders every pushing replay behind whatever ran before it.""" + if not (step.decode and step.small) or not capturing or breakable: + return False + if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): + return False + return mla_decode_branch + + +def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: + """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). + + ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers + and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input + here is fixed by that shape, so a captured step and its replays are classified alike. + """ + if num_tokens <= 0: + return None + requests = _decode_requests(attn_metadata, num_tokens) + if requests is not None: + return DecodeStep(num_tokens, requests, num_tokens // requests) + if num_tokens <= DECODE_MAX_TOKENS: + return DecodeStep(num_tokens) + return None + + +def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: + """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" + if attn_metadata.num_contexts != 0: + return None + requests = attn_metadata.num_generations + if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: + return None + tokens = num_tokens // requests + if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: + return None + seq_lens = getattr(attn_metadata, "seq_lens", None) + if seq_lens is not None and seq_lens.device.type == "cpu": + lens = seq_lens[:requests] + if lens.numel() != requests or bool((lens != tokens).any()): + return None + return requests + + +def _check_construction(model_config: ModelConfig) -> None: + """The settings this target is built for that are fixed before the first step.""" + capability = torch.cuda.get_device_capability() + assert capability == _SM, ( + f"this target requires sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" + ) + mapping = model_config.mapping + topology = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.enable_attention_dp, + ) + assert topology == (16, 16, 1, 4, 4, False), ( + "the tp16_moetp4ep4 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " + "moe_tensor_parallel_size 4, moe_expert_parallel_size 4 and enable_attention_dp false; the engine built " + f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" + ) + assert model_config.torch_dtype == torch.bfloat16, ( + f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" + ) + for layer_idx in range(1, 93): + quant = KimiK3MoERuntime._resolve_routed_quant_config(model_config, layer_idx) + if quant.quant_algo != QuantAlgo.NVFP4 or quant.group_size != 16: + raise ValueError( + f"NVFP4 target requires group-size-16 NVFP4 experts at layer {layer_idx}; " + f"got {quant.quant_algo}, group_size={quant.group_size}" + ) + strategy = model_config.allreduce_strategy + assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( + f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" + ) + + +class ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4(KimiLinearForCausalLM): + """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. + + It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through + weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model + keeps the built-in one's. + """ + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The V2 hybrid manager holds the KDA states; the step contract requires it.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + config = model_config.pretrained_config + assert getattr(config, "text_config", None) is not None, ( + "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " + "text_config" + ) + _check_construction(model_config) + text = _text_model_config(model_config) + spec_config = getattr(text, "spec_config", None) + assert ( + spec_config is None + or spec_config.spec_dec_mode.is_sa() + or spec_config.spec_dec_mode.is_dflash() + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" + # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 + # weight-read conversion of the shared / latent MLPs. + self._fp8_weight_read_moe_mlp = False + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + cfg = text.pretrained_config + SpecDecOneEngineForCausalLM.__init__( + self, + KimiLinearModel(text), + text, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self._own_spec_branch(text) + self._step_checked = False + # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the + # processor the speculative worker's target logits and a parallel drafter's logits call too. + stock_logits_processor = self.logits_processor + self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) + draft_model = getattr(self, "draft_model", None) + if getattr(draft_model, "logits_processor", None) is stock_logits_processor: + draft_model.logits_processor = self.logits_processor + # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which + # must therefore be the text config, as the built-in wrapper leaves it. + model_config._frozen = False + model_config.pretrained_config = self.config + model_config._frozen = True + + def _own_spec_branch(self, model_config: ModelConfig) -> None: + """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a + standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is + `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" + spec_config = model_config.spec_config + if spec_config is None: + return + mode = spec_config.spec_dec_mode + standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) + if ( + standalone_dspark + and self.draft_config is not None + and is_mla(self.draft_config.pretrained_config) + ): + # The reference NVFP4 checkpoint uses an MLA DSpark drafter. Its + # compressed context cache does not meet the GQA worker contract. + return + if ( + standalone_dspark + and self.draft_config is not None + and not is_mla(self.draft_config.pretrained_config) + ): + drafter = K3DSparkDrafter( + self.draft_config, dflash_attention_backend=spec_config.attention_backend + ) + if mode.is_parallel_draft(): + drafter.logits_processor = self.logits_processor + self.epilogue[self.epilogue.index(self.draft_model)] = drafter + self.draft_model = drafter + if mode.is_dflash(): + worker_cls = _spec_worker.KimiK3DFlashWorker + elif standalone_dspark: + worker_cls = _spec_worker.KimiK3DSparkWorker + else: + return + worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) + worker.set_draft_model(self.draft_model) + self.epilogue[self.epilogue.index(self.spec_worker)] = worker + self.spec_worker = worker + + def load_weights(self, weights, *args, **kwargs): + _weights.load(self, weights) + + def cache_derived_state(self) -> None: + """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of + every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: + a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites + are among them, and the drafter gets the state too.""" + super().cache_derived_state() + if self.model.decode_gemvs is not None: + return + sites = ("dense_gate_up", "dense_down") + if isinstance(self.draft_model, K3DSparkDrafter): + sites += ("drafter_qkv", "drafter_o", "drafter_gate_up", "drafter_down") + gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head, sites=sites) + self.model.decode_gemvs = gemvs + self.logits_processor.gemvs = gemvs + for layer in self.model.layers: + if not layer.is_moe: + layer.decode_gemvs = gemvs + if isinstance(self.draft_model, K3DSparkDrafter): + self.draft_model.decode_gemvs = gemvs + + def post_load_weights(self) -> None: + """Prepare BF16 decode state and MNNVL collectives before graph capture.""" + super().post_load_weights() + if getattr(self, "_nvfp4_decode_state_ready", False): + return + layers = self.model.layers + comm = None + if all(layer._mnnvl_allreduce() is not None for layer in layers): + comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, None) + for layer in layers: + layer.decode_comm = comm + layer.sandwich_oproj = False + _decode_comm.use_decode_one_shot(self) + self._gate_spec_worker_kernels(comm) + self._gate_drafter_comm(comm) + self._nvfp4_decode_state_ready = True + logger.info( + "Kimi K3 NVFP4 SM107: stock quantized attention and NVFP4 experts; " + "BF16 decode GEMVs and DSpark specializations; " + f"attention-residual MNNVL fusion={comm is not None}" + ) + + def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, + ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input + that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's + ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the + vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks + each step's own conditions. Returns the setting (False without such a worker).""" + worker = getattr(self, "spec_worker", None) + if not hasattr(worker, "k3_decode"): + return False + gemvs = self.model.decode_gemvs + worker.k3_decode = ( + comm is not None and gemvs is not None and gemvs.head_workspace is not None + ) + logger.info( + f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " + + ( + "on" + if worker.k3_decode + else "off (no MNNVL decode state or no k3_head_gemv LM head)" + ) + ) + return worker.k3_decode + + def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built + it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context + projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms + in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on + the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate + requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` + workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the + drafter took the state (False without such a drafter).""" + drafter = getattr(self, "draft_model", None) + if comm is None or not isinstance(drafter, K3DSparkDrafter): + return False + drafter.use_decode_comm(comm) + return True + + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """First-forward checks of the engine surface and the per-engine settings.""" + objects = { + "attn_metadata": attn_metadata, + "kv_cache_manager": attn_metadata.kv_cache_manager, + } + missing = [ + f"{owner}.{name}" + for owner, names in REQUIRED_ENGINE_FIELDS.items() + for name in names + if not hasattr(objects[owner], name) + ] + assert not missing, f"engine fields this target reads are missing: {missing}" + manager = attn_metadata.kv_cache_manager + assert isinstance(manager, MambaHybridCacheManagerV2), ( + "this target needs the V2 hybrid KV / state cache manager " + "(kv_cache_config.use_kv_cache_manager_v2); the engine built " + f"{type(manager).__name__}" + ) + self._step_checked = True + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + if kwargs.pop("multimodal_params", None): + raise ValueError( + "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" + ) + if not self._step_checked: + self._check_step_contract(attn_metadata) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py new file mode 100644 index 000000000000..abc1a1388213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py @@ -0,0 +1,1343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import functools +import inspect +from typing import Optional + +import torch +from torch import nn + +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.speculative.dflash import ( + DFlashSpecMetadata, + DFlashWorker, + dflash_allocated_ctx_limit, + dflash_noise_block_embedding, +) +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase +from tensorrt_llm.logger import logger + + +@functools.lru_cache(maxsize=None) +def _takes_ctx_rows_start_cls(cls) -> bool: + return "ctx_rows_start" in inspect.signature(cls.dflash_forward).parameters + + +def _takes_ctx_rows_start(draft_model) -> bool: + """Whether the draft model's ``dflash_forward`` accepts ``ctx_rows_start``.""" + return _takes_ctx_rows_start_cls(type(draft_model)) + + +def dflash_ctx_rows_kwargs(draft_model, ctx_rows_start: Optional[int]) -> dict: + """The ``ctx_rows_start`` keyword (from ``prepare_1st_drafter_inputs``) for a draft model whose ``dflash_forward`` + takes it; no keyword for another draft model or without a start.""" + if ctx_rows_start is None or not _takes_ctx_rows_start(draft_model): + return {} + return {"ctx_rows_start": ctx_rows_start} + + +def capture_view( + spec_metadata: Optional[SpecMetadata], layer_id: int, num_tokens: int +) -> Optional[torch.Tensor]: + """``layer_id``'s slot of the DFlash capture buffer for ``num_tokens`` rows (a strided view a kernel can write the + tap into), or None when ``spec_metadata`` is not DFlash's or the layer is not captured.""" + if ( + not isinstance(spec_metadata, DFlashSpecMetadata) + or spec_metadata.captured_hidden_states is None + ): + return None + i = spec_metadata._layer_to_idx.get(layer_id) + if i is None: + return None + return spec_metadata.captured_hidden_states[ + :num_tokens, i * spec_metadata.hidden_size : (i + 1) * spec_metadata.hidden_size + ] + + +class KimiK3DFlashWorker(DFlashWorker): + # Set by a Kimi K3 target running its decode kernels: decode steps then run the acceptance and the drafter's + # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context + # K/V with trtllm::k3_ctx_kv. + k3_decode = False + + def _mask_token_id(self, draft_model) -> int: + """The drafter's mask token id, resolved once.""" + if self._resolved_mask_token_id is None: + if ( + hasattr(self.spec_config, "mask_token_id") + and self.spec_config.mask_token_id is not None + ): + self._resolved_mask_token_id = self.spec_config.mask_token_id + elif hasattr(draft_model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.mask_token_id + elif hasattr(draft_model.model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.model.mask_token_id + else: + raise ValueError( + "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " + "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." + ) + return self._resolved_mask_token_id + + @staticmethod + def _trained_mask_embedding(draft_model, mask_token_id: int) -> Optional[torch.Tensor]: + """The drafter's own trained mask row, if it kept one for the mask id in use.""" + if mask_token_id != getattr(draft_model, "mask_token_id", None): + return None + return getattr(draft_model, "mask_token_embedding", None) + + def _k3_accept_applies( + self, logits, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Whether this step's acceptance and the drafter's inputs run as ``trtllm::k3_spec_accept``. + + That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no + rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target + logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block + table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. + """ + if not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return False + K = spec_metadata.runtime_draft_len + if ( + logits.dim() != 2 + or logits.dtype != torch.float32 + or logits.shape[0] != num_gens * (K + 1) + ): + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + embed = draft_model.draft_model_full.model.embed_tokens + return accept_op.supports( + logits.shape[1], num_gens, self._compute_block_size, K, embed.weight.shape[1] + ) + + def _k3_accept_step_applies( + self, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Everything ``_k3_accept_applies`` checks but the target logits.""" + if not self.k3_decode: + return False + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + + if num_contexts != 0 or not 0 < num_gens <= 8 or self.guided_decoder is not None: + return False + if ( + not spec_metadata.is_all_greedy_sample + or self._can_use_rejection_sampling(spec_metadata) + or getattr(spec_metadata, "enable_penalty", False) + or type(self)._reshape_draft_tokens_for_accept + is not SpecWorkerBase._reshape_draft_tokens_for_accept + or type(self)._reshape_logits_for_accept + is not SpecWorkerBase._reshape_logits_for_accept + ): + return False + K = spec_metadata.runtime_draft_len + draft_tokens = spec_metadata.draft_tokens + if ( + K <= 0 + or draft_tokens is None + or draft_tokens.dtype != torch.int32 + or draft_tokens.numel() != num_gens * K + ): + return False + mgr = attn_metadata.kv_cache_manager + mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) + if ( + not isinstance(mgr, MambaHybridCacheManagerV2) + or type(mgr).update_mamba_states is not MambaHybridCacheManagerV2.update_mamba_states + or not mgr.use_kda_replay_update + or getattr(mgr, "prev_num_accepted_tokens", None) is None + or getattr(mgr, "_dummy_request_mask", None) is None + or mamba_metadata is None + or mamba_metadata.state_indices.dtype != torch.int32 + ): + return False + if ( + self._ctx_block_tables is None + or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None + or getattr(attn_metadata, "kv_lens_cuda", None) is None + ): + return False + embed = draft_model.draft_model_full.model.embed_tokens + return getattr(embed, "tp_size", 1) == 1 and embed.weight.dtype == torch.bfloat16 + + def target_logits( + self, hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) -> torch.Tensor: + """This rank's bf16 vocabulary shard of the target logits, [rows, vocab / TP] (``_k3_head_shard``: the head + GEMM without its all-gather and fp32 cast), for a step whose acceptance runs as the sharded + ``trtllm::k3_spec_accept`` (see ``_k3_logits_shard``); the gathered fp32 logits otherwise. At such a step the + acceptance is their only reader: the one-model spec sampler takes the worker's tokens and rejects requests for + logits or log probabilities, so the shard is also what the step returns as ``"logits"``.""" + self._k3_step_shard = None + shard = self._k3_logits_shard( + hidden_states, lm_head, attn_metadata, spec_metadata, draft_model + ) + if shard is None: + return super().target_logits( + hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) + logits = self._k3_head_shard(logits_processor, lm_head, hidden_states) + self._k3_step_shard = (logits, shard) + return logits + + @staticmethod + def _k3_head_shard(logits_processor, lm_head, rows: torch.Tensor) -> torch.Tensor: + """This rank's bf16 vocabulary shard of ``lm_head(rows)``, without the head's all-gather: from the logits + processor's own head kernel where it has one that takes the rows (``lm_head_shard``; the Kimi K3 target's + ``gemm/k3_head_gemv``), so the shard holds the values of the processor's gathered logits; else from the head's + own GEMM.""" + head_shard = getattr(logits_processor, "lm_head_shard", None) + logits = None if head_shard is None else head_shard(rows, lm_head) + if logits is None: + logits = lm_head.apply_linear(rows, lm_head.bias) + return logits + + def _k3_logits_shard(self, hidden_states, lm_head, attn_metadata, spec_metadata, draft_model): + """``(workspace, first column)`` when this step's target logits stay vocabulary-sharded, else None. + + That needs a step ``_k3_accept_step_applies`` takes, plain TP whose all-reduces own an MNNVL workspace for + this mapping, and an unquantized, bias-free, unpadded, evenly split column-parallel bf16 lm_head with shards + the kernel splits. The exchange is collective, so every rank must decide alike: every condition is the + batch's or the configuration's. The workspace is allocated (collectively) and the kernel compiled outside + CUDA-graph capture only, so a graph whose warmup allocated no workspace is captured on the gathered path. + """ + if not self.k3_decode: + return None + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + mapping = self.mapping + if getattr(self, "_k3_shard_head", None) is not lm_head: + reason = self._k3_logits_shard_reason(lm_head) + if reason is None and MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None: + # Not kept: re-checked on later steps, since the model's MNNVL workspace may not exist yet. + logger.info_once( + "DFlash: the target logits are all-gathered (the TP all-reduces are not MNNVL)", + key="dflash_k3_logits_no_mnnvl", + ) + return None + if mapping is not None and mapping.tp_size > 1: + if reason is None: + logger.info( + f"DFlash: decode-step target logits stay vocabulary-sharded ({lm_head.weight.shape[0]} " + f"columns per rank); trtllm::k3_spec_accept exchanges the row maxima, and a request's " + f"logits post-processors get them gathered" + ) + else: + logger.info_once( + f"DFlash: the target logits are all-gathered ({reason})", key=reason + ) + self._k3_shard_head, self._k3_shard_ok = lm_head, reason is None + if not self._k3_shard_ok: + return None + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + K = spec_metadata.runtime_draft_len + if hidden_states.shape[0] != num_gens * (K + 1) or not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return None + embed = draft_model.draft_model_full.model.embed_tokens + columns = lm_head.weight.shape[0] + if not accept_op.supports( + columns, num_gens, self._compute_block_size, K, embed.weight.shape[1], mapping.tp_size + ): + return None + if torch.cuda.is_current_stream_capturing(): + ws = accept_op.existing_workspace(mapping) + if ws is None: + return None + else: + ws = accept_op.workspace(mapping) + return ws, mapping.tp_rank * columns + + def _k3_logits_shard_reason(self, lm_head): + """Why the target logits cannot stay vocabulary-sharded under this configuration and head (None if they can, + given an MNNVL workspace for the mapping, which ``_k3_logits_shard`` checks).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + if mapping is None or mapping.tp_size <= 1: + return "no TP" + if mapping.enable_attention_dp or mapping.has_cp() or mapping.has_pp(): + return "attention DP, CP or PP" + if ( + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or getattr(lm_head, "gather_output_sizes", None) is not None + ): + return "the lm_head is not an evenly split, unpadded column-parallel head" + if ( + lm_head.bias is not None + or lm_head.has_any_quant + or lm_head.weight.dtype != torch.bfloat16 + ): + return "the lm_head has a bias, is quantized or is not bf16" + if not accept_op.supports_columns(lm_head.weight.shape[0], mapping.tp_size): + columns = lm_head.weight.shape[0] + return f"trtllm::k3_spec_accept cannot split {columns} columns over {mapping.tp_size} ranks" + return None + + def _k3_static_arange(self, lo: int, hi: int) -> torch.Tensor: + """``arange(lo, hi)`` (int64), built once outside CUDA-graph capture.""" + cache = self.__dict__.setdefault("_k3_aranges", {}) + t = cache.get((lo, hi)) + if t is None: + t = torch.arange(lo, hi, dtype=torch.long, device="cuda") + if not torch.cuda.is_current_stream_capturing(): + cache[(lo, hi)] = t + return t + + def _k3_mask_row(self, draft_model) -> torch.Tensor: + """The noise block's mask embedding, bf16 [hidden], cached: the drafter's own trained mask row if it kept one + for the mask id in use, else the embedding's row (``dflash_noise_block_embedding`` reads the same row).""" + row = getattr(self, "_k3_mask_embedding", None) + if row is None: + embed = draft_model.draft_model_full.model.embed_tokens + mask_token_id = self._mask_token_id(draft_model) + trained = self._trained_mask_embedding(draft_model, mask_token_id) + if trained is not None: + row = trained.to(embed.weight.dtype).reshape(-1).contiguous() + else: + row = embed.weight[mask_token_id].contiguous() + if torch.cuda.is_current_stream_capturing(): + return row + self._k3_mask_embedding = row + return row + + def _k3_ctx_pool_key(self): + """The identity of the pool the ctx cache is bound to now (KV cache estimation rebinds it to a new one).""" + return ( + tuple(t.data_ptr() for t in self._ctx_kv_buf) if self._ctx_kv_buf is not None else None + ) + + def _k3_ctx_kv_applies(self, draft_model, projected: torch.Tensor) -> bool: + """Whether ``trtllm::k3_ctx_kv`` writes this step's context K/V: the manager-bound paged pool with K and V, + bf16, the fused K/V weight without bias or context input norm, k_norm, NeoX RoPE from flashinfer's fp32 + cache, a context the op can split (up to 64 tokens: B <= 8 requests of K + 1 <= 8, ``ctx_op.pick_split``), one + allocation for every layer's pool. Checked once per token count and + bound pool; the pool view is kept for the bound pool only, so a rebind never leaves the kernel writing the + replaced pool.""" + pool_key = self._k3_ctx_pool_key() + if getattr(self, "_k3_ctx_pool_bound", None) != pool_key: + self._k3_ctx_pool_bound = pool_key + self._k3_ctx_pool = None + self._k3_ctx_kv_ok = {} + cache = self._k3_ctx_kv_ok + n = projected.shape[0] + ok = cache.get(n) + if ok is not None: + return ok + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctx_kv import op as ctx_op + from tensorrt_llm._torch.models import modeling_dflash + + reason = None + if not (self._ctx_paged and self._ctx_block_tables is not None): + reason = "the context pool is not the manager-bound paged pool" + elif projected.dtype != torch.bfloat16 or self._ctx_kv_buf[0].dtype != torch.bfloat16: + reason = "not bf16" + else: + if draft_model._fused_kv_weight is None: + draft_model._build_fused_kv_buffers() + view = ctx_op.pool_view(list(self._ctx_kv_buf)) + if ( + draft_model._fused_kv_bias is not None + or getattr(draft_model, "_input_ln_eps", None) is not None + ): + reason = "K/V bias or context input norm" + elif draft_model._k_norm_stacked is None or not draft_model._is_neox: + reason = "no k_norm or not NeoX RoPE" + elif ( + modeling_dflash._flashinfer_rope is None + or draft_model._get_cos_sin_cache().shape[1] != 64 + ): + reason = "RoPE is not flashinfer's full-rotary fp32 cache" + elif view is None or self._ctx_kv_buf[0].size(1) != 2: + reason = "the layers' pools are not K/V views of one allocation" + elif not ctx_op.pick_split( + draft_model._fused_kv_weight.shape[0], + projected.shape[1], + draft_model._num_kv_heads, + self._draft_tokens_per_req, + n, + projected.device, + ): + reason = f"{n} context tokens (at most 64) or the K/V weight shape" + else: + self._k3_ctx_pool = view + ok = reason is None + if not ok: + logger.info(f"DFlash: context K/V on the Python path ({reason})") + if not torch.cuda.is_current_stream_capturing(): + cache[n] = ok + return ok + + def _k3_ctx_kv( + self, draft_model, projected, ctx_positions, num_accepted, slots, rows + ) -> torch.Tensor: + """``trtllm::k3_ctx_kv``: every drafter layer's context K/V of this step into the paged pool and + ``_ctx_len += num_accepted`` (clamped); returns the context length each request may advertise.""" + if self._k3_ctx_pool is None or self._k3_ctx_pool_bound != self._k3_ctx_pool_key(): + raise RuntimeError("k3_ctx_kv: the context pool was rebound after the path was chosen") + flat, layer_off, page_stride, kv_stride, head_stride = self._k3_ctx_pool + return torch.ops.trtllm.k3_ctx_kv( + projected.contiguous(), + draft_model._fused_kv_weight, + draft_model._k_norm_stacked, + draft_model._get_cos_sin_cache(), + ctx_positions, + num_accepted, + self._ctx_len, + slots, + rows, + self._ctx_block_tables, + self._ctx_block_counts, + flat, + layer_off, + page_stride, + kv_stride, + head_stride, + draft_model._k_norm_eps, + self._max_ctx, + self._ctx_page_size, + self._compute_block_size, + draft_model._num_kv_heads, + ) + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Called with this step's acceptance when ``k3_spec_accept`` produced it (a hook for drafter families).""" + + def _k3_accept(self, logits, attn_metadata, spec_metadata, draft_model, shard=None): + """``trtllm::k3_spec_accept`` for a step ``_k3_accept_applies`` accepted: the acceptance, the block table, + the KDA replay record, kv_lens + 1 and the drafter's inputs in one launch. ``shard``: the exchange + workspace and first column when ``logits`` are this rank's vocabulary shard (``_k3_logits_shard``). + Returns (accepted tokens, num accepted, {rewind, bonus, qpos, cpos, noise}).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + num_gens = attn_metadata.num_seqs + K = spec_metadata.runtime_draft_len + mgr = attn_metadata.kv_cache_manager + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + force = 0.0 if is_warmup else float(self.force_num_accepted_tokens) + mode, _, _ = accept_op.force_mode(force, K) + if mode == accept_op._kernel_module().FORCE_FRAC: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + else: + dummy_rng = getattr(self, "_k3_dummy_rng", None) + if dummy_rng is None: + dummy_rng = self._k3_dummy_rng = ( + torch.zeros(4, dtype=torch.float32, device=logits.device), + torch.zeros(2, dtype=torch.int64, device=logits.device), + ) + pool, counter = dummy_rng + embed = draft_model.draft_model_full.model.embed_tokens + accepted, num_acc, rewind, bonus, qpos, cpos, noise = torch.ops.trtllm.k3_spec_accept( + logits, + spec_metadata.draft_tokens.reshape(num_gens, K), + attn_metadata.draft_kv_cache_block_offsets, + self._ctx_pool_idx, + self._ctx_block_divisor, + self._ctx_block_counts, + self._ctx_block_tables, + mgr.prev_num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices, + mgr._dummy_request_mask, + attn_metadata.kv_lens_cuda, + self._batch_to_slot, + self._ctx_len, + self._max_ctx, + embed.weight, + self._k3_mask_row(draft_model), + pool, + counter, + force, + self._compute_block_size, + *self._k3_shard_args(shard), + ) + self._on_acceptance(accepted, num_acc, attn_metadata, spec_metadata) + return ( + accepted, + num_acc, + dict(rewind=rewind, bonus=bonus, qpos=qpos, cpos=cpos, noise=noise), + ) + + @staticmethod + def _k3_shard_args(shard) -> tuple: + """``trtllm::k3_spec_accept``'s exchange arguments for a vocabulary-sharded step (none otherwise).""" + if shard is None: + return () + ws, first_column = shard + return ( + ws["uc"], + ws["mc"], + ws["flags"], + ws["rank"], + ws["slots"], + ws["push_copies"], + first_column, + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + # Set when target_logits kept this step's logits vocabulary-sharded (only k3_spec_accept reads them). + step_shard, self._k3_step_shard = getattr(self, "_k3_step_shard", None), None + shard = None + if step_shard is not None: + if step_shard[0] is not logits: + raise RuntimeError( + "DFlash: the target logits are not the vocabulary shard target_logits returned" + ) + shard = step_shard[1] + + raw_logits = logits + K = spec_metadata.runtime_draft_len + + if K == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + # Lazy init buffers and attach worker reference for prepare() + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + # A decode step of the supported mode runs the acceptance and the drafter's inputs as one kernel + # (trtllm::k3_spec_accept), which also decodes the block table below. + if shard is None: + fused_accept = self._k3_accept_applies( + logits, attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ) + elif self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + fused_accept = True + else: + raise RuntimeError( + "DFlash: the target logits are vocabulary-sharded, but the fused acceptance no longer applies" + ) + # Before any store: prefill and decode both address pages through it. + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + if not fused_accept: + self._refresh_ctx_block_tables(attn_metadata, batch_size) + + # Save context lengths so both warmup and a failed forward can roll + # back the in-place _ctx_len updates made during drafting. + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # Never allocate the snapshot while capturing a CUDA graph: the + # clone would live in the graph memory pool, and its replay-time + # writes could alias blocks reused by later captures. Rollback is + # only meaningful for eager/warmup forwards anyway; a failure + # during capture aborts the graph itself, and captured ops do not + # mutate _ctx_len until replay. + self._saved_ctx_len = self._ctx_len.clone() + self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_req_ctx_pos = dict(self._req_ctx_pos) + self._ctx_len_restore_pending = True + + self._execute_guided_decoder_if_present(logits) + + prep = None + if fused_accept: + # Saved before the kernel updates kv_lens (warmup also snapshots kv_lens_cuda here). + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + accepted_tokens, num_accepted_tokens, prep = self._k3_accept( + logits, attn_metadata, spec_metadata, draft_model, shard + ) + else: + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + # Opt-in acceptance recording (env-gated; eager-mode measurement + # runs only). Skipped for CUDA-graph batches (capture/replay/warmup + # use synthetic requests and forbid the host sync). + if ( + self._accept_stats is not None + and num_gens > 0 + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + ): + self._accept_stats.on_accept( + spec_metadata.request_ids[num_contexts:batch_size], + num_accepted_tokens[num_contexts:batch_size].tolist(), + ) + + if fused_accept: + # The kernel recorded the KDA replay acceptance and advanced kv_lens_cuda. + self._kv_rewind_amount = prep["rewind"] + self._kv_rewind_nc = num_contexts + self._kv_rewind_bs = batch_size + self._kv_rewind_pending = True + attn_metadata.update_for_spec_dec() + else: + # Update GDN/Mamba recurrent states to the accepted token's state. + if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) + + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] + + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + ) + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + self._rebuild_batch_to_slot(spec_metadata.request_ids) + + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + prep=prep, + ) + + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + ctx_kv_cache=inputs["ctx_kv_cache"], + ctx_page_table=inputs["ctx_page_table"], + **dflash_ctx_rows_kwargs(draft_model, inputs["ctx_rows_start"]), + ) + + # Gather K logits per gen request from the block outputs. + # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. + block_size = self._compute_block_size + gen_hidden_states = self._draft_block_hidden_states( + draft_model, hidden_states_out, num_gens, block_size, K + ) + gen_logits = self._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) + + # DFlash 2: replace the independent per-position picks with one + # coherent path through the block (absent for plain DFlash). + if getattr(draft_model, "has_candidate_selector", False): + gen_logits = self._apply_dflash2_selector( + draft_model, + gen_logits, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + spec_metadata, + ) + vocab_size = gen_logits.shape[-1] + + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + ) + + # Opt-in confidence-calibration recording (env-gated). The + # confidence values come from an optional provider on the + # recorder (None here -> no calibration rows collected; the + # DSpark confidence-scheduled verification MR supplies the + # real provider, see accept_stats.py). Guarded off for + # CUDA-graph batches and d2t vocab-mapped drafters. + if ( + self._accept_stats is not None + and self._accept_stats.confidence_provider is not None + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + and self._d2t is None + ): + self._accept_stats.record_draft_confidence( + spec_metadata.request_ids[num_contexts:batch_size], + draft_model, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + gen_draft_tokens, + ) + + else: + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + gen_vocab = vocab_size if num_gens > 0 else None + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + self._rollback_guided_decoder_after_verify(num_accepted_tokens) + + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + # Restore context lengths after warmup; real runs keep the updates. + if is_warmup: + self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_host() + self._ctx_len_restore_pending = False + + outputs = { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + if shard is not None: + # "logits" is this rank's vocabulary shard: the engine gathers it for any logits post-processor. + outputs["logits_vocab_shard"] = True + return outputs + + def _draft_block_hidden_states( + self, + draft_model, + hidden_states_out: torch.Tensor, + num_gens: int, + block_size: int, + num_draft_tokens: int, + ) -> torch.Tensor: + """The block-output rows that produce the K draft logits per gen request (``_draft_slot_ids``). + + The slot ids depend only on the shapes, so they are built once per shape, outside CUDA-graph + capture, and cached. When they are one contiguous run of rows (a single request, or the + shift_label convention with a block of K slots, whose requests' slots follow each other) the + rows are returned as a view, without a gather; otherwise they are gathered. + """ + rows = hidden_states_out.shape[0] + key = (num_gens, block_size, num_draft_tokens, rows) + cache = self.__dict__.setdefault("_draft_block_rows", {}) + entry = cache.get(key) + if entry is None: + ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. + ids = ids.clamp(max=rows - 1) + if torch.cuda.is_current_stream_capturing(): + return hidden_states_out[ids] + host = ids.tolist() + if host and host == list(range(host[0], host[0] + len(host))): + entry = host[0] + else: + entry = ids + cache[key] = entry + if isinstance(entry, int): + return hidden_states_out[entry : entry + num_gens * num_draft_tokens] + return hidden_states_out[entry] + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata, + ) -> torch.Tensor: + """fp32 logits of the block positions over the full vocabulary (the TP lm_head all-gathers its + shards). A drafter family whose refinement and draft sampler reduce across ranks themselves may + return this rank's vocabulary shard instead.""" + return draft_model.logits_processor( + gen_hidden_states, draft_model.lm_head, attn_metadata, True + ) + + def _ctx_rows_start(self, num_contexts: int, num_gens: int) -> Optional[int]: + """The first row of the gen requests' page-table rows where those rows are one contiguous run, else None. + + The manager's block table is keyed by batch position, so the gen requests' rows are [num_contexts, + num_contexts + num_gens): a drafter can view them instead of gathering ``ctx_cache_batch_idx``'s rows. The + private arena's table is keyed by slot, which need not be contiguous. + """ + if num_gens > 0 and self._ctx_block_tables is not None: + return num_contexts + return None + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DFlashSpecMetadata, + draft_model: nn.Module, + total_target_tokens: int = 0, + prep: Optional[dict] = None, + ): + """Prepare inputs for DFlash's draft forward. + + ``prep``: ``trtllm::k3_spec_accept``'s bonus tokens, positions and noise embedding for a decode step it + took (see ``_k3_accept``); computed here otherwise. + + For gen requests, builds: + - noise_embedding: token embeddings for [accepted + mask] tokens + - query_positions: position IDs for the query tokens + - num_ctx_per_req: per-request context length in the pool + - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed + views of the persistent per-layer K/V pool. + - ctx_rows_start: where ctx_cache_batch_idx is the contiguous rows + [ctx_rows_start, ctx_rows_start + num_gens) of ctx_page_table, their + start (``_ctx_rows_start``), else None. + """ + num_contexts = attn_metadata.num_contexts + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + + mask_token_id = self._mask_token_id(draft_model) + + # Get the embed_tokens layer from the draft model + embed_tokens = draft_model.draft_model_full.model.embed_tokens + hidden_dim = ( + spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] + ) + + if num_gens > 0: + gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] + gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] + + total_tokens_per_req = self._draft_tokens_per_req # K+1 + K = spec_metadata.runtime_draft_len + + # Get captured multi-layer hidden states from spec_metadata + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + has_target_features = ( + captured_hs is not None + and hasattr(draft_model, "fc") + and hasattr(draft_model, "hidden_norm") + ) + + # Use cached block_size (resolved once on first call) + block_size = self._compute_block_size + query_tokens_per_req = block_size + + # Get slots for gen requests from pre-computed mapping + slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] + K_plus_1 = K + 1 + + if prep is not None: + gen_rows_out = self._k3_static_arange(num_contexts, num_contexts + num_gens) + offsets_kp1 = self._k3_static_arange(0, K_plus_1) + bonus = prep["bonus"] + query_position_ids = prep["qpos"] + ctx_position_ids = prep["cpos"] + noise_embed_2d = prep["noise"] + else: + gen_rows_out = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + + bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) + bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() + + ctx_len_gen = self._ctx_len[slots] + j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") + offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") + + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, + bonus, + mask_token_id, + query_tokens_per_req, + self._trained_mask_embedding(draft_model, mask_token_id), + ) + + # Accumulate new accepted features into context buffers + fused_ctx_kv = False + if has_target_features: + gen_start = attn_metadata.num_ctx_tokens + # Target now processes exactly K+1 tokens per gen req, so the + # captured slice is already the full set we need to project. + gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] + gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) + projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) + fused_ctx_kv = prep is not None and self._k3_ctx_kv_applies( + draft_model, projected_to_store + ) + if fused_ctx_kv: + num_ctx_fused = self._k3_ctx_kv( + draft_model, + projected_to_store, + ctx_position_ids, + gen_num_accepted, + slots, + gen_rows_out, + ) + # k3_ctx_kv adds the accepted tokens in the kernel; empty the dummy slot after + # it, as _advance_ctx_len does after the add on the torch path. + self._ctx_len[self._dummy_slot].zero_() + if has_target_features and not fused_ctx_kv: + gen_num_accepted_long = gen_num_accepted.long() + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) + if self._ctx_block_tables is not None: + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) + else: + col_idx = col_idx.clamp(max=self._max_ctx - 1) + + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) + col_flat = col_idx.reshape(-1) + proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] + pos_flat = ctx_position_ids.long().reshape(-1) + mask_1d = write_mask.reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + k_new, v_new = draft_model.precompute_context_kv( + proj_flat.to(cache_dtype), pos_flat + ) + mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + slot_long = slot_flat.long() + col_long = col_flat.long() + if self._ctx_paged: + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + self._advance_ctx_len(slots, gen_num_accepted_long) + + if fused_ctx_kv: + num_ctx_per_req_t = num_ctx_fused + else: + num_ctx_per_req_t = self._ctx_len[slots] + if not fused_ctx_kv and self._ctx_block_tables is not None: + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + noise_embedding = noise_embed_2d + query_positions = query_position_ids.long() + + # Update seq_lens for gen requests to K+1 (the number of tokens + # the target forward actually processed). + attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( + total_tokens_per_req + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req + else: + noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) + query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") + num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") + slots = torch.empty(0, dtype=torch.long, device="cuda") + gen_rows_out = slots + bonus = torch.empty(0, dtype=torch.long, device="cuda") + + return { + "noise_embedding": noise_embedding, + "query_positions": query_positions, + "num_ctx_per_req": num_ctx_per_req_t, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": bonus, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), + "ctx_rows_start": self._ctx_rows_start(num_contexts, num_gens), + } + + +class KimiK3DSparkWorker(KimiK3DFlashWorker, DSparkWorker): + # Whether this step's block logits are the Kimi K3 decode path's vocabulary shard (see ``_draft_block_logits``). + _k3_sharded_block_logits = False + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Block logits the Kimi K3 decode path kept vocab-sharded (see ``_draft_block_logits``) run the + whole chain as ``trtllm::k3_markov``, which also returns the greedy draft tokens and + next_new_tokens. Other TP-sharded logits slice markov_w2's rows to this rank's contiguous + shard and chain through the TP-aware global argmax. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # Unlike the d2t guard this one cannot move to set_draft_model: it + # keys on the runtime logits width, and reproducing that at init would + # duplicate the draft head's sharding rules. A standalone drafter + # borrows the target lm_head, whose gather_output defaults to True, so + # the logits normally arrive full-vocab and this branch is skipped. + self._k3_markov = None + k3_sharded, self._k3_sharded_block_logits = self._k3_sharded_block_logits, False + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + if k3_sharded: + return self._k3_markov_chain( + draft_model, gen_logits, first_prev_tokens, vocab_slice + ) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) + + def _keep_draft_logits_sharded(self, draft_model, spec_metadata, num_gens: int) -> bool: + """Keep the draft logits vocab-sharded: ``trtllm::k3_markov`` reduces across the ranks itself. + + The chain's global argmaxes are the greedy draft tokens, so the lm_head's all-gather of the + full logits and the draft sampler's gather are both skipped. Needs a Kimi K3 target running its + decode kernels (``k3_decode``), plain TP, greedy drafting (rejection sampling reads + full-vocabulary probabilities), an unquantized bias-free column-parallel bf16 lm_head whose + shards tile the Markov vocabulary, bf16 Markov weights of the kernel's rank, MNNVL (the model's + all-reduces own an MNNVL workspace for this mapping), and a block, shard and batch of + ``num_gens`` requests the kernel can split. + """ + if not self.k3_decode: + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + lm_head = getattr(draft_model, "lm_head", None) + if ( + lm_head is None + or not getattr(draft_model, "has_markov_head", False) + or mapping is None + or mapping.tp_size <= 1 + or mapping.enable_attention_dp + or spec_metadata.wants_advanced_draft_sampling + or getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "bias", None) is not None + or lm_head.weight.dtype != torch.bfloat16 + or draft_model.markov_w1.dtype != torch.bfloat16 + or draft_model.markov_w2.dtype != torch.bfloat16 + or lm_head.weight.shape[0] * mapping.tp_size != draft_model.markov_w2.shape[0] + or MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None + ): + return False + shard = lm_head.weight.shape[0] + block = spec_metadata.runtime_draft_len + rank = k3_markov_op._kernel_module().MARKOV_RANK + if ( + not 0 < block <= k3_markov_op.WORKSPACE_MAX_BLOCK + or draft_model.markov_w1.dim() != 2 + or draft_model.markov_w1.shape[1] != rank + or tuple(draft_model.markov_w2.shape[1:]) != (rank,) + ): + return False + if k3_markov_op.pick_grid(shard, block, num_gens) == 0: + logger.warning_once( + f"DSpark Markov head: trtllm::k3_markov cannot split a {shard}-row vocab shard for " + f"{num_gens} requests; the draft logits are all-gathered and the chain runs unfused.", + key=f"dspark_k3_markov_shard_{num_gens}", + ) + return False + return True + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + ) -> torch.Tensor: + """This rank's bf16 shard of the block logits when ``k3_markov`` takes them (it converts them to fp32 + exactly, as ``.float()`` would), from the drafter's logits processor's head kernel where it takes the rows, + else the head's own GEMM (``_k3_head_shard``); otherwise the base class' fp32 logits.""" + num_gens = gen_hidden_states.shape[0] // max(spec_metadata.runtime_draft_len, 1) + self._k3_sharded_block_logits = self._keep_draft_logits_sharded( + draft_model, spec_metadata, num_gens + ) + if self._k3_sharded_block_logits: + return self._k3_head_shard( + draft_model.logits_processor, draft_model.lm_head, gen_hidden_states + ) + return super()._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): + """The base acceptance; under ``k3_decode`` its outputs are kept for ``k3_markov``'s next_new_tokens.""" + accepted_tokens, num_accepted_tokens = super().sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + if self.k3_decode: + self._on_acceptance(accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata) + return accepted_tokens, num_accepted_tokens + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Keeps this step's acceptance for ``k3_markov``'s next_new_tokens (and the metadata whose KV lengths it + rewinds).""" + self._k3_acceptance = ( + accepted_tokens, num_accepted_tokens, attn_metadata.num_contexts, spec_metadata, attn_metadata + ) # fmt: skip + + def _k3_markov_chain( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + vocab_slice: slice, + ) -> torch.Tensor: + """The Markov chain as ``trtllm::k3_markov`` on this rank's shard of the block logits. + + Returns the corrected logits (fp32); keeps the kernel's greedy tokens for + ``sample_draft_tokens`` and its next_new_tokens (built from this step's acceptance) for + ``_prepare_next_new_tokens``. + """ + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + + acceptance = getattr(self, "_k3_acceptance", None) + if acceptance is None: + raise RuntimeError("DSpark Markov head: k3_markov needs this step's acceptance first") + accepted_tokens, num_accepted_tokens, num_contexts, spec_metadata, attn_metadata = ( + acceptance + ) + num_gens = gen_logits.shape[0] + # The draft forward's pending KV-length rewind (see _apply_kv_rewind_after_draft) runs in the kernel: every + # reader of kv_lens_cuda in the draft forward precedes it. In warmup the restore of kv_lens_cuda that + # follows overwrites it, as it would the Python rewind. + rewind = getattr(self, "_kv_rewind_amount", None) + fold = ( + getattr(self, "_kv_rewind_pending", False) + and rewind is not None + and getattr(attn_metadata, "kv_lens_cuda", None) is not None + and rewind.dtype == torch.int32 + and self._kv_rewind_bs - self._kv_rewind_nc == rewind.numel() <= num_gens + ) + corrected, tokens, next_new = k3_markov_op.markov_chain( + self.mapping, + gen_logits.contiguous(), + first_prev_tokens.long(), + draft_model.markov_w1, + draft_model.markov_w2[vocab_slice], + vocab_slice.start, + accepted_tokens, + num_accepted_tokens[num_contexts : num_contexts + num_gens], + spec_metadata.batch_indices_cuda[num_contexts : num_contexts + num_gens], + kv_lens=attn_metadata.kv_lens_cuda if fold else None, + rewind=rewind if fold else None, + rewind_first=self._kv_rewind_nc if fold else 0, + ) + if fold: + self._kv_rewind_amount = None + self._kv_rewind_pending = False + self._k3_markov = (corrected, tokens, next_new, accepted_tokens, num_accepted_tokens) + return corrected + + def sample_draft_tokens( + self, + logits, + spec_metadata, + batch_size, + *, + num_contexts=0, + draft_step=None, + mapping_lm_head_tp=None, + ): + """Greedy block drafts straight from ``k3_markov``. + + When ``logits`` are the corrected logits the kernel just returned, its per-position global + argmax (first maximum, lowest vocabulary index among equal values) is the token the + TP-gathered greedy sampler would pick. Anything else goes to the base sampler. + """ + chain = getattr(self, "_k3_markov", None) + if ( + chain is not None + and chain[0] is logits + and mapping_lm_head_tp is None + and not spec_metadata.wants_advanced_draft_sampling + ): + self._k3_markov_next = (chain[1], chain[2], chain[3], chain[4]) + return chain[1] + self._k3_markov_next = None + return super().sample_draft_tokens( + logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + draft_step=draft_step, + mapping_lm_head_tp=mapping_lm_head_tp, + ) + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """``k3_markov``'s next_new_tokens when the drafts are its tokens for the whole batch (no context + requests); otherwise the base assembly.""" + chain = getattr(self, "_k3_markov_next", None) + # The step ends here: release its acceptance (which holds the step's attention and spec metadata) and the + # kernel's outputs, so they do not keep a captured graph's pool or metadata alive after the step. + self._k3_markov_next = None + self._k3_acceptance = None + self._k3_markov = None + if ( + chain is not None + and chain[0] is next_draft_tokens + and chain[2] is accepted_tokens + and chain[3] is num_accepted_tokens + and next_draft_tokens.shape[0] == batch_size + ): + return chain[1] + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py new file mode 100644 index 000000000000..950a29442166 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight loading: Kimi K3 (NVFP4) / sm_107 / tp16_moetp4ep4. + +The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and +`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: + +* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the + routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under + moe_tp 4 x moe_ep 4), and checks its own key coverage. +* The vision tower and the projector are a predicted non-load: listed here, never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_utils import filter_weights + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") + + +def load(model, weights) -> None: + """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) + assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = _LANG_PREFIX + KimiLinearForCausalLM.load_weights(model, lm_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py index 854616aebd9f..09abb6fd26f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py @@ -20,7 +20,7 @@ from ..._router_index import NoTarget -# The one GPU architecture these targets are written for. sm is part of a +# The GPU architecture the MXFP4 targets are written for. sm is part of a # target's identity, not a knob: a different SM is a different target. The K3 # decode kernels (the tcgen05 GEMVs, the fused MoE) are built for GB200 only. _SM = (10, 0) @@ -64,6 +64,35 @@ def _quant_algo(pretrained_config: Any) -> Optional[QuantAlgo]: return quant_config.quant_algo +def _has_nvfp4_experts(pretrained_config: Any) -> bool: + """Require the mixed-precision checkpoint's explicit NVFP4 expert declarations.""" + from tensorrt_llm._torch.model_config import ModelConfig + + declared = getattr(pretrained_config, "quantization_config", None) + if declared is None: + return False + try: + global_quant, layers = ModelConfig.load_hf_quant_config(declared, "TRTLLM") + except ValueError: + return False + if global_quant.quant_algo != QuantAlgo.MIXED_PRECISION or not layers: + return False + for layer_idx in range(1, 93): + names = tuple( + f"{prefix}layers.{layer_idx}.{module}.experts" + for prefix in ("language_model.model.", "model.", "") + for module in ("block_sparse_moe", "mlp") + ) + if any(global_quant.is_module_excluded_from_quantization(name) for name in names): + return False + matches = [layers[name] for name in names if name in layers] + if not matches or any( + quant.quant_algo != QuantAlgo.NVFP4 or quant.group_size != 16 for quant in matches + ): + return False + return True + + def _parallel(m) -> Optional[str]: """Name the parallel topology, or None if no target implements it. @@ -120,6 +149,18 @@ def route( compute_capability == _SM and shape == _SHAPE and quant_algo == QuantAlgo.W4A16_MXFP4 ) + if ( + compute_capability == (10, 7) + and shape == _SHAPE + and parallel == "tp16_moetp4ep4" + and _has_nvfp4_experts(pretrained_config) + ): + from .kimi_k3_nvfp4__sm_107__tp16_moetp4ep4.modeling import ( + ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4, + ) + + return ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4 + if mxfp4_on_sm100 and parallel == "tp16_moetp4ep4": from .kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4, @@ -142,5 +183,6 @@ def route( f"split_set={getattr(mapping, 'moe_tp_ep_user_specified', False)} " f"attention_dp={mapping.enable_attention_dp}; the kimi_k3_vl targets are the MXFP4 " "Kimi K3 checkpoint on sm100 (GB200), 16 ranks of tensor-parallel attention, with the " - "experts split moe_tp 4 x moe_ep 4 or, set explicitly, moe_tp 16 x moe_ep 1" + "experts split moe_tp 4 x moe_ep 4 or, set explicitly, moe_tp 16 x moe_ep 1; " + "and group-size-16 NVFP4 Kimi K3 on sm107 with tp16 attention and moe_tp 4 x moe_ep 4" ) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 9e26eac92caa..7dc71313211a 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -3110,7 +3110,14 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: else: attention = layer.self_attn.mixer for name, linear in attention.named_children(): - if isinstance(linear, TrtllmLinear) and linear.has_fp8_block_scales: + if ( + isinstance(linear, TrtllmLinear) + and linear.has_fp8_block_scales + and not linear.use_cute_dsl_blockscaling_mm + and not linear.disable_deep_gemm + ): + # The weight-read specialization always uses DeepGEMM. + # Preserve Linear when another GEMM backend was requested. _swap_linear_to_fp8_weight_read(attention, name, (TrtllmLinear,)) fused = getattr(attention, "qkvg_proj", None) if isinstance(fused, _Fp8BlockScaleWeightReadLinear): diff --git a/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py b/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py index 077d1425121b..c3f84b894b0b 100644 --- a/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py +++ b/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py @@ -229,7 +229,7 @@ def _model( quantized=None, alias="language_model.model.", excluded=None, - cute_mm=True, + cute_mm=False, q_lora_rank=128, ): cfg = SimpleNamespace( @@ -609,3 +609,34 @@ def test_kda_small_checkpoint_fp8_projections_load_as_bf16(tp_size, tp_rank): assert isinstance(module, nn.Linear) assert module.weight.dtype == torch.bfloat16 torch.testing.assert_close(module.weight, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)), + reason="Kimi CuTe DSL checkpoint tests require SM100/SM103/SM107", +) +@pytest.mark.parametrize("kda", [True, False]) +def test_checkpoint_finalization_preserves_cute_dsl_backend(kda): + model = _model(kda=kda, cute_mm=True) + checkpoint = _checkpoint(model) + model.load_weights(checkpoint) + layer = model.model.layers[0] + attention = layer.linear_attn if kda else layer.self_attn.mixer + names = _KDA_PROJECTIONS[:5] if kda else _MLA_PROJECTIONS + for name in names: + module, _ = _projection(attention, name) + assert isinstance(module, Linear) + assert module.use_cute_dsl_blockscaling_mm + for module in model.modules(): + if isinstance(module, Linear): + module.post_load_weights() + hidden = torch.randn(8, 256, dtype=torch.bfloat16, device="cuda") + if kda: + expected = torch.cat([getattr(attention, name)(hidden) for name in names[:4]], dim=-1) + assert torch.isfinite(expected).all() + if attention.qkvg_proj is not None: + torch.testing.assert_close(attention.qkvg_proj(hidden), expected, rtol=0.02, atol=0.02) + else: + assert torch.isfinite(attention.kv_a_proj_with_mqa(hidden)).all() + assert torch.isfinite(attention.kv_b_proj(hidden)).all() diff --git a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py index 41976eaef202..d3b2be6a1877 100644 --- a/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py +++ b/tests/unittest/_torch/modeling_v2/activation/test_modeling_v2_k3_situ_mul.py @@ -14,9 +14,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul assert torch.cuda.is_available(), "k3_situ_mul requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py index 3a1657098d1b..359358609efc 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_drafter_attn_qknorm.py @@ -7,9 +7,9 @@ assert torch.cuda.is_available(), "k3_drafter_attn requires a CUDA device" -if torch.cuda.get_device_capability() != (10, 0): - # The CuTe DSL kernel uses tcgen05, TMA and clusters; it is certified on sm_100 (B200 / GB200) only. - pytest.skip("k3_drafter_attn is certified on sm_100 only", allow_module_level=True) +if torch.cuda.get_device_capability() not in ((10, 0), (10, 7)): + # The CuTe DSL kernel uses tcgen05, TMA and clusters on the K3 target architectures. + pytest.skip("k3_drafter_attn requires SM100 or SM107", allow_module_level=True) from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( # noqa: E402 k3_drafter_attn_qknorm, diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py index 4befd44bf4a4..8b57695fc4b2 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv.py @@ -14,9 +14,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv assert torch.cuda.is_available(), "k3_ctm_gemv requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py index b842a80f0c5a..0702b0e61adb 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_long.py @@ -17,9 +17,9 @@ ) assert torch.cuda.is_available(), "k3_ctm_gemv_long requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py index e61be03e4683..4f63eb93ed27 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_swiglu.py @@ -17,9 +17,9 @@ ) assert torch.cuda.is_available(), "k3_ctm_gemv_swiglu requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) M_ALL = range(1, 9) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py index af81feb06f7b..e6801fc031b3 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_ctm_gemv_wide.py @@ -20,9 +20,9 @@ ) assert torch.cuda.is_available(), "k3_ctm_gemv_wide requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) M_ALL = range(1, 65) diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py index b21f4c21c4ce..0bc5f80cc836 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_decode_gemv.py @@ -15,9 +15,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv assert torch.cuda.is_available(), "k3_decode_gemv requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) TOL = 8e-3 # max |y - ref| / max |ref|, ref an fp64 product diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py index e3e82157fdfc..4150878a0937 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_k3_head_gemv.py @@ -21,9 +21,9 @@ ) assert torch.cuda.is_available(), "k3_head_gemv requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) TOL = 8e-3 # max |y - ref| / max |ref|, ref an fp64 product diff --git a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py index 89b1016e6d70..97e235fdf411 100644 --- a/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py +++ b/tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_k3_embed_norm.py @@ -18,9 +18,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm assert torch.cuda.is_available(), "k3_embed_norm requires a CUDA device" -# The receipts are sm_100 ones; CI's other architectures skip this file. +# Exercise the original SM100 target and the NVFP4 SM107 target with the same references. pytestmark = pytest.mark.skipif( - torch.cuda.get_device_capability() != (10, 0), reason="certified on sm_100 only" + torch.cuda.get_device_capability() not in ((10, 0), (10, 7)), reason="requires SM100 or SM107" ) VOCAB, HIDDEN = 163840, 7168 diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index ae2731dbe63e..ca11a8d2f831 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -460,6 +460,211 @@ def test_kimi_k3_on_another_sm_does_not_match(): modeling_v2_resolve(config, _llm_args("require")) +def _k3_nvfp4_mixed_config(): + """ModelOpt's expert declarations, alongside FP8 attention projections.""" + return { + "quant_method": "modelopt", + "quantization": { + "quant_algo": "MIXED_PRECISION", + "group_size": 16, + "quantized_layers": { + **{ + f"language_model.model.layers.{layer}.mlp.experts": { + "quant_algo": "NVFP4", + "group_size": 16, + } + for layer in range(1, 93) + }, + "language_model.model.layers.1.self_attn.q_proj": {"quant_algo": "FP8_PB_WO"}, + }, + }, + } + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize("quant_method", ["modelopt", "modelopt_mixed"]) +def test_kimi_k3_nvfp4_routes_on_sm107(monkeypatch, mode, quant_method): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) + quantization = _k3_nvfp4_mixed_config() + if quant_method == "modelopt_mixed": + quantization = dict(quantization["quantization"], quant_method=quant_method) + config = _model_config(_k3_config(quantization=quantization), **_TP16_MOETP4EP4) + assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == ( + "models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4.modeling." + "ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4" + ) + + +@pytest.mark.parametrize("bad_layer", ["missing", "mxfp4", "wrong_group_size"]) +def test_kimi_k3_nvfp4_rejects_incomplete_expert_declarations(monkeypatch, bad_layer): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) + quantization = _k3_nvfp4_mixed_config() + layers = quantization["quantization"]["quantized_layers"] + key = "language_model.model.layers.92.mlp.experts" + if bad_layer == "missing": + del layers[key] + elif bad_layer == "mxfp4": + layers[key] = {"quant_algo": "W4A8_MXFP4_MXFP8", "group_size": 32} + else: + layers[key]["group_size"] = 32 + config = _model_config(_k3_config(quantization=quantization), **_TP16_MOETP4EP4) + with pytest.raises(ValueError): + modeling_v2_resolve(config, _llm_args("require")) + + +def test_kimi_k3_nvfp4_rejects_uniformly_quantized_attention(monkeypatch): + """NVFP4 attention weights are distinct from the supported FP8 attention checkpoint.""" + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) + config = _model_config(_k3_config(quantization=_K3_NVFP4_QUANTIZATION), **_TP16_MOETP4EP4) + with pytest.raises(ValueError, match="no target fits"): + modeling_v2_resolve(config, _llm_args("require")) + + +def test_kimi_k3_nvfp4_declines_excluded_experts(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) + quantization = _k3_nvfp4_mixed_config() + quantization["quantization"]["exclude_modules"] = ["*.layers.92.mlp.experts"] + config = _model_config(_k3_config(quantization=quantization), **_TP16_MOETP4EP4) + assert modeling_v2_resolve(config, _llm_args("auto")) is None + with pytest.raises(ValueError, match="no target fits"): + modeling_v2_resolve(config, _llm_args("require")) + + +def test_kimi_k3_nvfp4_constructor_uses_per_layer_quantization(): + import importlib + from types import SimpleNamespace + + target = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + "kimi_k3_nvfp4__sm_107__tp16_moetp4ep4.modeling" + ) + moe_runtime = target.KimiK3MoERuntime + + quant, layers = ModelConfig.load_hf_quant_config(_k3_nvfp4_mixed_config(), "TRTLLM") + config = SimpleNamespace(quant_config=quant, quant_config_dict=layers) + for layer_idx in (1, 92): + resolved = moe_runtime._resolve_routed_quant_config(config, layer_idx) + assert resolved.quant_algo == QuantAlgo.NVFP4 + assert resolved.group_size == 16 + del layers["language_model.model.layers.92.mlp.experts"] + with pytest.raises(ValueError, match="Missing NVFP4"): + moe_runtime._resolve_routed_quant_config(config, 92) + + +def test_kimi_k3_nvfp4_keeps_mla_dspark_worker(): + """The reference MLA drafter must retain its compressed-context-cache worker.""" + import importlib + from types import SimpleNamespace + + target = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + "kimi_k3_nvfp4__sm_107__tp16_moetp4ep4.modeling" + ) + drafter, worker = object(), object() + model = SimpleNamespace( + draft_config=SimpleNamespace( + pretrained_config=PretrainedConfig(kv_lora_rank=512, qk_rope_head_dim=64) + ), + draft_model=drafter, + spec_worker=worker, + epilogue=[drafter, worker], + ) + config = SimpleNamespace( + spec_config=SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dspark=lambda: True), + draft_is_embedded_in_target=False, + ) + ) + target.ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4._own_spec_branch(model, config) + assert model.draft_model is drafter + assert model.spec_worker is worker + assert model.epilogue == [drafter, worker] + + +@pytest.mark.parametrize("prefix", ["language_model.model.", "model.", ""]) +def test_kimi_k3_nvfp4_attention_quantization_preserves_config(prefix: str) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( + modeling as target, + ) + from tensorrt_llm.models.modeling_utils import QuantConfig + + projection = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) + global_quant = QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=QuantAlgo.NVFP4 + ) + config = ModelConfig( + quant_config=global_quant, + quant_config_dict={f"{prefix}layers.3.self_attn.q_a_proj": projection}, + ) + resolved = target.resolve_attention_quant_config(config, 3, "q_a_proj") + assert resolved.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + assert resolved.group_size == 128 + assert resolved.kv_cache_quant_algo == QuantAlgo.NVFP4 + assert projection.kv_cache_quant_algo is None + assert target.resolve_attention_quant_config(config, 3, "g_proj").quant_algo is None + + +def test_kimi_k3_nvfp4_attention_honors_exclusions() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( + modeling as target, + ) + from tensorrt_llm.models.modeling_utils import QuantConfig + + config = ModelConfig( + quant_config=QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, + kv_cache_quant_algo=QuantAlgo.NVFP4, + exclude_modules=["*.g_proj"], + ), + quant_config_dict={ + "model.layers.3.self_attn.g_proj": QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES) + }, + ) + resolved = target.resolve_attention_quant_config(config, 3, "g_proj") + assert resolved.quant_algo is None + assert resolved.kv_cache_quant_algo == QuantAlgo.NVFP4 + + +@pytest.mark.parametrize( + "algo,group_size", [(QuantAlgo.NVFP4, 16), (QuantAlgo.FP8_BLOCK_SCALES, 64)] +) +def test_kimi_k3_nvfp4_attention_rejects_unsupported_quantization( + algo: QuantAlgo, group_size: int +) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( + modeling as target, + ) + from tensorrt_llm.models.modeling_utils import QuantConfig + + config = ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), + quant_config_dict={ + "model.layers.3.self_attn.q_a_proj": QuantConfig(quant_algo=algo, group_size=group_size) + }, + ) + with pytest.raises(ValueError, match="quantization|128x128"): + target.resolve_attention_quant_config(config, 3, "q_a_proj") + + +def test_kimi_k3_nvfp4_attention_rejects_conflicting_block_sizes() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( + modeling as target, + ) + from tensorrt_llm.models.modeling_utils import QuantConfig + + config = ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), + quant_config_dict={ + f"{prefix}layers.3.self_attn.q_a_proj": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=group_size + ) + for prefix, group_size in (("language_model.model.", 128), ("model.", 64)) + }, + ) + with pytest.raises(ValueError, match="Conflicting"): + target.resolve_attention_quant_config(config, 3, "q_a_proj") + + # --- qwen4_exp: one-model MTP-Eagle over a linear chain, one rank --------------- _QWEN4_EXP_TARGET = "models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.modeling." diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 971087181b0a..d486111e565f 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -1964,6 +1964,93 @@ def test_nvfp4_experts_match_situ_reference(moe_backend, input_scale): assert cosine > 0.95, f"cosine={cosine.item()}, rel_l2={rel_l2.item()}" +@nvfp4_moe_supported +@pytest.mark.parametrize("num_tokens", [1, 8, 33]) +def test_modeling_v2_nvfp4_streamed_moe_matches_reference(num_tokens: int) -> None: + """Exercise the target's loader and router without allocating the full K3 model.""" + from transformers import PretrainedConfig + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( + modeling as target, + ) + from tensorrt_llm.models.modeling_utils import QuantConfig + + if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("The NVFP4 modeling_v2 target requires SM107") + torch.manual_seed(91) + cfg = _K3Config( + hidden_size=_TP_HIDDEN, + num_experts=_TP_EXPERTS, + moe_intermediate_size=_TP_INTERMEDIATE, + routed_expert_hidden_size=_TP_HIDDEN, + num_shared_experts=1, + latent_moe_use_norm=True, + ) + pretrained = PretrainedConfig(**dataclasses.asdict(cfg), torch_dtype=torch.bfloat16) + config = ModelConfig( + pretrained_config=pretrained, + mapping=Mapping(), + moe_backend="TRTLLM", + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION), + quant_config_dict={ + "layers.1.mlp.experts": QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) + }, + ) + runtime = target.KimiK3MoERuntime( + config, + cfg, + layer_idx=1, + aux_stream_dict={kind: torch.cuda.Stream() for kind in AuxStreamType}, + ).cuda() + gate = runtime.gate.float() + with torch.no_grad(): + gate.weight.normal_(std=0.05) + gate.e_score_correction_bias.normal_(std=0.1) + x = torch.randn(num_tokens, _TP_HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.5 + bank = [] + for _ in range(_TP_EXPERTS): + w1, w3 = [ + torch.randn(_TP_INTERMEDIATE, _TP_HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.05 + for _ in range(2) + ] + w2 = torch.randn_like(w1.t()) * 0.05 + bank.append(_quantize_expert_to_nvfp4(w1, w2.contiguous(), w3, 1.0)) + + streamed = runtime.routed_experts + # A full model allocates mixed-precision weights in its post-init hook. + streamed.create_weights() + streamed.cuda() + backend = streamed.backend + assert "trtllm_gen" in type(backend).__module__ + backend.quant_method.prepare_streaming_expert_load(backend) + spec = runtime.expert_ckpt_spec + checkpoint = { + f"experts.{expert}.{name}": tensor + for expert, tensors in enumerate(bank) + for name, tensor in tensors.items() + } + for expert in range(_TP_EXPERTS): + spec.loader(backend, "experts", expert, expert, checkpoint.__getitem__) + assert getattr(backend, spec.loaded_slots_attr) == set(range(_TP_EXPERTS)) + backend.process_weights_after_loading() + streamed.post_load_weights() + + whole = _make_nvfp4_moe(gate, moe_backend="TRTLLM") + _load_nvfp4_bank_for(whole, bank, "TRTLLM") + logits = gate.compute_logits(x) + actual = streamed(x, logits).float() + torch.testing.assert_close(actual, whole(x, logits).float(), rtol=0, atol=0) + + dequantized = [_dequantize_expert_bank(expert) for expert in bank] + w1, w2, w3 = map(list, zip(*dequantized)) + expected = _situ_reference_moe( + x, logits, gate.routing_method, w1, w2, w3, beta=4.0, linear_beta=25.0 + ) + relative_l2 = torch.linalg.vector_norm(actual - expected) / torch.linalg.vector_norm(expected) + assert torch.isfinite(actual).all() + assert relative_l2 < _SITU_NVFP4_REL_L2_MAX, relative_l2.item() + + def _plain_swiglu_reference_moe(x, router_logits, routing_method, w1, w2, w3): """Same routing/geometry as the SiTU reference, but plain SwiGLU. From 7cfaa4dd93d4c1c9166d6f8db2fcafb8a26a15b6 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:36:20 +0800 Subject: [PATCH 2/8] Share Kimi K3 modeling across FP4 formats and support SM103 (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 15 +- .../decode_comm.py | 0 .../decode_gemv.py | 0 .../decode_moe.py | 0 .../decode_moe.py => decode_moe_tp16.py} | 48 +- .../__init__.py | 2 +- .../modeling.py | 3529 +--------------- .../weights.py | 44 +- .../__init__.py | 2 +- .../decode_comm.py | 402 -- .../decode_gemv.py | 465 --- .../modeling.py | 3592 +--------------- .../spec_worker.py | 1343 ------ .../weights.py | 42 +- .../__init__.py | 3 + .../modeling.py | 21 + .../weights.py | 6 + .../__init__.py | 3 + .../modeling.py | 21 + .../weights.py | 6 + .../__init__.py | 3 + .../modeling.py | 21 + .../weights.py | 6 + .../__init__.py | 2 +- .../decode_comm.py | 402 -- .../decode_gemv.py | 467 --- .../modeling.py | 3009 +------------- .../spec_worker.py | 1343 ------ .../weights.py | 42 +- .../modeling_v2/models/kimi_k3_vl/modeling.py | 3694 +++++++++++++++++ .../modeling_v2/models/kimi_k3_vl/routing.py | 45 +- .../spec_worker.py | 0 .../modeling_v2/models/kimi_k3_vl/weights.py | 42 + .../cute_dsl_kernels/k3_fused_moe/op.py | 14 +- .../integration/test_lists/test-db/l0_cpu.yml | 2 + .../comm/_kimi_k3_decode_comm_op_matrix.py | 3 +- .../_kimi_k3_route_b_decode_moe_op_matrix.py | 7 +- .../modeling_v2/test_modeling_v2_claims.py | 21 +- .../test_modeling_v2_kimi_k3_construction.py | 46 +- .../test_modeling_v2_kimi_k3_decode_gemv.py | 2 +- .../test_modeling_v2_kimi_k3_decode_step.py | 12 +- .../test_modeling_v2_kimi_k3_drafter.py | 4 +- .../test_modeling_v2_kimi_k3_drift.py | 209 +- .../test_modeling_v2_kimi_k3_moe_routing.py | 11 +- .../test_modeling_v2_kimi_k3_route_b_moe.py | 4 +- .../modeling_v2/test_modeling_v2_routing.py | 74 +- .../_torch/moe/test_kimi_k3_situ_moe.py | 8 +- .../hw_agnostic/test_k3_decode_worker.py | 2 +- .../hw_agnostic/test_kimi_k3_capture_view.py | 2 +- .../test_kimi_k3_ctx_rows_start.py | 4 +- .../test_kimi_k3_draft_block_rows.py | 2 +- .../hw_agnostic/test_kimi_k3_drafter_comm.py | 12 +- .../test_kimi_k3_spec_worker_gate.py | 10 +- 53 files changed, 4159 insertions(+), 14910 deletions(-) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/{kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 => }/decode_comm.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/{kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 => }/decode_gemv.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/{kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 => }/decode_moe.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/{kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py => decode_moe_tp16.py} (94%) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/modeling.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/{kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 => }/spec_worker.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/weights.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 7dc176ff06da..6aba851fa95d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -1,7 +1,6 @@ # ModelingV2 -One self-contained modeling codebase per deployment target, beside the -built-in model zoo rather than inside it. +Explicit deployment targets beside the built-in model zoo. Where `_torch/models/modeling_deepseekv3.py` is one class serving V3, V3-Lite, R1 and V3.2 across every GPU generation and parallel topology, @@ -11,6 +10,18 @@ entries, sharing nothing with its siblings, and trusted through accuracy gates instead of shared abstractions. The one-to-one correspondence between `models//` here and `modeling_.py` there is the point of the exercise. +Kimi K3 keeps concrete checkpoint, SM and parallelism identities in thin target +classes and shares its implementation within `models/kimi_k3_vl/`. Its MXFP4 and +NVFP4 targets use the same decoder, attention residuals, BF16 GEMVs, speculative +worker and checkpoint streaming lifecycle. The expert checkpoint specification +selects the loader and finalization behavior. MXFP4 uses the K3 MXFP4/MXFP8 fused +decode kernels; NVFP4 uses the existing NVFP4 MoE kernels and stock quantized +attention, preserving FP8 projection metadata and the configured KV-cache format. +The MXFP4 fusion's group-32 E8M0 scales and MXFP8 intermediate differ from NVFP4's +group-16 scales and FP4 activation quantization, so the formats keep independent +MoE kernels. Shared model code does not imply shared accuracy or performance +certification: every hardware and quantization target needs its own gates. + ## Using it ```python diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_comm.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_comm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_comm.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_moe.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe_tp16.py similarity index 94% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe_tp16.py index e39c3ddeb0e7..b797c4fe6108 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe_tp16.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# >>> route B: this target's MoE engines (every expert on each rank) +# This target's MoE engines (every expert on each rank). """The decode path's latent MoE on the catalog's Kimi K3 MoE entries, with every expert on each rank (moe TP16 x EP1). **At most 8 tokens** (`MAX_TOKENS`), a MoE layer runs: @@ -41,7 +41,6 @@ launches without programmatic dependent launch), which the decode kernels a pushing step runs do; pushes run only in CUDA-graph replays, so no other step's kernels sit between two pushes. """ -# <<< route B from __future__ import annotations @@ -76,7 +75,7 @@ weight_supported, ) -# >>> route B: the one- and two-token engines +# The one- and two-token engines. from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe_m1 import ( K3MoeM1Layer, K3MoeM1State, @@ -89,8 +88,6 @@ k3_moe_m2, k3_moe_m2_push, ) - -# <<< route B from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_route_quant import k3_route_quant from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel @@ -104,34 +101,30 @@ # (comm/k3_sandwich_tail takes the TP16 tail weight [7168, 256 + 384]). TAIL_K_TILE = 256 -# >>> route B: the routed experts' SiTU caps (activation_situ_beta, activation_situ_linear_beta) that k3_moe_m1, +# The routed experts' SiTU caps (activation_situ_beta, activation_situ_linear_beta) that k3_moe_m1,. # k3_moe_m2 and k3_moe compile in (their kernels' SITU_GATE_CAP and SITU_LINEAR_CAP) ENGINE_SITU_CAPS = (4.0, 25.0) -# <<< route B @dataclass(eq=False) class K3DecodeMoe: - # >>> route B: the engines' workspaces; no wide build + # The engines' workspaces; no wide build. """What every MoE layer's decode path shares on one device: the TP group's ``K3MoeHeadWorkspace`` (the front's all-gather), the engines' workspaces (``k3_moe_m1``, ``k3_moe_m2`` and the ``k3_moe`` build for up to 8 tokens; the layers run one at a time on one stream), the TP group's ``MnnvlWorkspace``, and the TP group's ``K3LatentExchange`` for a pushing step's latent all-reduce (None: every step keeps the routed experts' all-reduce). ``wide`` is None: ``k3_moe``'s wide build does not fit 896 local experts. Built by `create`.""" - # <<< route B - head: K3MoeHeadWorkspace small: K3MoeState wide: K3MoeWideState mnnvl: MnnvlWorkspace exchange: Optional[K3LatentExchange] = None - # >>> route B: the one- and two-token engines' workspaces (``create`` builds both) + # The one- and two-token engines' workspaces (``create`` builds both). m1: Optional[K3MoeM1State] = None m2: Optional[K3MoeM2State] = None - # <<< route B - # >>> route B: every expert on each rank: the engines with their push builds for the exchange and the k3_moe build + # Every expert on each rank: the engines with their push builds for the exchange and the k3_moe build. # for up to 8 tokens; no wide build @classmethod def create( @@ -170,8 +163,6 @@ def create( K3MoeM2State.create(device, i_logical, i_tp, num_local, push=builds), ) - # <<< route B - def _experts(moe: nn.Module) -> tuple: """The routed experts' TRTLLM-Gen W4A8_MXFP4_MXFP8 buffers ``k3_moe`` reads in place.""" @@ -193,10 +184,9 @@ def layout_gaps(moe: nn.Module, tp_size: int, max_snapshots: int) -> list: gate_up, down = shared.gate_up_proj.weight, shared.down_proj.weight if not moe._reduce_routed_output: return ["a routed output the model does not reduce"] - # >>> route B: k3_moe_m1, k3_moe_m2 and k3_moe compile the SiTU caps in + # K3_moe_m1, k3_moe_m2 and k3_moe compile the SiTU caps in. if tuple(moe._situ_betas) != ENGINE_SITU_CAPS: return [f"SiTU caps {tuple(moe._situ_betas)}; the MoE engines compile {ENGINE_SITU_CAPS}"] - # <<< route B if (moe.num_experts, moe.top_k, moe.moe_hidden_size, moe.hidden_size) != (896, 16, 3584, 7168): return [ f"experts / top-k / latent / hidden {(moe.num_experts, moe.top_k, moe.moe_hidden_size)}" @@ -274,10 +264,9 @@ class K3DecodeMoeLayer: shared_cols: int small: K3MoeLayer wide: K3MoeLayer - # >>> route B: the one- and two-token engines' handles + # The one- and two-token engines' handles. m1: K3MoeM1Layer m2: K3MoeM2Layer - # <<< route B @classmethod def create( @@ -315,26 +304,24 @@ def create( lo=lo, width=width, shared_cols=gate_up.shape[0] // 2, - # >>> route B: the engines' handles; no wide build + # The engines' handles; no wide build. small=state.small.layer(*weights), wide=None, m1=state.m1.layer(*weights), m2=state.m2.layer(*weights), - # <<< route B ) def warm_up(self, moe: nn.Module) -> None: - # >>> route B: the engines; no wide build + # The engines; no wide build. """One call of every kernel of the decode path on zero inputs, so none compiles under a capture: the front (collective: every rank makes the same call), each engine (``k3_moe_m1``, ``k3_moe_m2``, ``k3_moe``) and, with the latent exchange, its push build, each push followed by the reduce (collective). Once per model: every layer's calls compile the same builds.""" - # <<< route B device = self.front_weight.device x = torch.zeros(1, moe.hidden_size, dtype=torch.bfloat16, device=device) ids, weights, x_fp8, x_sf, _ = self._front(moe, x) offset = moe.routed_experts.backend.slot_start - # >>> route B: every engine returning and, with the latent exchange, pushing (k3_moe's builds compile here), + # Every engine returning and, with the latent exchange, pushing (k3_moe's builds compile here),. # each push followed by the reduce; no wide build exchange = self.state.exchange for rows in (1, 2, 3): @@ -346,7 +333,6 @@ def warm_up(self, moe: nn.Module) -> None: torch.cuda.synchronize(device) self._routed(*args, offset, push=True) k3_latent_reduce(rows, exchange) - # <<< route B torch.cuda.synchronize(device) def _front(self, moe: nn.Module, x: torch.Tensor): @@ -361,16 +347,14 @@ def _front(self, moe: nn.Module, x: torch.Tensor): ) def takes(self, hidden_states: torch.Tensor, step, partial_tail: bool) -> bool: - # >>> route B: no wide step + # No wide step. """Whether this decode path runs the layer on ``step`` (bf16 rows ``hidden_states``): any step of at most `MAX_TOKENS` tokens, with or without ``partial_tail``.""" - # <<< route B rows = hidden_states.shape[0] if hidden_states.dtype != torch.bfloat16: return False - # >>> route B: no wide step (k3_moe's wide build does not fit 896 local experts) + # No wide step (k3_moe's wide build does not fit 896 local experts). return 0 < rows <= MAX_TOKENS - # <<< route B def forward( self, @@ -390,13 +374,12 @@ def forward( ids, weights, x_fp8, x_sf, shared_act = self._front(moe, hidden_states) offset = moe.routed_experts.backend.slot_start exchange = self.state.exchange - # >>> route B: the engine of the token count + # The engine of the token count. if push and exchange is not None: self._routed(x_fp8, x_sf, ids, weights, offset, push=True) latent = k3_latent_reduce(hidden_states.shape[0], exchange) else: latent = moe.routed_experts.all_reduce(self._routed(x_fp8, x_sf, ids, weights, offset)) - # <<< route B if partial_tail: return PendingTail( latent.contiguous(), @@ -415,7 +398,7 @@ def forward( ) return (up * scale + shared_out.float()).bfloat16() - # >>> route B: the engine of a token count + # The engine of a token count. def _routed(self, x_fp8, x_sf, ids, weights, offset: int, push: bool = False): """This rank's routed partial of the front's outputs on ``k3_moe_m1`` at one token, ``k3_moe_m2`` at two, else ``k3_moe`` (up to `MAX_TOKENS`); with ``push``, stored into every rank's latent exchange instead (None).""" @@ -435,7 +418,6 @@ def _routed(self, x_fp8, x_sf, ids, weights, offset: int, push: bool = False): return k3_moe_m2(x_fp8, x_sf, ids, weights, offset, self.m2) return k3_moe(x_fp8, x_sf, ids, weights, offset, self.small) - # <<< route B def _wide(self, moe: nn.Module, hidden_states: torch.Tensor, gemvs) -> torch.Tensor: """A wide decode step's MoE: this rank's unreduced share of the output, ``[M, hidden]`` bf16.""" x = hidden_states.contiguous() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py index 2b24a35badfd..e48707ebcc19 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 16 x moe_ep 1.""" +"""Kimi K3 MXFP4 / sm_100 / tp16_moetp16ep1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index b53cdb0e8482..575c6a578a97 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -1,3532 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# >>> route B: this target's layout, and decoding without speculation -"""ModelingV2 target: Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 16 x moe_ep 1. +"""Kimi K3 MXFP4, SM100, attention TP16 and MoE TP16 x EP1.""" -Kimi K3's language model: 93 layers, hidden 7168. Every fourth layer (3, 7, ..., 91) is MLA attention with 96 -query heads; the others are Kimi Delta Attention (KDA), a gated linear-attention recurrence with per-request state. -Layer 0's MLP is dense. Every other layer's MLP is a latent MoE: 896 routed experts, top-16, expert width 3584. -Attention residuals mix each sublayer's input from a bank of earlier outputs. The routed experts are MXFP4 (the -checkpoint's compressed-tensors default, read as W4A8_MXFP4_MXFP8); everything else is bf16, and so is the KV pool. - -`tp16_moetp16ep1` is `tensor_parallel_size: 16` with `moe_tensor_parallel_size: 16` and -`moe_expert_parallel_size: 1`, both set explicitly, no attention data parallelism, all 16 GPUs in one NVLink domain. -Attention is head-split (6 MLA query heads per rank), and every rank holds all 896 experts at a sixteenth of their -intermediate width (192 of 3072 values, zero-padded to 256 when they load). With the expert split left unset the -engine resolves the same sizes, but Kimi K3 then runs the experts expert-parallel over the 16 ranks, a layout no -target serves. This is Kimi K3's layout for decoding without speculation; `tp16_moetp4ep4` serves speculative -decoding. - -**This module is `tp16_moetp4ep4`'s, copied** (targets share no files) and changed only inside its marked route B -blocks; test_modeling_v2_kimi_k3_drift.py checks that the two modules match everywhere else. - -**Each step is classified once, on the host, from its shape** (`decode_step`, a `DecodeStep` or None), and the -classification decides which kernels each module runs: - -* **small**: at most 8 tokens (`DECODE_MAX_TOKENS`, one token tile), context requests included. The token-count - kernels take it: decode GEMVs, the sandwiches, the embedding and residual epilogues. -* **decode**: a pure decode step of R <= 8 generation requests and no context request; without speculation each - request has one token. The request-aware kernels take it: MLA attention and its KV store. - -Every other step (prefill, mixed steps, decode steps above those bounds) runs the **generic path**: this target's -text model (`KimiLinearModel` below: decoder layers, attention residuals, the MLA / KDA / MoE runtimes), computed as -the built-in Kimi K3 text model computes it, on stock modules and ops that have no catalog entries yet -(`UNCERTIFIED_GENERIC_CALLS` names them), except that the routed experts' top-16 is computed outside the TRTLLM-Gen -kernel (`KimiK3MoeRoutingMethod`), with the kernel's arithmetic. - -The text model hands each step's classification to its attention modules, which run a **decode step** on the K3 -decode kernels' catalog entries: - -* KDA (`K3DecodeKDA`): one token per request, the fused input projection and the plain decode in one - `ssm/k3_kda_decode_attn` launch. -* MLA (`K3DecodeMLA`): `attention/k3_mla_qkv` (the query path and the step's latent KV rows into the paged cache), - then `attention/k3_mla_attn_vb_out` (the attention, v_b and the output gate in one launch). - -The projections around them run on the decode GEMV sites of `decode_gemv.py` (the [W_a; W_g] projection, `o_proj` -on every classified step), as do the LM head, the embedding and layer 0's dense MLP. A classified step's -attention-residual epilogues (the selection and the RMSNorm after it) take the fused kernels up to one token tile. On -any step of at most 16 tokens, the post-attention all-reduce carries the residual update (`decode_comm.py`): at most 8 -tokens on an attention decode branch, o_proj, the all-reduce and the update are one `comm/k3_sandwich_oproj` kernel; -otherwise the attention's unreduced o_proj output goes through `comm/mnnvl_allreduce_attn_res`. The stock MNNVL -all-reduces send one-shot up to 4 MiB. The state those kernels share (the KDA projection's Lamport buffers, the MLA -attention workspace, the decode GEMVs' state, the TP group's MNNVL and sandwich workspaces) lives in typed objects this -target creates in `post_load_weights`, before any graph capture. - -A MoE layer on a step of at most 8 tokens runs `decode_moe.py` with every expert on each rank: `moe/k3_moe_front` -(this rank's head slice, its all-gather, the top-16 routing, the MXFP8 latent and the shared experts' gate_up + SiTU), -then the routed experts as `moe/k3_moe_m1` at one token, `moe/k3_moe_m2` at two and `moe/k3_moe` over all 896 experts -up to 8, the latent all-reduce and the row-parallel tail. On a pushing step (`DecodeStep.latent_push`) the engines -push their partials into the TP group's `K3LatentExchange` and `comm/k3_latent_reduce` sums them; other steps use the -routed experts' all-reduce, which sums in the same order. Its routing is the front's: the noaux_tc arithmetic, as on -the generic path, on the router logits of the front's head GEMV. Wider steps, and every step of a checkpoint whose -SiTU caps differ from the ones the engines compile in, run the generic path: `k3_moe`'s wide build does not fit 896 -local experts. - -**What this target asserts rather than adapts**: SM 10.0; the topology above, with the expert split set explicitly; -no speculative decoding; the MXFP4 checkpoint's quantization (W4A16_MXFP4 with no per-layer declarations, so the -routed experts run the W4A8_MXFP4_MXFP8 default and the excluded modules stay bf16); bf16 weights and a bf16 KV pool; -tokens_per_block 64 (the MLA generation kernels K3's 96 heads reach exist only at 64); the V2 hybrid KV / state -manager, which holds the KDA states, with block reuse off and fp32 recurrent states; an all-reduce strategy of AUTO or -MNNVL. The construction-time ones fail in `__init__`, the -per-engine ones on the first forward, each naming the setting. A layer the decode kernels do not take fails the -weight load. - -**Text only.** The checkpoint is the vision-language wrapper. This target builds and loads no vision tower (its -weights are a predicted non-load, `weights.py`), and a step carrying multimodal input raises. - -**No speculative decoding.** A speculative decoding config fails construction, naming `tp16_moetp4ep4`'s expert -split. The causal LM is still the stock one-engine shell the built-in model builds on, which without that config -builds no drafter and no worker. This target drops `tp16_moetp4ep4`'s KDA verify kernels and the per-token verify -states they keep; the text model's speculative hidden-state taps stay as `tp16_moetp4ep4` has them and never fire. -The DSpark drafter (`K3DSparkDrafter`) and the speculative worker (`spec_worker.py`), which `_own_spec_branch` puts in -place of the shell's stock ones, also stay as `tp16_moetp4ep4` has them and are never built here; the drafter's decode -GEMV sites are warmed with the others and never called. -""" -# <<< route B - -from __future__ import annotations - -import copy -import math -import os -from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( - k3_drafter_attn_qknorm, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( - k3_mla_attn_vb_out, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( - K3MlaAttnWorkspace, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import k3_mla_qkv - -# >>> route B: no KDA verify kernels, and no trtllm::kda_mtp_decode (the built-in KDA verify's) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import K3KdaBuffers -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( - k3_kda_decode_attn, -) -from tensorrt_llm._torch.attention.backends import AttentionMetadata -from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.decode_view import k3_mla_decode_view - -# <<< route B -from tensorrt_llm._torch.distributed import AllReduce from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM -from tensorrt_llm._torch.models.modeling_dspark import ( - GQADSparkForCausalLM, - draft_is_embedded_in_target, -) -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM -from tensorrt_llm._torch.models.modeling_utils import DecoderModel -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention -from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention -from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace -from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel -from tensorrt_llm._torch.modules.rms_norm import RMSNorm -from tensorrt_llm._torch.modules.situ import SituAndMul -from tensorrt_llm._torch.moe.fused_moe import ( - ConfigurableMoE, - SiTuActivation, - TRTLLMGenFusedMoE, - create_moe, -) -from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind -from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph -from tensorrt_llm._torch.pyexecutor.config_utils import is_mla -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 -from tensorrt_llm._torch.utils import AuxStreamType -from tensorrt_llm.functional import AllReduceStrategy -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - -from . import decode_comm as _decode_comm -from . import decode_gemv as _decode_gemv -from . import decode_moe as _decode_moe -from . import spec_worker as _spec_worker -from . import weights as _weights - -if TYPE_CHECKING: - from transformers import PretrainedConfig - - -# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, -# and the certification is per arch. -_SM = (10, 0) - -#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in -#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router -#: and fused-A GEMMs), then the decode kernels'. -REQUIRED_TRTLLM_OPS = ( - "attn_res_fwd", - "attn_res_rmsnorm_fwd", - "attn_res_add_rmsnorm_fwd", - "attn_res_add_rmsnorm_persistent_fwd", - "kda_prefill", - "kda_decode", - # >>> route B: no KDA verify (kda_mtp_decode, k3_kda_attn, k3_kda_verify) - "dsv3_router_gemm_op", - "dsv3_fused_a_gemm_op", - "k3_kda_decode_attn", - # <<< route B - "k3_mla_qkv", - "k3_mla_attn_vb_out", - # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). - "k3_decode_gemv", - "k3_ctm_gemv_wide", - "k3_ctm_gemv_long", - "k3_situ_mul", - "k3_head_gemv", - "k3_embed_norm", - "allgather", - # The decode path's collectives (decode_comm.py) and MoE (decode_moe.py). - "mnnvl_allreduce_attn_res", - "k3_sandwich_oproj", - "k3_sandwich_tail", - "k3_moe_front", - "k3_moe", - "k3_latent_reduce", - "k3_route_quant", - "mnnvl_allgather_split", - # >>> route B: the MoE engines of one and two tokens - "k3_moe_m1", - "k3_moe_m2", - # <<< route B - # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces - # with the residual add and RMSNorm (the split context projection's with hidden_norm). - "k3_ctm_gemv", - "k3_ctm_gemv_swiglu", - "k3_drafter_attn_qknorm", - "k3_sandwich_plain", - "mnnvl_fusion_allreduce", -) - -#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. -#: A renamed field upstream fails here, loudly, instead of reading as a default. -REQUIRED_ENGINE_FIELDS = { - "attn_metadata": ( - "num_contexts", - "num_generations", - "num_seqs", - "num_tokens", - "seq_lens", - "tokens_per_block", - "kv_cache_manager", - "mamba_metadata", - ), - "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), -} - -#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every -#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry -#: leaves the list when a catalog entry replaces it. -UNCERTIFIED_GENERIC_CALLS = ( - # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. - "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", - "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", - "tensorrt_llm._torch.models.modeling_utils.DecoderModel", - # The text model's stock modules. - "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", - "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", - # >>> route B: no trtllm::kda_mtp_decode registration (the built-in KDA verify's) - # <<< route B - "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", - "tensorrt_llm._torch.moe.fused_moe.create_moe", - "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", - "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", - "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", - "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", - "tensorrt_llm._torch.modules.situ.SituAndMul", - "tensorrt_llm._torch.modules.rms_norm.RMSNorm", - "tensorrt_llm._torch.distributed.AllReduce", - "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", - # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries - # do not take a block, its context projection without the TP group's collective state, its context k / v, its - # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. - "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", - "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", - "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", - "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", -) - -# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware -# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). -DECODE_MAX_TOKENS = 8 -MAX_REQUESTS = 8 -MAX_TOKENS_PER_REQUEST = 8 - -# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own -# get_model_defaults sets it for the same reason). -_TOKENS_PER_BLOCK = 64 - -_LANG_PREFIX = "language_model." - - -# ---------------------------------------------------------------------------------------------------------------------- -# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and -# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, -# `K3DecodeMLA` below). -# ---------------------------------------------------------------------------------------------------------------------- - -# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections -# instead of the min-latency fused GEMM op (read once at import). -_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( - os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" -) - - -# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table -# holds cos=1/sin=0), but the chunked-context path indexes the table by -# absolute position, so it must cover max_position_embeddings (~512MB per -# backend for the 1M-position checkpoint); a smaller table is read out of -# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context -# deployments. -_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" - - -class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): - """DeepSeek-V3 routing, computed outside the TRTLLM-Gen MoE kernel. - - The kernel's own top-16 of 896 experts is the slow part of a generic MoE call of a few tokens. Outside it, the - scheduler routes with the backend's fused route + MXFP8 quantize (``trtllm::k3_route_quant``) up to 64 tokens and - with ``noaux_tc_op`` above, and the kernel takes the expert ids and weights. Both compute the kernel's arithmetic - and break ties the same way; a bf16 weight can differ by one ulp where it lies within fp32 rounding of a bf16 - rounding boundary. From about a thousand tokens on, ``noaux_tc_op`` takes longer than the kernel's routing. - """ - - @property - def requires_separated_routing(self) -> bool: - return True - - -class KimiK3MoEGate(nn.Module): - """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" - - def __init__( - self, - config: Any, - *, - logits_gemm_dtype: torch.dtype | None = None, - device: torch.device | None = None, - ) -> None: - super().__init__() - self.config = config - self.top_k = config.num_experts_per_token - self.num_experts = config.num_experts - self.routed_scaling_factor = config.routed_scaling_factor - self.moe_router_activation_func = config.moe_router_activation_func - self.num_expert_group = getattr(config, "num_expert_group", 1) - self.topk_group = getattr(config, "topk_group", 1) - self.moe_renormalize = config.moe_renormalize - self.gating_dim = config.hidden_size - - assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( - "K3 MoE gate supports sigmoid or softmax scoring only" - ) - - # The checkpoint stores the gate weight in bf16. Storing it in bf16 - # permits the single bf16xbf16 router GEMM while retaining fp32 output. - weight_dtype = logits_gemm_dtype or torch.float32 - self.weight = nn.Parameter( - torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) - ) - self.e_score_correction_bias = nn.Parameter( - torch.empty(self.num_experts, dtype=torch.float32, device=device) - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" - hidden_2d = hidden_states.reshape(-1, self.gating_dim) - if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: - return torch.ops.trtllm.dsv3_router_gemm_op( - hidden_2d.contiguous(), - self.weight.t(), - bias=None, - out_dtype=torch.float32, - ) - return torch.nn.functional.linear( - hidden_2d.type(torch.float32), - self.weight.type(torch.float32), - None, - ) - - @property - def routing_method(self) -> KimiK3MoeRoutingMethod: - """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" - if self.moe_router_activation_func != "sigmoid": - raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") - if not self.moe_renormalize: - raise ValueError( - "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." - ) - return KimiK3MoeRoutingMethod( - top_k=self.top_k, - n_group=self.num_expert_group, - topk_group=self.topk_group, - routed_scaling_factor=self.routed_scaling_factor, - callable_e_score_correction_bias=lambda: self.e_score_correction_bias, - is_fused=True, - ) - - -class KimiK3RMSNorm(nn.Module): - """RMSNorm matching the Kimi checkpoint implementation's rounding.""" - - def __init__( - self, - hidden_size: int, - eps: float = 1e-6, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - ) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) - self.eps = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states_float = hidden_states.to(torch.float32) - variance = hidden_states_float.pow(2).mean(-1, keepdim=True) - hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) - return self.weight * hidden_states_float.to(input_dtype) - - -def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: - """Return the finite SiTu betas required by the routed-expert kernels.""" - config_situ_beta = getattr(cfg, "activation_situ_beta", None) - situ_beta = 1.0 if config_situ_beta is None else config_situ_beta - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - if situ_linear_beta is None: - raise ValueError( - "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " - "None means an identity linear branch that the fused kernels cannot represent." - ) - if situ_beta <= 0 or situ_linear_beta <= 0: - raise ValueError( - f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." - ) - return float(situ_beta), float(situ_linear_beta) - - -def _get_text_config(pretrained_config: "PretrainedConfig"): - """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" - if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( - not hasattr(pretrained_config, "linear_attn_config") - and hasattr(pretrained_config, "text_config") - ): - return pretrained_config.text_config - return pretrained_config - - -def _is_kda_layer(cfg, layer_idx: int) -> bool: - return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] - - -def _is_mla_layer(cfg, layer_idx: int) -> bool: - return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] - - -KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" - - -_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" - - -KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" - - -_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" - - -KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" - - -KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" - - -_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") - - -# Turning the port on should not require knowing which of the three topologies -# is the right one: "1" selects the measured policy (``split``), so enabling the -# feature and choosing the policy are one action through one variable. -_ATTN_RES_TOPOLOGY_ON = "1" - - -def _read_attn_res_topology() -> str: - """Default ``per_token``: the persistent kernel is opt-in. - - ``1`` is the only accepted on-value and resolves to ``split``. The named - topologies stay available for measurement: ``persistent`` uses the - persistent kernel at every shape it implements, ``per_token`` at none. - """ - raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") - if raw == _ATTN_RES_TOPOLOGY_ON: - return "split" - if raw not in _ATTN_RES_TOPOLOGIES: - # Loudly, for the same reason as the token ceiling below: a mistyped A/B - # arm that silently fell back to the default would measure one side - # twice and report no difference. - raise ValueError( - f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " - f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " - f"(which means 'split'), got {raw!r}" - ) - return raw - - -def _read_fused_attn_res_max_tokens() -> int: - """Resolved after the topology, because its default follows it. - - With the persistent port off -- the default -- the ceiling is 1, the - pre-existing gate: the fused epilogue is taken at the single-token decode - shape and nowhere else. Enabling the port raises it to 32, the top of the - measured range, so that "off" keeps meaning "unchanged". - """ - default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" - raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) - try: - value = int(raw) - except ValueError: - # Failing loudly matters more than usual here: a mistyped A/B arm that - # silently fell back to the default would measure the candidate twice - # and report no difference. - raise ValueError( - f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" - ) from None - if value < 1: - raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") - return value - - -_ATTN_RES_TOPOLOGY = _read_attn_res_topology() - - -_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() - - -def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: - """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. - - No token ceiling: the persistent grid is sized by the SM count, not by the - token count, so prefill is the case it exists for. - """ - del M # deliberately unused; see above - return H == 7168 and 2 <= N <= 9 - - -def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: - """Pick between the two fused kernels for this call site. - - ``persistent`` takes the persistent kernel at every shape it implements; - ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in - for the prefill/decode boundary. Shapes the persistent kernel does not - implement fall through to the caller's existing gate and land on the - unfused path. - """ - if not _persistent_attn_res_applicable(M, H, N): - return False - if _ATTN_RES_TOPOLOGY == "persistent": - return True - return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens - - -def _apply_attn_res_fused( - prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm -) -> Optional[torch.Tensor]: - """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. - - Returns ``None`` when the call falls outside the fused kernel's - contract (dtype/shape/arch) so the caller can use the exact fp32 reference - instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` - layout. Candidate order matches the reference: snapshots first, the - running prefix sum last. - """ - if ( - prefix_sum.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not block_residual.is_cuda - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): - return None - try: - attn_res_op = torch.ops.trtllm.attn_res_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output, _rsigma, _probs, _logits = attn_res_op( - layer_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - ) - return output.reshape(M, H) - - -def _rms_norm_eps(norm: nn.Module) -> float: - if hasattr(norm, "eps"): - return float(norm.eps) - return float(norm.variance_epsilon) - - -def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: - """Report whether the fused path was actually reached, once per shape. - - ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that - the shape gate let the call through, and a rejected call looks exactly like - a disabled one in the logs. Emitted at debug level: it is once per distinct - shape, not once per process, so it is a diagnostic rather than a summary. - """ - logger.debug_once( - f"Kimi K3 attn-res fusion [{site}]: " - f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", - key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", - ) - - -def _apply_attn_res_rmsnorm_fused( - prefix_sum: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Optional[torch.Tensor]: - """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, - default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" - if max_fused_tokens is None: - max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS - if ( - prefix_sum.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not block_residual.is_cuda - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - N = K + 1 - # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, - # which is the measured window; larger token counts have not been measured - # and fall back to the unfused add + attn_res_fwd + RMSNorm path. - if _use_persistent_attn_res(M, H, N, max_fused_tokens): - try: - persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd - except (AttributeError, RuntimeError): - return None - _, output = persistent_op( - prefix_sum.reshape(M, 1, H).contiguous(), - None, - block_residual.reshape(K, M, 1, H).contiguous(), - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - ) - _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) - return output.reshape(M, H) - - if M > max_fused_tokens or H != 7168 or N > 12: - _note_attn_res_fusion("attn_res+norm", False, M, H, N) - return None - try: - attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output = attn_res_rmsnorm_op( - layer_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - early_trigger=True, - ) - _note_attn_res_fusion("attn_res+norm", True, M, H, N) - return output.reshape(M, H) - - -def _apply_attn_res_add_rmsnorm_fused( - prefix_sum: torch.Tensor, - addend: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in - ``_apply_attn_res_rmsnorm_fused``). - - The production residual add produces a BF16 tensor that remains live across - the following MLP. The kernel therefore returns that materialized, - BF16-rounded prefix sum alongside the normalized attention-residual output, - while avoiding a separate add launch and a re-read of the intermediate by - attention-residual selection. - """ - if max_fused_tokens is None: - max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS - if ( - prefix_sum.dtype is not torch.bfloat16 - or addend.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not addend.is_cuda - or not block_residual.is_cuda - or prefix_sum.shape != addend.shape - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - N = K + 1 - # Same measured window as _apply_attn_res_rmsnorm_fused above. - if _use_persistent_attn_res(M, H, N, max_fused_tokens): - try: - persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd - except (AttributeError, RuntimeError): - return None - updated_prefix_sum, output = persistent_op( - prefix_sum.reshape(M, 1, H).contiguous(), - addend.reshape(M, 1, H).contiguous(), - block_residual.reshape(K, M, 1, H).contiguous(), - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - ) - _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) - return updated_prefix_sum.reshape(M, H), output.reshape(M, H) - - if M > max_fused_tokens or H != 7168 or N > 12: - _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) - return None - try: - attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - addend_kernel = addend.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - updated_prefix_sum, output = attn_res_add_rmsnorm_op( - layer_kernel, - addend_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - early_trigger=True, - ) - _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) - return updated_prefix_sum.reshape(M, H), output.reshape(M, H) - - -def _apply_attn_res( - prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm -) -> torch.Tensor: - """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). - - prefix_sum: ``[num_tokens, hidden_size]`` - block_residual: ``[num_snapshots, num_tokens, hidden_size]`` - - Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's - contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. - Only the fallback boundary restores the HF ``[M, K, H]`` layout. - """ - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) - if fused is not None: - return fused - block_residual_hf = block_residual.transpose(0, 1) - v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) - v_float = v.float() - variance = v_float.pow(2).mean(-1, keepdim=True) - k = v_float * torch.rsqrt(variance + norm.eps) - score_weight = norm.weight.float() * proj.weight.squeeze(0).float() - scores = (k * score_weight).sum(-1) - probs = scores.softmax(-1).unsqueeze(1) - hidden_states = torch.matmul(probs, v_float).squeeze(1) - return hidden_states.to(v.dtype) - - -def _apply_attn_res_and_rmsnorm( - prefix_sum: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> torch.Tensor: - """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the - fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_rmsnorm_fused( - prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens - ) - if fused is not None: - return fused - return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) - - -def _apply_attn_res_add_and_rmsnorm( - prefix_sum: torch.Tensor, - addend: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in - ``_apply_attn_res_and_rmsnorm``).""" - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_add_rmsnorm_fused( - prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens - ) - if fused is not None: - return fused - updated_prefix_sum = prefix_sum + addend - return updated_prefix_sum, _apply_attn_res_and_rmsnorm( - updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens - ) - - -# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where -# they are faster than the add -> attn_res -> RMSNorm chain. -_WIDE_ATTN_RES_MAX_TOKENS = 32 - - -def _attn_res_max_tokens(step: DecodeStep) -> int: - """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step - ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" - return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS - - -_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") - - -# The routed experts' module path prefixes. ``exclude_modules`` matches with -# wildcards and walks ancestor prefixes, so an empty prefix would widen what -# matches instead of just missing. -_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") - - -# Routed-expert quantization used when the checkpoint declares nothing per -# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors -# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that -# checkpoint is what this default has always served. -_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 - - -def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: - backend.quant_method.load_packed_mxfp4_expert( - backend, - global_expert_id=expert_idx, - local_slot_id=local_slot_id, - w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), - w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), - w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), - w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), - w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), - w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), - ) - - -class _K3ExpertCkptSpec(NamedTuple): - """How one routed-expert quantization is spelled and loaded.""" - - # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. - kinds: Tuple[str, ...] - loader: Callable[..., None] - # Set of filled slots the loader maintains, checked after the load. - loaded_slots_attr: str - # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). - needs_layer_finalize: bool - - -_K3_EXPERT_CKPT_SPECS = { - QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( - kinds=("weight_packed", "weight_scale"), - loader=_load_packed_mxfp4_expert, - loaded_slots_attr="_packed_mxfp4_loaded_slots", - needs_layer_finalize=False, - ), -} - - -def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: - spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) - if spec is None: - raise NotImplementedError( - f"Kimi K3 routed experts are quantized as {quant_algo}, for which " - "no per-expert checkpoint layout is known. Supported: " - f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." - ) - return spec - - -class KimiK3MoERuntime(nn.Module): - """Kimi K3 latent MoE block backed by ConfigurableMoE.""" - - def __init__( - self, - model_config: ModelConfig, - cfg, - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ): - """Build the routed experts and the shared expert for one MoE layer. - - ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: - the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields - that ``ModelConfig`` does not carry. - - ``aux_stream_dict`` is shared across every layer of the model, so the - streams reached through it are borrowed and must not be synchronized - or reassigned here. - """ - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = cfg.hidden_size - self.num_experts = cfg.num_experts - self.top_k = cfg.num_experts_per_token - self.moe_hidden_size = cfg.routed_expert_hidden_size - # ValueError (not assert): these guard unsupported checkpoint - # configurations and must stay active under ``python -O``. - if self.moe_hidden_size is None: - raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") - if not getattr(cfg, "latent_moe_use_norm", False): - raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") - - situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) - self._situ_betas = (situ_beta, situ_linear_beta) - dtype = torch.bfloat16 - - # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 - # accumulate/output (checkpoint stores the gate weight in bf16; saves a - # per-layer input cast + fp32 splitK pair on the bs1 decode path). - # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. - _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") - _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True - self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) - - routed_moe_model_config = self._routed_moe_model_config(model_config) - routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) - # Resolved here so ``load_weights`` reads the checkpoint layout off the - # module instead of re-deriving it at each of its three call sites. - self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) - routed_moe_kwargs = dict( - routing_method=self.gate.routing_method, - num_experts=self.num_experts, - hidden_size=self.moe_hidden_size, - intermediate_size=cfg.moe_intermediate_size, - dtype=dtype, - # Kimi owns the latent reduction so it can order that collective - # after the shared expert's auxiliary-stream reduction. - reduce_results=False, - model_config=routed_moe_model_config, - override_quant_config=routed_quant_config, - layer_idx=layer_idx, - aux_stream_dict=aux_stream_dict, - # Let CommunicationFactory select the best available strategy. - communication_method=None, - activation=SiTuActivation( - gate_softcap=situ_beta, - linear_softcap=situ_linear_beta, - ), - # A request that silently degraded to CUTLASS would be benchmarked - # as if it were the backend that was asked for, and the decline is - # easy to trigger: MegaMoE has its own token / top-k limits and is - # EP-only, and CuteDSL declines on activation shape, SM version and - # the CuTe DSL dependency. As measured once: a CUTEDSL request - # was turned down on every one of the 92 MoE layers, on all 16 - # ranks, and still produced correct text and a zero exit -- the - # only trace was a warning line per layer. Fail in the resolver - # instead, which reports the rejection trail. - # - # CUTLASS is absent on purpose: it is the fallback target, so - # "degraded to CUTLASS" is not a thing that can happen to it. - allow_backend_degradation=routed_moe_model_config.moe_backend - not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), - ) - self._check_trtllm_situ_quant( - routed_moe_model_config.moe_backend, routed_quant_config.quant_algo - ) - - self.routed_experts = create_moe(**routed_moe_kwargs) - if not isinstance(self.routed_experts, ConfigurableMoE): - raise RuntimeError( - "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." - ) - if self.routed_experts.layer_load_balancer is not None: - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming does not yet support " - "dynamic EPLB or replicated expert slots." - ) - local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) - if local_expert_ids != list( - range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) - ): - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming currently requires a " - "contiguous static expert partition." - ) - self.local_expert_ids = tuple(local_expert_ids) - self.experts_per_rank = len(local_expert_ids) - self.expert_lo = local_expert_ids[0] - self.expert_hi = self.expert_lo + self.experts_per_rank - - shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts - shared_model_config = copy.copy(model_config) - shared_model_config.quant_config = QuantConfig() - # Direct MoE-TP leaves both branches as partials for one concatenated - # all-reduce. - use_shared_tp = model_config.mapping.tp_size > 1 - self._reduce_routed_output = ( - use_shared_tp - and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM - ) - if self._reduce_routed_output and self.routed_experts.all_reduce is None: - raise RuntimeError( - "Kimi K3 direct MoE tensor parallelism requires the " - "ConfigurableMoE all-reduce even when reduce_results=False." - ) - self.shared_experts = GatedMLP( - hidden_size=cfg.hidden_size, - intermediate_size=shared_intermediate, - bias=False, - activation=SituAndMul( - beta=situ_beta, - linear_beta=situ_linear_beta, - use_fused_activation=True, - ), - dtype=dtype, - config=shared_model_config, - reduce_output=use_shared_tp, - layer_idx=layer_idx, - is_shared_expert=True, - ) - # Side stream (+ fork/join events) for overlapping shared-expert - # compute with the routed chain. Only engaged when multi-stream is - # active (CUDA graphs on); otherwise both run in order on the default - # stream. - self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] - self.moe_main_event = torch.cuda.Event() - self.moe_shared_event = torch.cuda.Event() - self.routed_expert_down_proj = nn.Linear( - cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype - ) - self.routed_expert_up_proj = nn.Linear( - self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype - ) - # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer - # fallback is the same fp32-variance eager math as KimiK3RMSNorm). - self.routed_expert_norm = RMSNorm( - hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. - self.decode_moe: Optional[_decode_moe.K3DecodeMoeLayer] = None - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - @staticmethod - def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: - if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): - return projection(hidden_states) - return torch.ops.trtllm.dsv3_fused_a_gemm_op( - hidden_states, projection.weight.t(), None, None - ) - - @staticmethod - def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: - # >>> route B: this target's split - """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit - ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (16 x 1, asserted at - construction).""" - # <<< route B - return mapping.moe_tp_size, mapping.moe_ep_size - - @staticmethod - def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: - """Routed-expert quantization for ``layer_idx``: the MXFP4 checkpoint declares nothing per layer (asserted at - construction), so the experts keep the ``W4A8_MXFP4_MXFP8`` default. - - An exclusion outranks the default: ``create_weights`` treats an override - as authoritative over anything ``__post_init__`` wrote, so this return - value stands in for both quantization passes and exclusion is the one - that runs second. It is matched as a pattern, so it is asked only about - real module names. - """ - quant_config = model_config.quant_config - if quant_config is not None and any( - quant_config.is_module_excluded_from_quantization( - f"{prefix}layers.{layer_idx}.{suffix}" - ) - for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES - for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES - ): - logger.debug( - "Kimi K3 layer %d routed experts: excluded from quantization, " - "keeping them unquantized", - layer_idx, - ) - return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) - - logger.debug( - "Kimi K3 layer %d routed experts: no per-layer quant config in the " - "checkpoint, defaulting to %s", - layer_idx, - _K3_DEFAULT_ROUTED_QUANT_ALGO, - ) - return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) - - @staticmethod - def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: - """Reject a routed-expert format trtllm-gen has no fused SiTu cubin for. - - trtllm-gen has fused SiTu FC1 cubins for two input formats and no - standalone SiTu activation kernel, so anything else has to die here - rather than in a cubin lookup deep inside the runner. Checked against - the resolved backend, not the K3 architecture branch, because the - generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also - land on TRTLLM. - - The admitted set is read off the backend rather than restated here, - because restating it is what broke. This guard was written in #17865 - when MXFP4 was the only fused SiTu drop; #17940 then added the NVFP4 - (group-16 ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``) cubins and updated - ``TRTLLMGenFusedMoE``'s set without touching this copy. For the week - in between, an NVFP4 K3 checkpoint could not start at all -- and not - only when TRTLLM was asked for by name, because - ``ModelConfig.resolve_moe_backend`` sends every K3 architecture to - TRTLLM, so the default AUTO configuration hit this raise too. The unit - tests did not catch it: they call ``create_moe`` directly and never - reach this guard, so the kernel path stayed green while the model path - was closed. - - A staticmethod, not an inline block, so that the invariant is - reachable from a test without constructing the whole runtime. - """ - situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() - if moe_backend != "TRTLLM" or quant_algo in situ_supported: - return - supported = ", ".join(sorted(algo.name for algo in situ_supported)) - raise ValueError( - f"Kimi K3 routed experts are quantized as {quant_algo}, which the " - "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " - f"exist only for {supported}. Set moe_config.backend to CUTLASS " - "or MEGAMOE_CUTEDSL." - ) - - @staticmethod - def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: - """Build a private routed-expert mapping without mutating the shared - config, with the split of ``_select_moe_tp_ep``.""" - # Every backend here declares ``ActivationType.SiTu`` in its - # ``activation_support``; the list is not a preference order. CUTEDSL - # joined once its act-fusion kernel grew the SiTU epilogue. - supported_backends = { - "CUTLASS", - "TRTLLM", - "CUTEDSL", - "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - } - if model_config.moe_backend not in supported_backends: - raise ValueError( - "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " - "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " - f"got {model_config.moe_backend!r}." - ) - if model_config.moe_load_balancer is not None: - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming does not yet support " - "EPLB or replicated expert slots." - ) - mapping = model_config.mapping - - moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) - logger.info_once( - f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " - f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", - key="kimi_k3_moe_tp_ep_split", - ) - - mapping_dict = mapping.to_dict() - mapping_dict["moe_cluster_size"] = 1 - mapping_dict["moe_tp_size"] = moe_tp - mapping_dict["moe_ep_size"] = moe_ep - routed_mapping = Mapping.from_dict(mapping_dict) - - routed_model_config = copy.copy(model_config) - routed_model_config._frozen = False - routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) - routed_model_config.mapping = routed_mapping - routed_model_config.moe_backend = model_config.moe_backend - # MegaMoE uses this value as global DP SymmBuffer capacity, then - # divides it by EP size for the per-rank allocation. Other backends - # keep the user-configured value as their MoE chunking bound. - # Preserve an explicitly larger capacity. - if routed_model_config.moe_backend in { - "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - }: - default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size - configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) - if configured_moe_max_num_tokens < default_moe_max_num_tokens: - logger.info_once( - "Kimi K3 MegaMoE raises moe_max_num_tokens from " - f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " - "because the global DP SymmBuffer requires capacity for " - "max_num_tokens * dp_size.", - key=( - "kimi_k3_megamoe_capacity_override_" - f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" - ), - ) - routed_model_config.moe_max_num_tokens = max( - configured_moe_max_num_tokens, - default_moe_max_num_tokens, - ) - routed_model_config._frozen = True - return routed_model_config - - def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: - """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel - share (``partial_tail``): the decode path takes the step with that tail.""" - return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) - - def forward( - self, - hidden_states: torch.Tensor, - all_rank_num_tokens=None, - partial_tail: bool = False, - *, - step: DecodeStep, - ) -> Union[torch.Tensor, _decode_comm.PendingTail]: - """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification - (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where - ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at - most 8 tokens, a tensor on a wide decode step.""" - decode = self.decode_moe - if decode is not None and decode.takes(hidden_states, step, partial_tail): - return decode.forward( - self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push - ) - if partial_tail: - raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") - return self.forward_generic(hidden_states, all_rank_num_tokens) - - def forward_generic( - self, hidden_states: torch.Tensor, all_rank_num_tokens=None - ) -> torch.Tensor: - identity = hidden_states - router_logits = self.gate.compute_logits(hidden_states) - moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None +from tensorrt_llm.quantization.mode import QuantAlgo - def _routed_output(): - # Latent down/up projections via the min-latency fused GEMM op: - # at <=16 tokens (decode graphs) it runs a single pipelined - # bf16 kernel per projection instead of cuBLAS's split-K GEMV + - # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); - # for larger token counts the op falls back to cuBLAS internally. - # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear - # (A/B escape hatch). When the FP8 weight-read conversion has - # replaced the projection module, call it directly: its weight is - # an e4m3 buffer the bf16 dsv3 op must not read, and its forward - # is already a single fused GEMM (fp8_swap_ab_gemm). - routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) - y = self.routed_experts( - routed_in, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - ) - if self._reduce_routed_output: - return y - # Communication-backed paths return a complete routed result. - y = self.routed_expert_norm(y) - return self._routed_projection(y, self.routed_expert_up_proj) +from ..modeling import K3TargetSpec, ModelingV2KimiK3 +from ..modeling import _check_construction as _check_target_construction - # Shared experts depend only on the block input, so overlap their GEMMs - # with the routed dispatch/expert/combine chain. Multi-stream engages - # only under CUDA graphs; otherwise both branches run in order on the - # default stream. The shared GatedMLP includes its output all-reduce on - # the auxiliary stream. The join below must precede the routed - # all-reduce: concurrent collectives on different streams can corrupt - # SYMM_MEM all-reduce state. - routed_out, shared_out = maybe_execute_in_parallel( - _routed_output, - lambda: self.shared_experts(identity), - self.moe_main_event, - self.moe_shared_event, - self.shared_expert_stream, - disable_on_compile=True, - ) - if self._reduce_routed_output: - routed_latent = moe_all_reduce(routed_out) - routed_latent = self.routed_expert_norm(routed_latent) - routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) - return routed_out + shared_out - - -def resolve_attention_quant_config( - config: ModelConfig | None, layer_idx: int, projection: str -) -> QuantConfig: - """Resolve a checkpoint projection, including mixed-precision exclusions.""" - if config is None: - return QuantConfig() - global_config = config.quant_config or QuantConfig() - names = [ - f"{prefix}layers.{layer_idx}.self_attn.{projection}" - for prefix in ("language_model.model.", "model.", "") - ] - if any(global_config.is_module_excluded_from_quantization(name) for name in names): - return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) - declarations = config.quant_config_dict or {} - matches = [declarations[name] for name in names if name in declarations] - if matches: - selected = matches[0] - if any(match.quant_algo != selected.quant_algo for match in matches[1:]): - raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") - elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: - selected = QuantConfig() - else: - selected = global_config - if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): - raise ValueError( - f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " - f"quantization {selected.quant_algo}" - ) - if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): - raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") - result = copy.copy(selected) - result.kv_cache_quant_algo = global_config.kv_cache_quant_algo - return result - - -class KimiMLARuntime(nn.Module): - """Wraps K3 MLA and applies its external TP output reduction.""" - - def __init__( - self, - cfg: "PretrainedConfig", - layer_idx: int, - model_config: ModelConfig, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ) -> None: - super().__init__() - - max_positions = int( - os.environ.get( - _KIMI_K3_MLA_MAX_POSITIONS_ENV, - cfg.max_position_embeddings, - ) - ) - self.layer_idx = layer_idx - # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the - # final output reduction in this wrapper so the output gate remains - # between attention and the row-parallel o_proj. - mapping = model_config.mapping - reduce_output = mapping.tp_size > 1 - self._o_allreduce = ( - AllReduce( - mapping=mapping, - strategy=model_config.allreduce_strategy, - dtype=torch.bfloat16, - ) - if reduce_output - else None - ) - attention_config = copy.copy(model_config) - attention_config._frozen = False - attention_config.quant_config_dict = { - name: resolve_attention_quant_config(model_config, layer_idx, name) - for name in ( - "q_a_proj", - "kv_a_proj_with_mqa", - "q_b_proj", - "kv_b_proj", - "g_proj", - "o_proj", - ) - } - attention_config.quant_config = QuantConfig( - kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo - if model_config.quant_config is not None - else None - ) - attention_config._frozen = model_config._frozen - self.mixer = K3DecodeMLA( - hidden_size=cfg.hidden_size, - num_heads=cfg.num_attention_heads, - q_lora_rank=cfg.q_lora_rank, - kv_lora_rank=cfg.kv_lora_rank, - qk_nope_head_dim=cfg.qk_nope_head_dim, - qk_rope_head_dim=cfg.qk_rope_head_dim, - v_head_dim=cfg.v_head_dim, - rms_norm_eps=cfg.rms_norm_eps, - dtype=torch.bfloat16, - layer_idx=layer_idx, - use_output_gate=cfg.mla_use_output_gate, - max_position_embeddings=max_positions, - model_config=attention_config, - aux_stream_dict=aux_stream_dict, - ) - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" - return self.mixer.will_run_decode_branch(attn_metadata, step) - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - reduce_output: bool = True, - project_output: bool = True, - ) -> torch.Tensor: - """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the - mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" - # MLA.forward takes position_ids first; K3 is NoPE, so pass None. - out = self.mixer( - None, hidden_states, attn_metadata, step=step, project_output=project_output - ) - return self._reduce(out) if project_output and reduce_output else out - - def forward_generic( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - reduce_output: bool = True, - ) -> torch.Tensor: - out = self.mixer.forward_generic(None, hidden_states, attn_metadata) - return self._reduce(out) if reduce_output else out - - def _reduce(self, out: torch.Tensor) -> torch.Tensor: - # Head-sharded TP: sum the row-sharded o_proj partials across - # the head-shard group. - return self._o_allreduce(out) if self._o_allreduce is not None else out - - -class KimiLinearDecoderLayer(nn.Module): - def __init__( - self, - model_config: ModelConfig, - cfg, - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ): - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = cfg.hidden_size - dtype = torch.bfloat16 - - self.is_kda = _is_kda_layer(cfg, layer_idx) - is_mla = _is_mla_layer(cfg, layer_idx) - if self.is_kda == is_mla: - raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") - - if self.is_kda: - projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") - attention_config = copy.copy(model_config) - attention_config._frozen = False - attention_config.quant_config_dict = { - name: resolve_attention_quant_config(model_config, layer_idx, name) - for name in projection_names - } - attention_config._frozen = model_config._frozen - self.linear_attn = K3DecodeKDA( - cfg, - layer_idx, - mapping=model_config.mapping, - allreduce_strategy=model_config.allreduce_strategy, - aux_stream=aux_stream_dict[AuxStreamType.Attention], - model_config=attention_config, - ) - else: - self.self_attn = KimiMLARuntime( - cfg, - layer_idx, - model_config=model_config, - aux_stream_dict=aux_stream_dict, - ) - - self.is_moe = ( - cfg.num_experts is not None - and layer_idx >= cfg.first_k_dense_replace - and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 - ) - if self.is_moe: - self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) - else: - situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) - # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's - # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take - # (decode_gemv.SITES); otherwise it stays within one node. - spans_nodes = self._mnnvl_allreduce() is not None - if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: - self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) - mlp_model_config = copy.copy(model_config) - mlp_model_config.quant_config = QuantConfig() - # K3's dense layer is BF16, so a unit block size gives the same - # subgroup selection as DeepSeek-V3. - self.mlp = GatedMLP( - hidden_size=cfg.hidden_size, - intermediate_size=cfg.intermediate_size, - bias=False, - activation=SituAndMul( - beta=situ_beta, - linear_beta=situ_linear_beta, - use_fused_activation=True, - ), - dtype=dtype, - config=mlp_model_config, - overridden_tp_size=self.mlp_tp_size, - reduce_output=self.mlp_tp_size > 1, - layer_idx=layer_idx, - ) - self._situ = (situ_beta, situ_linear_beta) - # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics - # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied - # after downcast, use_gemma=False). - self.input_layernorm = RMSNorm( - hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.post_attention_layernorm = RMSNorm( - hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - - # Attention residual scheme (always on for K3). The res norms stay - # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by - # _apply_attn_res and the fused attn_res op, never called as - # modules. - self.attn_res_block_size = cfg.attn_res_block_size - assert self.attn_res_block_size is not None, ( - "Kimi K3 runtime expects attn_res_block_size to be set" - ) - self.self_attention_res_norm = KimiK3RMSNorm( - cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) - self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set - # by the target's post_load_weights. - self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None - self.sandwich_oproj = False - - def forward( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - attn_metadata: AttentionMetadata, - step: DecodeStep, - capture: Optional[Tuple[Any, int]] = None, - prenormed: bool = False, - pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, - defer_moe_tail: bool = False, - ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: - """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). - - ``block_residual`` is a preallocated snapshot bank in kernel-native - ``[K_max, M, H]`` layout. Returns the running prefix sum and the - number of valid bank rows. - - ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux - stream for the layer BEFORE this one: the aggregated stream for layer j - is by definition what its next consumer sees, so the mixture computed - below already is it. Reading it here beats recomputing it, and is only - possible because K3 asserts pp_size == 1 -- layer j+1 is always local. - PP support would need a recompute at the rank boundary. - - ``step`` is the step's classification (``decode_step``), handed to the - attention module. - - ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is - this layer's input norm, and the layer's input, the step's embedding, - already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). - - The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens - runs the attention's all-reduce and the residual update in one - collective once the target has built its ``decode_comm``: with o_proj - (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, - the step and the attention's decode branch, else on the attention's - unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). - - ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the - previous layer's MoE output, unreduced (``defer_moe_tail``); then - ``hidden_states`` is the prefix sum without it, and this layer's - pre-attention step reduces and adds it: a ``PendingTail`` in one - ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its - all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer - the ``sandwich_tail`` call also stores the tap into the capture slot - where the speculative metadata exposes it (``capture_view``); other - tapped layers take the tap after the step. - - ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` - instead, ``partial`` this layer's MoE output unreduced, for the next - consumer to reduce. - """ - prefix_sum = hidden_states - valid_block_residual = block_residual[:num_snapshots] - attn_res_max_tokens = _attn_res_max_tokens(step) - tail = ( - pending_moe_partial - if isinstance(pending_moe_partial, _decode_comm.PendingTail) - else None - ) - # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum - # straight into the bank row it snapshots. - snapshot_row = None - if tail is not None and self.layer_idx % self.attn_res_block_size == 0: - snapshot_row = block_residual[num_snapshots] - # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the - # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). - tap_view = None - if tail is not None and capture is not None: - tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) - - if prenormed: - assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 - assert pending_moe_partial is None and capture is None - elif tail is not None: - hidden_states, prefix_sum = self.decode_comm.sandwich_tail( - tail, - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - updated_out=snapshot_row, - tap=tap_view, - tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, - ) - elif pending_moe_partial is not None: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - attn_res_max_tokens, - ) - elif capture is not None: - # The mixture tap needs the PRE-norm value, which the fused - # attn-res + RMSNorm kernel does not expose. Keep the two steps - # split on captured layers only and fuse everywhere else. - if num_snapshots > 0: - hidden_states = _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers - # acceptance, silently. hidden_states is the pre-norm attn_res mixture; - # prefix_only wants the running prefix, already in hand as prefix_sum. - tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - hidden_states = self.input_layernorm(hidden_states) - elif num_snapshots > 0: - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - attn_res_max_tokens, - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - if capture is not None and pending_moe_partial is not None and tap_view is None: - # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output - # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. - tapped = ( - _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - if _AUX_ATTN_RES_STREAM_ENABLED - else prefix_sum - ) - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - - if self.layer_idx % self.attn_res_block_size == 0: - if not prenormed and snapshot_row is None: - block_residual[num_snapshots].copy_(prefix_sum) - num_snapshots += 1 - valid_block_residual = block_residual[:num_snapshots] - prefix_sum = None - attention = self.linear_attn if self.is_kda else self.self_attn - comm = self.decode_comm - if comm is not None and comm.takes_post_attention(hidden_states, step): - if ( - self.sandwich_oproj - and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS - and attention.will_run_decode_branch(attn_metadata, step) - ): - core = attention(hidden_states, attn_metadata, step=step, project_output=False) - hidden_states, prefix_sum = comm.sandwich_oproj( - core, - self._o_proj().weight, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) - hidden_states, prefix_sum = comm.allreduce_attn_res( - partial, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - if comm is not None and step.wide: - partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) - attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) - else: - attention_output = attention(hidden_states, attn_metadata, step=step) - if prefix_sum is None: - prefix_sum = attention_output - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - attn_res_max_tokens, - ) - else: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - attention_output, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - attn_res_max_tokens, - ) - all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) - if self.is_moe and defer_moe_tail: - partial = self.block_sparse_moe( - hidden_states, all_rank_num_tokens, partial_tail=True, step=step - ) - return prefix_sum, num_snapshots, partial - if self.is_moe: - hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) - else: - hidden_states = self._dense_mlp(hidden_states, step) - - prefix_sum = prefix_sum + hidden_states - return prefix_sum, num_snapshots - - def forward_generic( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - attn_metadata: AttentionMetadata, - capture: Optional[Tuple[Any, int]] = None, - ) -> Tuple[torch.Tensor, int]: - prefix_sum = hidden_states - valid_block_residual = block_residual[:num_snapshots] - if capture is not None: - if num_snapshots > 0: - hidden_states = _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - hidden_states = self.input_layernorm(hidden_states) - elif num_snapshots > 0: - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - if self.layer_idx % self.attn_res_block_size == 0: - block_residual[num_snapshots].copy_(prefix_sum) - num_snapshots += 1 - valid_block_residual = block_residual[:num_snapshots] - prefix_sum = None - attention = self.linear_attn if self.is_kda else self.self_attn - comm = self.decode_comm - if comm is not None and comm.takes_post_attention(hidden_states, None): - partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) - hidden_states, prefix_sum = comm.allreduce_attn_res( - partial, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - elif prefix_sum is None: - prefix_sum = attention.forward_generic(hidden_states, attn_metadata) - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - attention.forward_generic(hidden_states, attn_metadata), - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - if self.is_moe: - all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) - hidden_states = self.block_sparse_moe.forward_generic( - hidden_states, all_rank_num_tokens - ) - else: - hidden_states = self.mlp(hidden_states) - return prefix_sum + hidden_states, num_snapshots - - def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: - """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode - kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" - gemvs = self.decode_gemvs - if gemvs is not None and step.small: - out = gemvs.dense_mlp( - hidden_states, - self.mlp.gate_up_proj.weight, - self.mlp.down_proj.weight, - *self._situ, - ) - if out is not None: - return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out - return self.mlp(hidden_states) - - def _mnnvl_allreduce(self): - """The MNNVL all-reduce of this layer's attention output, or None.""" - attention = self.linear_attn if self.is_kda else self.self_attn - return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) - - def _o_proj(self) -> nn.Module: - """This layer's attention output projection (row parallel).""" - return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj - - def _o_allreduce(self) -> AllReduce: - """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" - return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce - - def accepts_moe_partial(self, num_snapshots: int) -> bool: - """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target - built its collective state, and the snapshot bank is not empty.""" - return num_snapshots > 0 and self.decode_comm is not None - - def skip_forward( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. - - ``modeling_utils.skip_forward()`` only drops a module's weights when it - finds this attribute, so without it the layer-wise benchmarks would - allocate all 93 layers instead of the profiled slice. - """ - return hidden_states, block_residual - - -class KimiLinearModel(DecoderModel): - def __init__(self, model_config: ModelConfig): - super().__init__(model_config) - cfg = _get_text_config(model_config.pretrained_config) - self._text_cfg = cfg - dtype = torch.bfloat16 - - # Attention and MoE phases are sequential, so their branch-overlap - # roles share one stream; MoE-internal overlap roles remain separate. - aux_stream_list = [torch.cuda.Stream() for _ in range(4)] - self.aux_stream_dict = { - AuxStreamType.Attention: aux_stream_list[0], - AuxStreamType.MoeShared: aux_stream_list[0], - AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], - AuxStreamType.MoeBalancer: aux_stream_list[2], - AuxStreamType.MoeOutputMemset: aux_stream_list[3], - } - - self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) - self.layers = nn.ModuleList( - [ - KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) - for layer_idx in range(cfg.num_hidden_layers) - ] - ) - self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) - - # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) - # by _apply_attn_res and the fused attn_res op. - self.output_attn_res_norm = KimiK3RMSNorm( - cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - self.num_attn_res_snapshots = ( - cfg.num_hidden_layers + cfg.attn_res_block_size - 1 - ) // cfg.attn_res_block_size - # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the - # weights are final. None: every step embeds on the generic path. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - # Which convention the drafter tap is on is not recoverable from the - # served output -- a mismatch only lowers acceptance -- so state it once - # at construction rather than leaving it to be inferred from an AL. - logger.info_once( - "Kimi K3 aux hidden capture: mode=" - f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " - f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", - key="kimi_k3_aux_capture_mode", - ) - - # >>> route B: no per-token KDA verify states (kda_draft_records): no speculative decoding - @property - def kda_draft_records(self) -> bool: - """False: this target decodes without speculation, so the hybrid cache manager keeps no KDA state after - each verify token, and a captured step's latent-push decision (``_latent_push``) reads that.""" - return False - - # <<< route B - - def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, - from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from - load-time state, so every rank and every replay of the graph decides alike.""" - if not torch.cuda.is_current_stream_capturing(): - return False - kda = [layer.linear_attn for layer in self.layers if layer.is_kda] - mla = [layer.self_attn for layer in self.layers if not layer.is_kda] - return latent_push( - step, - capturing=True, - breakable=is_in_breakable_cuda_graph(), - kda_draft_records=self.kda_draft_records, - kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), - mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), - ) - - def _defer_moe_tail( - self, - i: int, - hidden_states: torch.Tensor, - num_snapshots: int, - spec_metadata, - capture_set, - step: DecodeStep, - ) -> bool: - """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE - takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after - the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps - its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated - tail.""" - layer = self.layers[i] - if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): - return False - if spec_metadata is not None and ( - capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) - ): - return False - return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - spec_metadata=None, - **kwargs, - ) -> torch.Tensor: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] - step = decode_step(attn_metadata, num_tokens) - if step is None: - return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) - if self._latent_push(attn_metadata, step): - step = replace(step, latent_push=True) - # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. - prenormed = None - if ( - inputs_embeds is None - and self.decode_gemvs is not None - and len(self.layers) > 0 - and self.num_attn_res_snapshots > 0 - ): - table = self.embed_tokens.weight - block_residual = table.new_empty( - self.num_attn_res_snapshots, num_tokens, table.shape[1] - ) - prenormed = self.decode_gemvs.embed_norm( - input_ids, table, self.layers[0].input_layernorm, block_residual - ) - if prenormed is not None: - hidden_states = prenormed - else: - hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - block_residual = hidden_states.new_empty( - self.num_attn_res_snapshots, - hidden_states.shape[0], - hidden_states.shape[1], - ) - num_snapshots = 0 - capture_set = getattr(spec_metadata, "_capture_layer_set", None) - # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) - # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. - pending_moe_partial = None - for i, layer in enumerate(self.layers): - defer_moe_tail = self._defer_moe_tail( - i, hidden_states, num_snapshots, spec_metadata, capture_set, step - ) - outputs = layer( - hidden_states, - block_residual, - num_snapshots, - attn_metadata, - step=step, - capture=self._capture(spec_metadata, capture_set, i), - prenormed=i == 0 and prenormed is not None, - pending_moe_partial=pending_moe_partial, - defer_moe_tail=defer_moe_tail, - ) - if defer_moe_tail: - hidden_states, num_snapshots, pending_moe_partial = outputs - else: - (hidden_states, num_snapshots), pending_moe_partial = outputs, None - - if isinstance(pending_moe_partial, _decode_comm.PendingTail): - normed, _ = self.layers[-1].decode_comm.sandwich_tail( - pending_moe_partial, - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - ) - return normed - if pending_moe_partial is not None: - _, normed = _apply_attn_res_add_and_rmsnorm( - hidden_states, - _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - _attn_res_max_tokens(step), - ) - return normed - return self._final_norm( - hidden_states, - block_residual, - num_snapshots, - spec_metadata, - capture_set, - _attn_res_max_tokens(step), - ) - - def _generic_forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor], - inputs_embeds: Optional[torch.FloatTensor], - spec_metadata, - ) -> torch.Tensor: - hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - block_residual = hidden_states.new_empty( - self.num_attn_res_snapshots, - hidden_states.shape[0], - hidden_states.shape[1], - ) - num_snapshots = 0 - capture_set = getattr(spec_metadata, "_capture_layer_set", None) - for i, layer in enumerate(self.layers): - hidden_states, num_snapshots = layer.forward_generic( - hidden_states, - block_residual, - num_snapshots, - attn_metadata, - capture=self._capture(spec_metadata, capture_set, i), - ) - return self._final_norm( - hidden_states, block_residual, num_snapshots, spec_metadata, capture_set - ) - - def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: - # DFlash/DSpark hidden-state capture. The drafter is distilled on - # the aggregated stream value -- the pre-norm softmax mixture its - # next consumer sees -- not on the raw prefix sum a layer returns, - # which is SGLang's fallback for models without the - # attention-residual scheme. Capturing the prefix sum costs 4.5pt - # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). - # The tap fires inside layer i+1, which computes that tensor - # anyway; see its forward docstring. Ground truth: SGLang - # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 - # aggregate_stream_torch. - if ( - spec_metadata is not None - and i > 0 - and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) - ): - return (spec_metadata, self.layers[i - 1].layer_idx) - return None - - def _final_norm( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - spec_metadata, - capture_set, - max_fused_tokens: Optional[int] = None, - ) -> torch.Tensor: - # The last layer has no successor, so this one recompute is - # unavoidable -- output-side score weights, matching SGLang's - # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set - # against 93 layers; kept so a set that does include the final layer - # gets the right tensor rather than the raw prefix sum. - if spec_metadata is not None and len(self.layers) > 0: - last = self.layers[-1] - if capture_set is None or last.layer_idx in capture_set: - tail = ( - _apply_attn_res( - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - ) - if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED - else hidden_states - ) - spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) - - return _apply_attn_res_and_rmsnorm( - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - max_fused_tokens, - ) - - -# ---------------------------------------------------------------------------------------------------------------------- -# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. -# ---------------------------------------------------------------------------------------------------------------------- - - -class K3DecodeKDA(KimiKDALinearAttention): - # >>> route B: no verify kernels - """Kimi K3's KDA attention: the built-in module, with the decode kernels on the steps they take. - - * A decode step of one token per request runs the fused input projection and the plain decode in one - ``ssm/k3_kda_decode_attn`` launch. - * On every step ``decode_step`` classifies, ``o_proj`` runs on the ``o_proj`` decode GEMV site where it takes the - rows, then the module's all-reduce. - - Every other step runs the built-in module. The kernels read one ``[q | k | v | g | f_a | b]`` weight, built at the - checkpoint load from the module's own, and the device's ``K3KdaBuffers`` and decode GEMVs' state, which the - target sets in ``post_load_weights``. - """ - - # <<< route B - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - # The fused [q | k | v | g | f_a | b | pad] projection weight; the six projections' weights and the built-in - # fused [q | k | v | g] and [f_a | b | pad] ones are views of it. - self.k3_proj_weight: Optional[torch.Tensor] = None - # The fused projection's Lamport buffers: one set per device, shared by every KDA layer. - self.k3_buffers: Optional[K3KdaBuffers] = None - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - @property - def takes_k3_kernels(self) -> bool: - """Whether the decode kernels can run this layer: its fused projection weight is built.""" - return self.k3_proj_weight is not None - - def finalize_decode_weights(self) -> None: - """The built-in fused weights, then one ``[q | k | v | g | f_a | b | pad]`` weight of both.""" - super().finalize_decode_weights() - assert ( - self.use_full_rank_gate - and self.gate_lower_bound is not None - and self._qkvg_proj_weight is not None - and self._bfa_proj_weight is not None - and self._qkvg_proj_weight.dtype == self._bfa_proj_weight.dtype == torch.bfloat16 - ), ( - f"Kimi K3 KDA layer {self.layer_idx}: the decode kernels read the bf16 fused projections the built-in " - "module builds on CUDA at head dim 128, with a full-rank output gate and a gate lower bound" - ) - rows = self._qkvg_proj_weight.shape[0] - with torch.no_grad(): - fused = self._merge_projection_weights( - (self.q_proj, self.k_proj, self.v_proj, self.g_proj, self.f_a_proj, self.b_proj), - pad_rows_to=8, - ) - self.k3_proj_weight = fused - self._qkvg_proj_weight, self._bfa_proj_weight = fused[:rows], fused[rows:] - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether ``forward`` runs ``step`` on the decode branch: outside a breakable CUDA graph.""" - return not is_in_breakable_cuda_graph() - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - reduce_output: bool = True, - project_output: bool = True, - ) -> torch.Tensor: - """``forward_generic`` where ``will_run_decode_branch`` does not hold. On the decode branch: the plain decode - on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the built-in dispatch; then - ``o_proj`` on its decode GEMV site and the TP all-reduce. - - ``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce), and, on the decode branch only, - ``project_output=False`` the post-o_norm core ``[N, H * 128]`` (no ``o_proj``).""" - if not self.will_run_decode_branch(attn_metadata, step): - return self.forward_generic(hidden_states, attn_metadata, reduce_output) - if ( - step.decode - and step.tokens_per_request == 1 - and self.takes_k3_kernels - and self.k3_buffers is not None - ): - core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) - else: - core = self._forward_impl(hidden_states, attn_metadata) - if not project_output: - return core.reshape(-1, self.proj_size) - return self._k3_project_output(core, reduce_output) - - def forward_generic( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - reduce_output: bool = True, - ) -> torch.Tensor: - if reduce_output: - return super().forward(hidden_states, attn_metadata) - return self.o_proj( - self._builtin_core(hidden_states, attn_metadata).reshape(-1, self.proj_size) - ) - - def _builtin_core( - self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata - ) -> torch.Tensor: - """The built-in forward's core ``[N, H, 128]``: inside a breakable CUDA graph from the eager - ``maybe_bcg_kda_core_inplace``, else from ``_forward_impl``.""" - if self.register_to_config and is_in_breakable_cuda_graph(): - core = hidden_states.new_empty( - (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 - ) - maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) - return core - return self._forward_impl(hidden_states, attn_metadata) - - def _k3_project_output(self, core: torch.Tensor, reduce_output: bool = True) -> torch.Tensor: - """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then, with - ``reduce_output``, the TP all-reduce.""" - core2d = core.reshape(-1, self.proj_size) - out = None - if self.decode_gemvs is not None: - out = self.decode_gemvs.project("o_proj", core2d, self.o_proj.weight) - if out is None: - if reduce_output: - return self._project_output(core) - out = self.o_proj(core2d) - if reduce_output and self._o_allreduce is not None: - out = self._o_allreduce(out) - return out - - def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: - """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; - each slot's conv window and state advance in place.""" - mamba_metadata = attn_metadata.mamba_metadata - slots = getattr(mamba_metadata, "generation_state_indices", None) - if slots is None: - slots = mamba_metadata.state_indices[: x.shape[0]] - layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) - w_q, w_k, w_v = self._get_mtp_conv_weights() - core = k3_kda_decode_attn( - x.contiguous(), - self.k3_proj_weight, - self.f_b_proj.weight, - w_q, - w_k, - w_v, - self._A_log_f32, - self._dt_bias_f32, - self._onorm_w_f32, - layer_cache.conv, - layer_cache.temporal, - slots, - self.k3_buffers, - float(self.gate_lower_bound), - self.head_k_dim**-0.5, - float(self.o_norm.eps), - ) - # Speculative decoding's replay caches keep their committed conv window in step with the pool's. - self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) - return core - - # >>> route B: no verify kernels (a verify step needs speculative decoding) - # <<< route B - - -class K3DecodeMLA(KimiK3MLAAttention): - """Kimi K3's MLA attention: the built-in module, with a decode step's attention on the decode kernels. - - A decode step runs ``x [W_a; W_g]^T`` with the gate columns through a sigmoid (the ``mla_ag`` decode GEMV site - where it takes the rows, else one GEMM), then ``attention/k3_mla_qkv`` (the q_a / kv_a RMSNorms, q_b and the k_b - absorption into the fused query, the step's latent rows stored into the paged cache) and - ``attention/k3_mla_attn_vb_out`` (the attention over the paged cache, v_b and the output gate in one launch), then - ``o_proj`` (the ``o_proj`` site where it takes the rows, else the module). Every other step, and a decode step whose - cache the kernels do not read (``k3_mla_decode_view`` says why), runs the built-in module. - - ``[W_a; W_g]`` is built at load from the module's weights, which become views of it. The attention workspace is - the device's ``K3MlaAttnWorkspace``; it and the decode GEMVs' state are set by the target in - ``post_load_weights``. - """ - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - # [W_a; W_g]: the fused q_a / kv_a projection's rows, then the output gate's. - self.k3_ag_weight: Optional[torch.Tensor] = None - # The decode attention's workspace: one per device, shared by every MLA layer. - self.k3_workspace: Optional[K3MlaAttnWorkspace] = None - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - def post_load_weights(self) -> None: - """The built-in post-load, then ``[W_a; W_g]`` (once: CUDA graphs captured since read it).""" - super().post_load_weights() - if self.k3_ag_weight is not None: - return - gaps = self._k3_layout_gaps() - assert not gaps, ( - f"Kimi K3 MLA layer {self.layer_idx}: the decode kernels do not take {'; '.join(gaps)}" - ) - qkv_a, gate = self.kv_a_proj_with_mqa, self.g_proj - rows = qkv_a.weight.shape[0] - with torch.no_grad(): - fused = torch.cat([qkv_a.weight, gate.weight]) - qkv_a.weight = nn.Parameter(fused[:rows], requires_grad=False) - gate.weight = nn.Parameter(fused[rows:], requires_grad=False) - self.k3_ag_weight = fused - - def _k3_layout_gaps(self) -> list: - """What of this layer the decode kernels do not take (empty when they take all of it).""" - if not (self.use_output_gate and self.fuse_qkv_a_proj and not self.is_lite): - return ["a layer without the output gate or the fused q_a / kv_a projection"] - linears = (self.kv_a_proj_with_mqa, self.g_proj, self.q_b_proj, self.o_proj) - checks = ( - (not self.mapping.has_cp_helix(), "helix context parallelism"), - (not self.apply_rotary_emb and not self.llama_4_scaling, "RoPE or llama-4 scaling"), - (self.sparse_attn_hooks is None, "sparse attention"), - ( - self.kv_cache_dtype != "fp8_ds_mla" - and not getattr(self.mqa, "has_fp8_kv_cache", False) - and not getattr(self.mqa, "has_fp4_kv_cache", False), - "a quantized KV cache", - ), - ( - all(m.weight.dtype == torch.bfloat16 and m.bias is None for m in linears) - and self.k_b_proj_trans.dtype == self.v_b_proj.dtype == torch.bfloat16, - "projections other than bf16 and unbiased", - ), - ( - not getattr(self.q_a_layernorm, "is_nvfp4", False) - and not getattr(self.kv_a_layernorm, "use_gemma", False), - "an NVFP4 q_a norm or a Gemma kv_a norm", - ), - ( - self.num_heads_tp % 6 == 0 - and self.kv_lora_rank == 512 - and self.qk_rope_head_dim == 64 - and self.q_lora_rank == 1536, - f"{self.num_heads_tp} heads, latent {self.kv_lora_rank}, rope {self.qk_rope_head_dim}, " - f"q_lora {self.q_lora_rank}", - ), - ) - return [why for ok, why in checks if not ok] - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether ``forward`` runs ``step`` on the decode kernels (see ``_k3_step_view``). Only there does it return - the gated attention output before ``o_proj``; the KV writes forbid running the attention twice, so a caller - that needs it asks first.""" - return self._k3_step_view(attn_metadata, step) is not None - - def _k3_step_view(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> Optional[dict]: - """The paged-cache view the decode kernels read on ``step``: a decode step, this layer's ``[W_a; W_g]`` and - workspace built, outside a breakable CUDA graph, and a cache ``k3_mla_decode_view`` takes. Else None.""" - if not ( - step.decode - and self.k3_ag_weight is not None - and self.k3_workspace is not None - and not is_in_breakable_cuda_graph() - ): - return None - return self._k3_decode_view(attn_metadata, step.num_tokens) - - def forward( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - project_output: bool = True, - ) -> torch.Tensor: - """``forward_generic``, except on a decode step whose cache the decode kernels read - (``will_run_decode_branch``). There only, ``project_output=False`` returns the gated attention output - ``[M, H * 128]``, the input of ``o_proj``.""" - view = self._k3_step_view(attn_metadata, step) - if view is None: - if not project_output: - raise ValueError("project_output=False needs a step will_run_decode_branch takes") - return self.forward_generic(position_ids, hidden_states, attn_metadata) - x = hidden_states[: step.num_tokens].contiguous() - rows = self.kv_a_proj_with_mqa.weight.shape[0] - gemvs = self.decode_gemvs - ag = None if gemvs is None else gemvs.project("mla_ag", x, self.k3_ag_weight) - if ag is None: - ag = torch.nn.functional.linear(x, self.k3_ag_weight) - ag[:, rows:].sigmoid_() - fused_q = k3_mla_qkv( - ag, - self.q_a_layernorm.weight, - float(self.q_a_layernorm.variance_epsilon), - self.q_b_proj.weight, - self.k_b_proj_trans, - self.kv_a_layernorm.weight, - float(self.kv_a_layernorm.variance_epsilon), - view["pool"], - view["row_stride"], - view["page_table"], - view["page_offset"], - view["seq_len"], - ) - attn_output = self.create_output(x, 0) - k3_mla_attn_vb_out( - fused_q, - view["pool"], - view["row_stride"], - view["page_table"], - view["page_offset"], - view["seq_len"], - view["softmax_scale"], - self.v_b_proj, - attn_output, - self.k3_workspace, - gate=ag, - gate_col0=rows, - ) - if not project_output: - return attn_output - out = None if gemvs is None else gemvs.project("o_proj", attn_output, self.o_proj.weight) - if out is None: - out = self._project_output([attn_output], position_ids, attn_metadata, None) - return out - - def forward_generic( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - return super().forward(position_ids, hidden_states, attn_metadata) - - def _k3_decode_view(self, attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[dict]: - """The paged latent cache as the decode kernels read it this step, or None when they do not read it (the - reason is logged once).""" - view = k3_mla_decode_view(self.mqa, attn_metadata, num_tokens) - if isinstance(view, str): - logger.info_once( - f"Kimi K3 MLA: the built-in path for a decode step the decode kernels do not read ({view})", - key=f"k3_mla_decode_view_{view}", - ) - return None - return view - - -# ---------------------------------------------------------------------------------------------------------------------- -# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. -# ---------------------------------------------------------------------------------------------------------------------- - -#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at -#: max_draft_len 7. -DRAFTER_ATTN_SPLITS = frozenset( - {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} -) -# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k -# RMSNorm's epsilon and the NeoX RoPE's base. -_DRAFTER_HEADS = (6, 1) -_DRAFTER_HEAD_DIM = 64 -_DRAFTER_PAGE = 64 -_DRAFTER_EPS = 1e-5 -_DRAFTER_ROPE_BASE = 10000.0 - - -def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: - """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input - feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" - if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: - return None - width = in_features // tp_size - return tp_rank * width, (tp_rank + 1) * width - - -class K3FcSlice(nn.Module): - """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is - ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's - ranks is ``fc``'s output.""" - - def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: - super().__init__() - self.weight = nn.Parameter(weight, requires_grad=False) - self.start = start - self.end = end - - @classmethod - def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: - """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns - (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" - columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) - if columns is None: - return None - start, end = columns - return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the - GEMM reads the column block in place, through the rows' stride.""" - return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) - - -class K3DSparkDrafter(GQADSparkForCausalLM): - """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. - - A block the entries take runs, per layer: - - * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; - * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to - its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the - worker writes the accepted rows' context k / v before a block reads them; - * the output projection on ``drafter_o``, then the module's all-reduce; - * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the - module's all-reduce. - - The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state - (below); a projection whose site does not take its rows runs its module. Every other block runs the stock - ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or - normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under - CUDA-graph capture. - - Where the target hands the drafter the TP group's collective state (``use_decode_comm``): - - * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and - ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; - * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the - post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the - block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the - projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection - after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is - followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers - whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. - - Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the - stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); - the context k / v and the Markov head stay upstream code. - """ - - def __init__( - self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" - ) -> None: - super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. - self._k3_take: Optional[bool] = None - # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes - # only these. - self._k3_attn_ran: set = set() - # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). - self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None - # The zero residual of the split context projection's all-reduce, up to a decode step's rows. - self._k3_zero_rows: Optional[torch.Tensor] = None - # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms - # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. - self._k3_norms_fuse = False - self._k3_sandwich_forms: frozenset = frozenset() - - def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: - """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the - target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are - loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. - - * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they - split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` - sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. - * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and - RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose - shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not - compile is not used.""" - self.decode_comm = comm - self._k3_split_fc() - self._k3_norms_fuse = self._k3_norms_take_comm() - forms = set() - if self._k3_norms_fuse: - layers = self.model.layers - for form, swiglu, weights in ( - ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), - ("down", True, [layer.mlp.down_proj.weight for layer in layers]), - ): - if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( - weights[0], swiglu=swiglu - ): - forms.add(form) - self._k3_sandwich_forms = frozenset(forms) - logger.info( - "Kimi K3 DSpark drafter: residual adds and RMSNorms " - + ( - f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " - "mnnvl_fusion_allreduce)" - if self._k3_norms_fuse - else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" - ) - ) - - def _k3_norms_take_comm(self) -> bool: - """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and - post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no - quantized output), and every output and down projection is row parallel with its all-reduce.""" - hidden = self.config.hidden_size - - def plain(norm: nn.Module) -> bool: - weight = getattr(norm, "weight", None) - return ( - isinstance(weight, torch.Tensor) - and weight.dtype == torch.bfloat16 - and tuple(weight.shape) == (hidden,) - and hasattr(norm, "variance_epsilon") - and not getattr(norm, "use_gemma", True) - and not getattr(norm, "is_nvfp4", True) - and not getattr(norm, "return_hp_output", True) - ) - - def reduces(linear: nn.Module) -> bool: - return ( - getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" - and getattr(linear, "reduce_output", False) - and getattr(linear, "all_reduce", None) is not None - ) - - return plain(self.model.norm) and all( - plain(layer.input_layernorm) - and plain(layer.post_attention_layernorm) - and reduces(layer.self_attn.o_proj) - and reduces(layer.mlp.down_proj) - for layer in self.model.layers - ) - - def _k3_tp_all_reduce(self) -> Optional[nn.Module]: - """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" - o_proj = self.model.layers[0].self_attn.o_proj - if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": - return None - return getattr(o_proj, "all_reduce", None) - - def _k3_split_fc(self) -> None: - """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero - residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is - a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP - all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" - fc = getattr(self, "fc", None) - if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): - return - mapping = self.model_config.mapping - weight = getattr(fc, "weight", None) - sliced = None - if ( - isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.dtype == torch.bfloat16 - and getattr(fc, "bias", None) is None - and self._k3_tp_all_reduce() is not None - ): - sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) - if sliced is None: - logger.info( - "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " - f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" - ) - return - self.fc = sliced - self._k3_zero_rows = weight.new_zeros( - MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] - ) - logger.info( - f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " - f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" - ) - - def load_weights(self, weights, weight_mapper=None, **kwargs): - """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" - result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) - self._k3_split_fc() - return result - - def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. - - With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group - with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode - step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then - ``hidden_norm``. Otherwise the stock projection.""" - fc = self.fc - if not isinstance(fc, K3FcSlice): - return super().project_target_hidden(hidden_states) - partial = fc(hidden_states.to(fc.weight.dtype)) - rows = partial.shape[0] - zeros = self._k3_zero_rows - if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): - normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) - return normed - if rows == 0: - return self.hidden_norm(partial) - return self.hidden_norm(self._k3_tp_all_reduce()(partial)) - - def dflash_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_k_cache: torch.Tensor, - ctx_v_cache: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor] = None, - ctx_page_table: Optional[torch.Tensor] = None, - ctx_rows_start: Optional[int] = None, - ) -> torch.Tensor: - """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class - docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows - ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows - as a view instead of gathering them.""" - keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) - if keys is None: - return super().dflash_forward( - noise_embedding, - query_positions, - num_ctx_per_req, - ctx_k_cache, - ctx_v_cache, - ctx_cache_batch_idx, - ctx_kv_cache, - ctx_page_table, - ) - out = self._k3_block_forward( - noise_embedding, - query_positions, - num_ctx_per_req, - ctx_cache_batch_idx, - ctx_kv_cache, - ctx_page_table, - ctx_rows_start, - ) - if not torch.cuda.is_current_stream_capturing(): - self._k3_attn_ran |= keys - return out - - def _k3_layers_take(self) -> bool: - """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head - layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, - non-causal unwindowed block attention, and no block convolution or post-attention gate.""" - if self._k3_take is None: - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - take = ( - self._use_fused_qk_norm_rope - and not self.has_block_conv - and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate - ) - for layer_idx, layer in enumerate(self.model.layers): - attn, mlp = layer.self_attn, layer.mlp - pos = getattr(attn, "pos_embd_params", None) - rope = getattr(pos, "rope", None) - norms = (attn.q_norm, attn.k_norm) - take = take and ( - (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS - and attn.head_dim == _DRAFTER_HEAD_DIM - and getattr(attn, "is_qk_norm", False) - and not getattr(attn, "use_gemma_rms_norm", True) - and rope is not None - and pos.is_neox - and getattr(pos, "mrope_section", None) is None - and rope.theta == _DRAFTER_ROPE_BASE - and all( - norm.variance_epsilon == _DRAFTER_EPS - and norm.weight.dtype == torch.bfloat16 - and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) - for norm in norms - ) - and attn.qkv_proj.bias is None - and attn.o_proj.bias is None - and mlp.activation is torch.nn.functional.silu - and mlp.swiglu_limit is None - and mlp.swiglu_alpha in (None, 1.0) - and mlp.swiglu_beta in (None, 0.0) - and mlp.gate_up_proj.bias is None - and mlp.down_proj.bias is None - and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) - ) - self._k3_take = bool(take) - logger.info( - f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " - + ( - "on k3_drafter_attn_qknorm and the drafter GEMV sites" - if self._k3_take - else "on the stock block forward (a layer the drafter entries do not take)" - ) - ) - return self._k3_take - - def _k3_block_keys( - self, - noise_embedding: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor], - ctx_page_table: Optional[torch.Tensor], - ) -> Optional[frozenset]: - """The attention's compile keys for this block when the drafter entries take it, else None.""" - if ( - self.dflash_attention_backend != "TRTLLM" - or ctx_kv_cache is None - or ctx_page_table is None - or noise_embedding.dtype != torch.bfloat16 - or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS - or is_in_breakable_cuda_graph() - or not self._k3_layers_take() - ): - return None - num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads - page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) - dense = ( - num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, - _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, - ) - strides = set() - for layer_idx in range(len(self.model.layers)): - cache = ctx_kv_cache[layer_idx] - if ( - cache.dtype != torch.bfloat16 - or cache.dim() != 5 - or tuple(cache.shape[1:]) != page - or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) - or cache.stride(0) % 8 - ): - return None - strides.add(cache.stride(0)) - keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) - if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: - return None - return keys - - def _k3_block_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: torch.Tensor, - ctx_page_table: torch.Tensor, - ctx_rows_start: Optional[int] = None, - ) -> torch.Tensor: - batch, block = noise_embedding.shape[:2] - rows = batch * block - ctx_len = num_ctx_per_req[:batch].to(torch.int32) - if ctx_rows_start is None: - page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) - else: - # The batch's page-table rows are one contiguous run: a view, no gather. - page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] - positions = query_positions.reshape(-1).contiguous() - hidden = noise_embedding.reshape(rows, -1) - layers = self.model.layers - comm = self._k3_fused_comm(hidden) - residual = None - if comm is not None: - # The fused all-reduces read the residual and return the updated one: the block's input serves as the - # first residual, uncopied. - residual = hidden - normed = layers[0].input_layernorm(hidden) - for layer_idx, layer in enumerate(layers): - attn = layer.self_attn - if comm is None: - if residual is None: - residual = hidden.clone() - normed = layer.input_layernorm(hidden) - else: - normed, residual = layer.input_layernorm(hidden, residual) - qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) - out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) - k3_drafter_attn_qknorm( - qkv, - attn.q_norm.weight, - attn.k_norm.weight, - positions, - attn.q_norm.variance_epsilon, - attn.pos_embd_params.rope.theta, - ctx_kv_cache[layer_idx], - page_table, - ctx_len, - attn.num_heads, - attn.num_key_value_heads, - out, - ) - if comm is not None: - # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after - # the last layer the final norm. - next_norm = ( - layers[layer_idx + 1].input_layernorm - if layer_idx + 1 < len(layers) - else self.model.norm - ) - normed, residual = self._k3_project_norm( - comm, out, attn.o_proj, residual, layer.post_attention_layernorm - ) - normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) - continue - hidden = self._k3_project("drafter_o", out, attn.o_proj) - hidden, residual = layer.post_attention_layernorm(hidden, residual) - hidden = self._k3_mlp(layer.mlp, hidden) - if comm is not None: - return normed - out, _ = self.model.norm(hidden, residual) - return out - - def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: - """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its - all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace - holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" - comm = self.decode_comm - if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): - return None - return comm - - def _k3_project_norm( - self, - comm: _decode_comm.K3DecodeComm, - x: torch.Tensor, - linear: nn.Module, - residual: torch.Tensor, - norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, - the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the - projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` - where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" - if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( - x, linear.weight, residual, norm - ): - return comm.sandwich_plain(x, linear.weight, residual, norm) - gemvs = self.decode_gemvs - partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) - if partial is None: - partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) - return comm.allreduce_norm(partial, residual, norm) - - def _k3_mlp_norm( - self, - comm: _decode_comm.K3DecodeComm, - mlp: nn.Module, - x: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down - projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul - form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down - projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; - then ``comm/mnnvl_fusion_allreduce``.""" - gemvs = self.decode_gemvs - gate_up = ( - None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) - ) - partial = None - if gate_up is not None: - if "down" in self._k3_sandwich_forms and comm.takes_plain( - gate_up, mlp.down_proj.weight, residual, norm, swiglu=True - ): - return comm.sandwich_plain( - gate_up, mlp.down_proj.weight, residual, norm, swiglu=True - ) - partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) - if partial is None: - partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) - return comm.allreduce_norm(partial, residual, norm) - - def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: - """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel - projection's all-reduce.""" - y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) - if y is None: - return linear(x) - return _row_parallel_reduce(linear, y) - - def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: - """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on - ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" - gemvs = self.decode_gemvs - gate_up = ( - None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) - ) - down = ( - None - if gate_up is None - else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) - ) - if down is None: - return mlp(x) - return _row_parallel_reduce(mlp.down_proj, down) - - -def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: - """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" - all_reduce = getattr(linear, "all_reduce", None) - if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: - return all_reduce(partial) - return partial - - -# ---------------------------------------------------------------------------------------------------------------------- -# The target: step classification, the construction checks and the class route returns. -# ---------------------------------------------------------------------------------------------------------------------- - - -def _text_model_config(model_config: ModelConfig) -> ModelConfig: - """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. - - The checkpoint names its language-model modules `language_model.`; the text model's are ``, with - `layers.*` under `model.`. - """ - config = model_config.pretrained_config - text = copy.copy(model_config) - text._frozen = False - text.pretrained_config = config.text_config - excluded = text.quant_config.exclude_modules - if excluded: - text.quant_config = copy.copy(text.quant_config) - renamed = [] - for name in excluded: - if name.startswith(_LANG_PREFIX): - name = name[len(_LANG_PREFIX) :] - if name.startswith("layers."): - name = "model." + name - renamed.append(name) - text.quant_config.exclude_modules = renamed - text.skip_create_weights_in_init = True - text._frozen = True - return text - - -@dataclass(frozen=True) -class DecodeStep: - """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` - generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" - - num_tokens: int - num_requests: Optional[int] = None - tokens_per_request: Optional[int] = None - #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` - #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. - latent_push: bool = False - - @property - def small(self) -> bool: - """Whether the step fits one token tile of the token-count kernels.""" - return self.num_tokens <= DECODE_MAX_TOKENS - - @property - def decode(self) -> bool: - """Whether the step is a pure decode step the request-aware kernels take.""" - return self.num_requests is not None - - @property - def wide(self) -> bool: - """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode - layout's MoE head and tail, on M-general ops.""" - return self.decode and not self.small - - -def latent_push( - step: DecodeStep, - *, - capturing: bool, - breakable: bool, - kda_draft_records: bool, - kda_decode_kernels: bool, - mla_decode_branch: bool, -) -> bool: - """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at - most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a - breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels - (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA - layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). - Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel - between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a - graph launch orders every pushing replay behind whatever ran before it.""" - if not (step.decode and step.small) or not capturing or breakable: - return False - if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): - return False - return mla_decode_branch - - -def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: - """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). - - ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers - and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input - here is fixed by that shape, so a captured step and its replays are classified alike. - """ - if num_tokens <= 0: - return None - requests = _decode_requests(attn_metadata, num_tokens) - if requests is not None: - return DecodeStep(num_tokens, requests, num_tokens // requests) - if num_tokens <= DECODE_MAX_TOKENS: - return DecodeStep(num_tokens) - return None - - -def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: - """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" - if attn_metadata.num_contexts != 0: - return None - requests = attn_metadata.num_generations - if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: - return None - tokens = num_tokens // requests - if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: - return None - seq_lens = getattr(attn_metadata, "seq_lens", None) - if seq_lens is not None and seq_lens.device.type == "cpu": - lens = seq_lens[:requests] - if lens.numel() != requests or bool((lens != tokens).any()): - return None - return requests +_TARGET = K3TargetSpec(sm=(10, 0), quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, moe_tp=16, moe_ep=1) def _check_construction(model_config: ModelConfig) -> None: - """The settings this target is built for that are fixed before the first step.""" - capability = torch.cuda.get_device_capability() - assert capability == _SM, ( - f"this target is certified on sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" - ) - mapping = model_config.mapping - topology = ( - mapping.world_size, - mapping.tp_size, - mapping.pp_size, - mapping.moe_tp_size, - mapping.moe_ep_size, - mapping.enable_attention_dp, - ) - # >>> route B: this target's topology, its expert split set explicitly, and no speculative decoding - assert topology == (16, 16, 1, 16, 1, False), ( - "the tp16_moetp16ep1 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " - "moe_tensor_parallel_size 16, moe_expert_parallel_size 1 and enable_attention_dp false; the engine built " - f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" - ) - assert mapping.moe_tp_ep_user_specified, ( - "the tp16_moetp16ep1 target needs moe_tensor_parallel_size 16 and moe_expert_parallel_size 1 set " - "explicitly; with the split unset the engine resolves the same sizes, but Kimi K3 then runs the experts " - "expert-parallel over the 16 ranks" - ) - assert getattr(model_config, "spec_config", None) is None, ( - "the tp16_moetp16ep1 target decodes without speculation; for DSpark, DFlash or SA set " - "moe_tensor_parallel_size 4 and moe_expert_parallel_size 4 (the tp16_moetp4ep4 target)" - ) - # <<< route B - assert model_config.torch_dtype == torch.bfloat16, ( - f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" - ) - kv_algo = model_config.quant_config.kv_cache_quant_algo - assert kv_algo is None, ( - f"this target's MLA kernels read a bf16 KV pool; kv_cache_config.dtype resolved to {kv_algo}" - ) - quant_algo = model_config.quant_config.quant_algo - assert quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( - "this target loads the MXFP4 checkpoint, whose compressed-tensors config the model config reads as " - "W4A16_MXFP4 with no per-layer declarations (its routed experts run the W4A8_MXFP4_MXFP8 default); the " - f"engine read {quant_algo} with {len(model_config.quant_config_dict or {})} per-layer declarations" - ) - strategy = model_config.allreduce_strategy - assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( - f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" - ) - - -# >>> route B: this target's name -class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(KimiLinearForCausalLM): - # <<< route B - """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. - - It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through - weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model - keeps the built-in one's. - """ - - @classmethod - def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: - """The V2 hybrid manager holds the KDA states; the step contract requires it.""" - return "V2" - - def __init__(self, model_config: ModelConfig): - config = model_config.pretrained_config - assert getattr(config, "text_config", None) is not None, ( - "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " - "text_config" - ) - _check_construction(model_config) - text = _text_model_config(model_config) - # >>> route B: no speculative decoding (_check_construction), so the one-engine shell builds no drafter - # <<< route B - # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 - # weight-read conversion of the shared / latent MLPs. - self._fp8_weight_read_moe_mlp = False - self.mapping_with_cp = None - self._repurposed_tp_mapping = None - cfg = text.pretrained_config - SpecDecOneEngineForCausalLM.__init__( - self, - KimiLinearModel(text), - text, - hidden_size=cfg.hidden_size, - vocab_size=cfg.vocab_size, - ) - self._own_spec_branch(text) - self._step_checked = False - # >>> route B: no drafter to hand the LM head to - # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. - self.logits_processor = _decode_gemv.K3LogitsProcessor(self.logits_processor) - # <<< route B - # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which - # must therefore be the text config, as the built-in wrapper leaves it. - model_config._frozen = False - model_config.pretrained_config = self.config - model_config._frozen = True - - def _own_spec_branch(self, model_config: ModelConfig) -> None: - """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a - standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is - `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" - spec_config = model_config.spec_config - if spec_config is None: - return - mode = spec_config.spec_dec_mode - standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) - if ( - standalone_dspark - and self.draft_config is not None - and not is_mla(self.draft_config.pretrained_config) - ): - drafter = K3DSparkDrafter( - self.draft_config, dflash_attention_backend=spec_config.attention_backend - ) - if mode.is_parallel_draft(): - drafter.logits_processor = self.logits_processor - self.epilogue[self.epilogue.index(self.draft_model)] = drafter - self.draft_model = drafter - if mode.is_dflash(): - worker_cls = _spec_worker.KimiK3DFlashWorker - elif standalone_dspark: - worker_cls = _spec_worker.KimiK3DSparkWorker - else: - return - worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) - worker.set_draft_model(self.draft_model) - self.epilogue[self.epilogue.index(self.spec_worker)] = worker - self.spec_worker = worker - - def load_weights(self, weights, *args, **kwargs): - _weights.load(self, weights) - - def cache_derived_state(self) -> None: - """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of - every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: - a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites - are among them, and the drafter gets the state too.""" - super().cache_derived_state() - if self.model.decode_gemvs is not None: - return - gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head) - self.model.decode_gemvs = gemvs - self.logits_processor.gemvs = gemvs - for layer in self.model.layers: - if not layer.is_moe: - layer.decode_gemvs = gemvs - if isinstance(self.draft_model, K3DSparkDrafter): - self.draft_model.decode_gemvs = gemvs - - def post_load_weights(self) -> None: - """The state the decode kernels share, built once per device before any CUDA-graph capture and handed to - every layer of its kind: the KDA projection's Lamport buffers and the MLA decode attention's workspace; the - decode GEMVs' state (built by ``cache_derived_state``) handed to every attention module; and, where every - attention all-reduce runs over MNNVL, the TP group's collective state (``K3DecodeComm``, collective: every - rank builds it here) handed to every layer, with whether its sandwich takes the layer's o_proj, the decode - path's one-shot ceiling on every stock MNNVL all-reduce (``use_decode_one_shot``), and the MoE decode path - (``_build_decode_moe``); then the speculative worker's decode kernels (``_gate_spec_worker_kernels``) and the - DSpark drafter's collectives (``_gate_drafter_comm``).""" - super().post_load_weights() - kda = [layer.linear_attn for layer in self.model.layers if layer.is_kda] - mla = [layer.self_attn.mixer for layer in self.model.layers if not layer.is_kda] - if kda[0].k3_buffers is not None: - return # built by an earlier call; CUDA graphs captured since hold it - device = self.model.embed_tokens.weight.device - buffers = K3KdaBuffers.create(device) - workspace = K3MlaAttnWorkspace.create(device, mla[0].num_heads_tp // 6) - for module in kda: - module.k3_buffers = buffers - for module in mla: - module.k3_workspace = workspace - for module in kda + mla: - module.decode_gemvs = self.model.decode_gemvs - layers = self.model.layers - comm = None - if all(layer._mnnvl_allreduce() is not None for layer in layers): - for layer in layers: - layer.sandwich_oproj = _decode_comm.K3DecodeComm.takes_oproj( - layer._o_proj(), self.model.num_attn_res_snapshots - ) - oproj = next((layer._o_proj().weight for layer in layers if layer.sandwich_oproj), None) - comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, oproj) - for layer in layers: - layer.decode_comm = comm - _decode_comm.use_decode_one_shot(self) - moe_layers = 0 if comm is None else self._build_decode_moe(comm) - logger.info( - # >>> route B: no verify kernels - "Kimi K3 decode kernels: KDA on k3_kda_decode_attn " - # <<< route B - f"({sum(m.takes_k3_kernels for m in kda)} / {len(kda)} layers take them), MLA on k3_mla_qkv and " - f"k3_mla_attn_vb_out ({len(mla)} layers); the post-attention all-reduce of at most " - f"{_decode_comm.AR_ATTN_RES_MAX_TOKENS} tokens " - + ( - "with the residual update: k3_sandwich_oproj with o_proj at most " - f"{_decode_comm.SANDWICH_MAX_TOKENS} decode tokens " - f"({sum(layer.sandwich_oproj for layer in layers)} / {len(layers)} layers), else " - "mnnvl_allreduce_attn_res" - if comm is not None - else "unfused (an attention all-reduce does not run over MNNVL)" - ) - # >>> route B: this target's MoE engines - + f"; MoE on k3_moe_front, k3_moe_m1 / k3_moe_m2 / k3_moe and the row-parallel tail ({moe_layers} layers)" - # <<< route B - ) - self._gate_spec_worker_kernels(comm) - self._gate_drafter_comm(comm) - - def _build_decode_moe(self, comm: _decode_comm.K3DecodeComm) -> int: - """The MoE decode path (``decode_moe.py``) on every MoE layer it takes: the shared state (collective: the - front's head workspace), each layer's decode weights with its latent norm folded into the latent up - projection, the decode GEMVs' state; then one call of every kernel of the path (collective: the front and - the sandwich tail). Every rank builds it here. Returns the number of MoE layers it takes.""" - mapping = self.model_config.mapping - moes = [layer.block_sparse_moe for layer in self.model.layers if layer.is_moe] - gaps = { - id(moe): _decode_moe.layout_gaps( - moe, mapping.tp_size, self.model.num_attn_res_snapshots - ) - for moe in moes - } - takes = [moe for moe in moes if not gaps[id(moe)]] - for reason in sorted({gap for moe in moes for gap in gaps[id(moe)]}): - logger.info_once( - f"Kimi K3 MoE decode path off on some layers: {reason}", - key=f"k3_decode_moe_off_{reason}", - ) - if not takes: - return 0 - backend = takes[0].routed_experts.backend - state = _decode_moe.K3DecodeMoe.create( - mapping, - backend.w3_w1_weight.device, - backend.w3_w1_weight.shape[1] // 2, - backend.expert_size_per_partition, - comm.mnnvl, - # >>> route B: the experts' logical slice (192 of the loader's 256), which k3_moe_m1 / k3_moe_m2 stream - i_logical=( - getattr(backend.quant_method, "intermediate_size_per_partition_lean", None) - or backend.intermediate_size_per_partition - ), - # <<< route B - ) - for moe in takes: - _decode_moe.fold_latent_norm(moe) - moe.decode_moe = _decode_moe.K3DecodeMoeLayer.create( - moe, state, mapping.tp_rank, mapping.tp_size - ) - moe.decode_gemvs = self.model.decode_gemvs - first = takes[0].decode_moe - first.warm_up(takes[0]) - logger.info( - "Kimi K3 MoE decode path: the latent all-reduce at <= 8 tokens on " - + ( - "the routed experts' push form and k3_latent_reduce on pushing steps (a pure decode step captured " - "into a CUDA graph on the decode kernels), else the routed experts' all-reduce" - if state.exchange is not None - else f"the routed experts' all-reduce (no latent exchange at TP {mapping.tp_size})" - ) - ) - comm.compile_tail(takes[0].moe_hidden_size, first.shared_cols, first.tail_weight) - return len(takes) - - def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: - """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, - ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input - that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's - ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the - vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks - each step's own conditions. Returns the setting (False without such a worker).""" - worker = getattr(self, "spec_worker", None) - if not hasattr(worker, "k3_decode"): - return False - gemvs = self.model.decode_gemvs - worker.k3_decode = ( - comm is not None and gemvs is not None and gemvs.head_workspace is not None - ) - logger.info( - f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " - + ( - "on" - if worker.k3_decode - else "off (no MNNVL decode state or no k3_head_gemv LM head)" - ) - ) - return worker.k3_decode + _check_target_construction(model_config, _TARGET) - def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: - """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built - it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context - projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms - in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on - the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate - requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` - workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the - drafter took the state (False without such a drafter).""" - drafter = getattr(self, "draft_model", None) - if comm is None or not isinstance(drafter, K3DSparkDrafter): - return False - drafter.use_decode_comm(comm) - return True - def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: - """First-forward checks of the engine surface and the per-engine settings.""" - objects = { - "attn_metadata": attn_metadata, - "kv_cache_manager": attn_metadata.kv_cache_manager, - } - missing = [ - f"{owner}.{name}" - for owner, names in REQUIRED_ENGINE_FIELDS.items() - for name in names - if not hasattr(objects[owner], name) - ] - assert not missing, f"engine fields this target reads are missing: {missing}" - assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( - f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " - f"{attn_metadata.tokens_per_block}" - ) - manager = attn_metadata.kv_cache_manager - assert isinstance(manager, MambaHybridCacheManagerV2), ( - "this target needs the V2 hybrid KV / state cache manager " - "(kv_cache_config.use_kv_cache_manager_v2); the engine built " - f"{type(manager).__name__}" - ) - assert not manager.enable_block_reuse, ( - "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" - ) - kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) - state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype - assert state_dtype == torch.float32, ( - "this target's KDA kernels keep fp32 recurrent states (kv_cache_config.mamba_ssm_cache_dtype float32 or " - f"auto); the engine built a {state_dtype} state pool" - ) - self._step_checked = True +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(ModelingV2KimiK3): + """Concrete deployment of the shared Kimi K3 text model.""" - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: bool = False, - spec_metadata=None, - resource_manager=None, - **kwargs, - ) -> torch.Tensor: - if kwargs.pop("multimodal_params", None): - raise ValueError( - "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" - ) - if not self._step_checked: - self._check_step_contract(attn_metadata) - return super().forward( - attn_metadata, - input_ids, - position_ids, - inputs_embeds, - return_context_logits, - spec_metadata, - resource_manager, - **kwargs, - ) + target_spec = _TARGET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py index 66a1f4f72049..24233d194709 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py @@ -1,44 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# >>> route B: this target's expert slice -"""Weight loading: Kimi K3 (MXFP4) / sm_100 / tp16_moetp16ep1. +"""The shared Kimi K3 text checkpoint loader.""" -The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and -`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: - -* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the - routed experts one at a time, keeps this rank's slice (a sixteenth of the width of every expert under moe_tp 16 x - moe_ep 1: 192 of 3072, zero-padded to 256 for the MoE kernels' tiles), and checks its own key coverage. -* The vision tower and the projector are a predicted non-load: listed here, never read. -* Any other key fails the load, naming it, rather than being dropped. -""" -# <<< route B - -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.models.modeling_utils import filter_weights - -_LANG_PREFIX = "language_model." - -# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. -PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") - - -def load(model, weights) -> None: - """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" - unknown = sorted( - k - for k in weights.keys() - if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) - ) - assert not unknown, ( - f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " - f"first {unknown[:5]}" - ) - lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) - assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" - checkpoint_dir = getattr(weights, "checkpoint_dir", None) - if checkpoint_dir is not None: - lm_weights.checkpoint_dir = checkpoint_dir - lm_weights.checkpoint_prefix = _LANG_PREFIX - KimiLinearForCausalLM.load_weights(model, lm_weights) +from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py index 474e865afce4..e1e4bb7413de 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4.""" +"""Kimi K3 MXFP4 / sm_100 / tp16_moetp4ep4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py deleted file mode 100644 index f5327daf9474..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py +++ /dev/null @@ -1,402 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode path's residual-update collectives on the catalog's Kimi K3 MNNVL and sandwich entries. - -A layer's post-attention step adds the attention output to the running prefix sum, selects the attention residual -and applies the post-attention RMSNorm. Under TP the attention output is a sum over the TP group, so the step is that -all-reduce followed by the residual update. On a step of at most `AR_ATTN_RES_MAX_TOKENS` tokens (wide decode steps -aside) one collective runs both: - -* `K3DecodeComm.sandwich_oproj`, `comm/k3_sandwich_oproj`: o_proj, its all-reduce and the residual update in one - kernel, at most `SANDWICH_MAX_TOKENS` tokens of an o_proj of the TP16 per-rank shape [7168, 768]. The attention - hands over its gated o_proj input. -* `K3DecodeComm.allreduce_attn_res`, `comm/mnnvl_allreduce_attn_res`, everywhere else: the one-shot MNNVL all-reduce - of the attention's unreduced o_proj partial, with the residual update as its epilogue. - -The pre-attention step after a MoE layer whose row-parallel tail was handed on (a `PendingTail`, `decode_moe.py`) is -the same kind of collective: `K3DecodeComm.sandwich_tail`, `comm/k3_sandwich_tail`, runs the tail GEMV, its -all-reduce and the next layer's residual update (the final norm's, after the last layer) in one kernel. - -The DSpark drafter (`K3DSparkDrafter`) runs its collectives with a plain residual add and RMSNorm on the same state: - -* `K3DecodeComm.sandwich_plain`, `comm/k3_sandwich_plain`: a row-parallel projection of at most `SANDWICH_MAX_TOKENS` - tokens (a layer's attention output projection, or its MLP's SiLU-and-mul and down projection), its all-reduce, the - residual add and the RMSNorm in one kernel. -* `K3DecodeComm.allreduce_norm`, `comm/mnnvl_fusion_allreduce`: the all-reduce of an unreduced projection output of - any token count the MNNVL workspace holds, with the residual add and the RMSNorm in its epilogue (the context - projection's with a zero residual). - -The state is one `MnnvlWorkspace` and one `K3SandwichWorkspace` of the TP group (`K3DecodeComm.create`): collective -over the group and eager, built by the target in `post_load_weights` before any CUDA-graph capture. Every rank must -make the same calls on each in the same order. Which call a step takes is decided from its token count and kind and -from the load-time layout alone, which every rank of the group shares. - -The plain MNNVL all-reduces the decode path keeps (the stock modules') send one-shot up to -`DECODE_AR_ONE_SHOT_MAX_BYTES` (`use_decode_one_shot`), except a wide decode step's (`wide_all_reduce`). -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import NamedTuple, Optional, Tuple - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_oproj import ( - K3SandwichWorkspace, - k3_sandwich_oproj, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_plain import ( - k3_sandwich_plain, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_tail import ( - k3_sandwich_tail, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allreduce_attn_res import ( - mnnvl_allreduce_attn_res, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_fusion_allreduce import ( - mnnvl_fusion_allreduce, - required_buffer_bytes, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( - MnnvlWorkspace, -) - -# The kernel's support predicate: metadata reads only. -from tensorrt_llm._torch.cute_dsl_kernels.k3_sandwich import op as _sandwich_op -from tensorrt_llm._torch.distributed import AllReduceParams - -# The most tokens of a step whose post-attention all-reduce runs one-shot with the residual update in its epilogue. -AR_ATTN_RES_MAX_TOKENS = 16 - -# The most tokens the sandwich kernel takes (one token tile). -SANDWICH_MAX_TOKENS = _sandwich_op.MAX_TOKENS - -# The MNNVL workspace's buffer size: a one-shot call pushes T x 7168 bf16 from each of 16 ranks, 3.5 MiB at T = 16. -MNNVL_BUFFER_BYTES = 4 << 20 - -# The one-shot ceiling of the stock MNNVL all-reduces on the decode path: 8 tokens x 7168 x 16 ranks x 2 B is -# 1.75 MiB, which the stock 1 MiB ceiling would send two-shot. -DECODE_AR_ONE_SHOT_MAX_BYTES = 4 << 20 - -# A wide decode step's ceiling: the stock 1 MiB. At 16 ranks two-shot is faster for every wide step's rows. -WIDE_AR_ONE_SHOT_MAX_BYTES = 1 << 20 - - -def use_decode_one_shot(model: nn.Module) -> None: - """Every stock MNNVL all-reduce of ``model`` sends one-shot up to `DECODE_AR_ONE_SHOT_MAX_BYTES`. Each grows its - workspace on its first eager call of a larger size, before the capture of that size.""" - for module in model.modules(): - mnnvl = getattr(module, "mnnvl_allreduce", None) - if mnnvl is not None: - mnnvl.one_shot_max_bytes = DECODE_AR_ONE_SHOT_MAX_BYTES - mnnvl.early_trigger = True - - -def skip_all_reduce() -> AllReduceParams: - """All-reduce parameters under which a row-parallel module returns this rank's unreduced output (its projection - without its all-reduce): a ``Linear``'s ``all_reduce_params``, a ``GatedMLP``'s ``final_all_reduce_params``.""" - return AllReduceParams(enable_allreduce=False) - - -def wide_all_reduce(all_reduce: nn.Module, x: torch.Tensor) -> torch.Tensor: - """``all_reduce(x)`` of a wide decode step (a stock ``AllReduce`` module, no fusion): its MNNVL all-reduce with - the `WIDE_AR_ONE_SHOT_MAX_BYTES` ceiling, else the module itself.""" - mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) - if mnnvl is not None: - out = mnnvl( - x.contiguous(), AllReduceParams(), one_shot_max_bytes=WIDE_AR_ONE_SHOT_MAX_BYTES - ) - if out is not None: - return out - return all_reduce(x) - - -class PendingTail(NamedTuple): - """A MoE layer's row-parallel tail left to its consumer's fused pre-attention step (``K3DecodeComm.sandwich_tail``): - the reduced latent ``[T, 3584]``, the shared experts' activation, the tail weight ``[latent up columns | padding | - shared down]``, this rank's first latent column and the latent norm's epsilon.""" - - latent: torch.Tensor - act: torch.Tensor - weight: torch.Tensor - lo: int - lat_eps: float - - -def _eps(norm: nn.Module) -> float: - """The epsilon of a KimiK3RMSNorm or ``torch.nn.RMSNorm`` (``eps``; for torch's None, its default: the machine - epsilon of the weight's dtype) or a stock RMSNorm (``variance_epsilon``).""" - eps = norm.eps if hasattr(norm, "eps") else norm.variance_epsilon - return float(torch.finfo(norm.weight.dtype).eps if eps is None else eps) - - -def _res_args(res_proj: nn.Module, res_norm: nn.Module, out_norm: nn.Module) -> tuple: - """The residual update's weights and epsilons in the entries' order.""" - return ( - res_proj.weight.reshape(-1), - res_norm.weight, - out_norm.weight, - _eps(res_norm), - _eps(out_norm), - ) - - -@dataclass(eq=False) -class K3DecodeComm: - """The decode path's collective state for the TP group: the `MnnvlWorkspace` and the `K3SandwichWorkspace` the - post-attention steps run on. Built by `create`; owned by the target and shared by every layer.""" - - mnnvl: MnnvlWorkspace - sandwich: K3SandwichWorkspace - - @classmethod - def create(cls, mapping, oproj_weight: Optional[torch.Tensor] = None) -> "K3DecodeComm": - """The state for ``mapping``'s TP group. Collective: every rank of the group calls it at the same point, - eagerly, before any CUDA-graph capture; each workspace fails on every rank or on none. - - ``oproj_weight``: an o_proj weight ``takes_oproj`` holds for. The sandwich kernel compiles here for its - shape, with one call on a zero row of a zero weight, so no capture compiles it; the call advances the - sandwich workspace on every rank alike.""" - state = cls( - MnnvlWorkspace.create(mapping, MNNVL_BUFFER_BYTES), - K3SandwichWorkspace.create(mapping), - ) - if oproj_weight is not None: - weight = torch.zeros_like(oproj_weight) - hidden = weight.shape[0] - ones = weight.new_ones(hidden) - k3_sandwich_oproj( - weight.new_zeros(1, weight.shape[1]), - weight, - None, - weight.new_zeros(0, 1, hidden), - weight.new_zeros(hidden), - ones, - ones, - 1e-6, - 1e-6, - state.sandwich, - ) - torch.cuda.synchronize(weight.device) - return state - - def compile_tail(self, latent_size: int, act_size: int, tail_weight: torch.Tensor) -> None: - """Compile the sandwich tail kernel for a MoE tail of ``latent_size`` latent and ``act_size`` shared columns - and ``tail_weight``'s shape, with one call on a zero row of a zero weight, before any capture. Collective: every - rank of the group makes the call; it advances the sandwich workspace on every rank alike.""" - weight = torch.zeros_like(tail_weight) - hidden = weight.shape[0] - ones = weight.new_ones(hidden) - k3_sandwich_tail( - weight.new_zeros(1, latent_size), - weight.new_zeros(1, act_size), - weight, - 0, - 1e-6, - None, - weight.new_zeros(0, 1, hidden), - weight.new_zeros(hidden), - ones, - ones, - 1e-6, - 1e-6, - self.sandwich, - ) - torch.cuda.synchronize(weight.device) - - def takes_post_attention(self, hidden_states: torch.Tensor, step) -> bool: - """Whether one of this state's collectives runs the post-attention step of a layer whose attention input is - ``hidden_states``: at most `AR_ATTN_RES_MAX_TOKENS` bf16 rows of a hidden size the MNNVL entry takes, on any - step but a wide decode step (``step.wide``), whose post-attention update stays the all-reduce and the fused - add + attn_res + RMSNorm.""" - if step is not None and step.wide: - return False - if hidden_states.dim() != 2 or hidden_states.dtype != torch.bfloat16: - return False - rows, hidden = hidden_states.shape - return ( - hidden % 1024 == 0 - and hidden <= 8192 - and 0 < rows <= min(AR_ATTN_RES_MAX_TOKENS, self.mnnvl.max_one_shot_tokens(hidden)) - ) - - @staticmethod - def takes_oproj(o_proj: nn.Module, max_snapshots: int) -> bool: - """Whether ``sandwich_oproj`` takes the layer of output projection ``o_proj`` on a step of at most - `SANDWICH_MAX_TOKENS` tokens, decided once the weights are final: a bias-free, contiguous bf16 weight of the - kernel's shape (the TP16 per-rank [7168, 768]) and a snapshot bank of at most ``max_snapshots`` rows the - kernel's candidate count holds.""" - weight = getattr(o_proj, "weight", None) - return ( - getattr(o_proj, "bias", None) is None - and isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.is_cuda - and _sandwich_op.supports(weight.new_empty((1, weight.shape[1])), weight, max_snapshots) - ) - - @staticmethod - def takes_tail( - latent: torch.Tensor, act: torch.Tensor, tail_weight: torch.Tensor, max_snapshots: int - ) -> bool: - """Whether ``sandwich_tail`` takes a MoE tail of these tensors' shapes (rows of ``latent`` and ``act``, at most - `SANDWICH_MAX_TOKENS`; ``tail_weight``) with a snapshot bank of at most ``max_snapshots`` rows: the TP16 - per-rank shapes.""" - return _sandwich_op.supports_tail(latent, act, tail_weight, max_snapshots) - - def allreduce_attn_res( - self, - partial: torch.Tensor, - prefix_sum: Optional[torch.Tensor], - block_residual: torch.Tensor, - res_proj: nn.Module, - res_norm: nn.Module, - out_norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` in one ``comm/mnnvl_allreduce_attn_res`` call: ``updated = prefix_sum + - allreduce(partial)`` (the sum alone without ``prefix_sum``), ``normed = out_norm(attn_res(block_residual..., - updated))``, the attention residual selected with ``res_proj`` and ``res_norm``. ``block_residual`` holds the - valid snapshots, ``[S, T, H]``.""" - return mnnvl_allreduce_attn_res( - partial.contiguous(), - prefix_sum, - block_residual, - *_res_args(res_proj, res_norm, out_norm), - self.mnnvl, - ) - - def sandwich_tail( - self, - pending: PendingTail, - prefix_sum: Optional[torch.Tensor], - block_residual: torch.Tensor, - res_proj: nn.Module, - res_norm: nn.Module, - out_norm: nn.Module, - updated_out: Optional[torch.Tensor] = None, - tap: Optional[torch.Tensor] = None, - tap_updated: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` as ``allreduce_attn_res`` of a MoE layer's row-parallel tail ``pending`` - (``[RMSNorm(latent)[:, lo:lo + 224] | act] @ weight.T``), in one ``comm/k3_sandwich_tail`` call. - ``updated_out``: a bf16 ``[T, H]`` tensor the call stores ``updated`` into (the consumer's snapshot bank row), - returned as ``updated``. ``tap``: a bf16 ``[T, H]`` view (a speculative capture slot) the call also stores - the pre-norm attention-residual mixture into, or ``updated`` with ``tap_updated``.""" - return k3_sandwich_tail( - pending.latent, - pending.act, - pending.weight, - pending.lo, - pending.lat_eps, - prefix_sum, - block_residual, - *_res_args(res_proj, res_norm, out_norm), - self.sandwich, - tap=tap, - tap_updated=tap_updated, - updated_out=updated_out, - ) - - def sandwich_oproj( - self, - core: torch.Tensor, - o_weight: torch.Tensor, - prefix_sum: Optional[torch.Tensor], - block_residual: torch.Tensor, - res_proj: nn.Module, - res_norm: nn.Module, - out_norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` as ``allreduce_attn_res`` of ``core @ o_weight.T``, in one ``comm/k3_sandwich_oproj`` - call: ``core`` is this rank's gated o_proj input ``[T, 768]``, ``o_weight`` its o_proj slice. Bit for bit - o_proj followed by ``allreduce_attn_res``.""" - return k3_sandwich_oproj( - core.contiguous(), - o_weight, - prefix_sum, - block_residual, - *_res_args(res_proj, res_norm, out_norm), - self.sandwich, - ) - - def compile_plain(self, weight: torch.Tensor, swiglu: bool = False) -> bool: - """Compile the plain sandwich (``sandwich_plain``) for a projection of ``weight``'s shape (with ``swiglu``, a - down projection after the SiLU-and-mul) with one call on a zero row of a zero weight, before any capture; - False, with nothing launched, where the kernel does not take that shape. Collective: every rank of the group - makes the call; it advances the sandwich workspace on every rank alike.""" - zero = torch.zeros_like(weight) - hidden, k_in = zero.shape - x = zero.new_zeros(1, 2 * k_in if swiglu else k_in) - residual = zero.new_zeros(1, hidden) - ones = zero.new_ones(hidden) - if not _sandwich_op.supports_plain(x, zero, residual, ones, swiglu): - return False - k3_sandwich_plain(x, zero, residual, ones, 1e-6, self.sandwich, swiglu=swiglu) - torch.cuda.synchronize(zero.device) - return True - - @staticmethod - def takes_plain( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - swiglu: bool = False, - ) -> bool: - """Whether ``sandwich_plain`` takes the call: at most `SANDWICH_MAX_TOKENS` contiguous bf16 rows ``x`` of a - row-parallel slice ``weight`` [7168, K] (K a multiple of 128 up to 896; with ``swiglu`` K 896 and ``x`` the - ``[gate | up]`` rows, 2 K wide), ``residual`` [rows, 7168], and ``norm``'s bf16 [7168] weight.""" - return _sandwich_op.supports_plain(x, weight, residual, norm.weight, swiglu) - - def sandwich_plain( - self, - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - swiglu: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` in one ``comm/k3_sandwich_plain`` call: ``updated = residual + allreduce(x @ - weight.T)`` (with ``swiglu``, of ``silu_and_mul(x) @ weight.T``), ``normed = norm(updated)`` for a plain - RMSNorm ``norm``. Bit for bit, by the kernel's statement, the projection on ``k3_ctm_gemv`` at split 1 (with - ``swiglu``, ``k3_ctm_gemv_swiglu`` at split 2) followed by ``allreduce_norm``'s one-shot call.""" - return k3_sandwich_plain( - x.contiguous(), - weight, - residual.contiguous(), - norm.weight, - _eps(norm), - self.sandwich, - swiglu=swiglu, - ) - - def takes_allreduce_norm(self, rows: int, hidden: int) -> bool: - """Whether ``allreduce_norm`` takes ``rows`` bf16 rows of ``hidden`` columns: the MNNVL workspace holds the - call at the decode path's one-shot ceiling (`DECODE_AR_ONE_SHOT_MAX_BYTES`, two-shot above it).""" - if rows <= 0 or hidden <= 0 or hidden % 8: - return False - world, buffer_bytes = self.mnnvl.world_size, self.mnnvl.buffer_bytes - need = required_buffer_bytes( - rows, hidden, world, torch.bfloat16, DECODE_AR_ONE_SHOT_MAX_BYTES - ) - two_shot = rows * hidden * world * 2 > DECODE_AR_ONE_SHOT_MAX_BYTES - return need <= buffer_bytes and not (two_shot and buffer_bytes % 32) - - def allreduce_norm( - self, partial: torch.Tensor, residual: torch.Tensor, norm: nn.Module - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` in one ``comm/mnnvl_fusion_allreduce`` call: ``updated = residual + - allreduce(partial)``, ``normed = norm(updated)`` for a plain RMSNorm ``norm``, sent one-shot up to - `DECODE_AR_ONE_SHOT_MAX_BYTES`. ``partial`` is this rank's unreduced ``[rows, hidden]`` bf16 output of a - row-parallel projection, of a shape ``takes_allreduce_norm`` holds for.""" - return mnnvl_fusion_allreduce( - partial.contiguous(), - self.mnnvl, - DECODE_AR_ONE_SHOT_MAX_BYTES, - residual.contiguous(), - norm.weight, - _eps(norm), - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py deleted file mode 100644 index 301b709643b3..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py +++ /dev/null @@ -1,465 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. - -* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest - at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, - `gemm/k3_ctm_gemv_wide`, `gemm/k3_ctm_gemv_long`, `gemm/k3_ctm_gemv` or `gemm/k3_ctm_gemv_swiglu`, and, where the - site lists it, more rows (up to `WIDE_ROWS`) on `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused - projections (MLA's [W_a; W_g] with the gate rows through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention - output projection, the built-in MLA path's q_a / kv_a, q_b and gate projections, the MoE decode path's projections - (`decode_moe.py`), two of them with fp32 outputs, and the DSpark drafter's q / k / v, output, gate / up and down - projections (`K3DSparkDrafter`). -* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on - `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the - stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the - drafter's logits on the same head take it too. -* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual - bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. -* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: - gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then - runs the down projection's all-reduce. - -Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. - -A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. -`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The -embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call -whose kernel has not run eagerly is refused. - -The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one -stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import Dict, Iterable, Optional, Set - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( - k3_ctm_gemv_long, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( - k3_ctm_gemv_swiglu, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( - k3_ctm_gemv_wide, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( - K3HeadGemvWorkspace, - k3_head_gemv, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm - -# The kernels' support predicates: metadata reads only. -from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op -from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op -from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op -from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op -from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE -from tensorrt_llm._utils import mpi_disabled - -# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 -# requests of 8 tokens). -MAX_ROWS = 8 -WIDE_ROWS = 64 - - -@dataclass(frozen=True) -class Site: - """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows - ("decode", "wide", "long", "ctm" for k3_ctm_gemv, or "swiglu" for k3_ctm_gemv_swiglu, whose rows are the - ``[gate | up]`` input of the SiLU-and-mul, 2 ``k`` wide), and k3_ctm_gemv_wide at MAX_ROWS+1..``wide_rows`` rows - where ``wide``. Output columns from ``sig_col0`` on are stored through a sigmoid; ``out_fp32``: an fp32 output - (k3_ctm_gemv_wide only). ``split`` / ``ring`` / ``push``: the CTM kernels' CTAs per 128-row weight tile, - k3_ctm_gemv_long's weight-ring stages, and whether the partial sums are pushed to each row's owner.""" - - n: int - k: int - small: str - wide: bool = False - sig_col0: int = -1 - split: int = 0 - ring: int = 0 - push: bool = False - out_fp32: bool = False - wide_rows: int = WIDE_ROWS - - -SITES: Dict[str, Site] = { - # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), - # the gate rows through a sigmoid. - "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), - # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). - "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), - # The attention output projection (row parallel; 6 heads x 128 in). - "o_proj": Site(7168, 768, "decode", wide=True), - # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. - "kv_a": Site(2112, 7168, "decode"), - "q_b": Site(1152, 1536, "wide"), - "g_proj": Site(768, 7168, "wide"), - # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. - "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), - "dense_down": Site(7168, 2112, "long", split=2, ring=6), - # The MoE decode path (decode_moe.py): this rank's head slice [latent down 224 | router 56] with an fp32 output, - # the shared experts' gate_up, the row-parallel tail [latent up 224 | padding 32 | shared down 384] and the - # replicated tail's latent up projection with an fp32 output. The wide kernel takes the head slice and the tail up - # to 32 rows, where it is faster than the stock GEMM. - "moe_head": Site(280, 7168, "wide", wide=True, out_fp32=True, wide_rows=32), - "moe_shared_gate_up": Site(768, 7168, "wide", wide=True), - "moe_tail": Site(7168, 640, "wide", wide=True, wide_rows=32), - "moe_up": Site(7168, 3584, "wide", out_fp32=True), - # The DSpark drafter's projections (K3DSparkDrafter; 6 query heads, 1 KV head of 64, MLP width 896): the fused - # q / k / v, the output projection (row parallel), gate_up [gate 896 | up 896], and the down projection with the - # SiLU-and-mul (row parallel). - "drafter_qkv": Site(512, 7168, "long", split=8, ring=6, push=True), - "drafter_o": Site(7168, 384, "ctm", split=1, push=True), - "drafter_gate_up": Site(1792, 7168, "long", split=8, ring=6, push=True), - "drafter_down": Site(7168, 896, "swiglu", split=2, push=True), -} - - -def _width(spec: Site) -> int: - """The row width a site's kernel reads: ``k``, or the ``[gate | up]`` input of a SiLU-and-mul site, 2 ``k``.""" - return 2 * spec.k if spec.small == "swiglu" else spec.k - - -def _wide_tile(rows: int) -> int: - from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel - - return k3_ctm_gemv_kernel.wide_tile(rows) - - -def _run( - spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor -) -> Optional[torch.Tensor]: - """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" - if kernel == "decode": - if spec.sig_col0 >= 0 or spec.out_fp32 or not _decode_op.supports(x2d, weight): - return None - return k3_decode_gemv(x2d, weight) - if kernel == "wide": - if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, spec.out_fp32): - return None - return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0, out_fp32=spec.out_fp32) - if spec.out_fp32: - return None - if kernel == "ctm": - if spec.sig_col0 >= 0 or not _ctm_op.supports(x2d, weight, spec.split): - return None - return k3_ctm_gemv(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) - if kernel == "swiglu": - if spec.sig_col0 >= 0 or not _ctm_op.supports_swiglu(x2d, weight, spec.split): - return None - return k3_ctm_gemv_swiglu(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) - # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. - sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count - if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( - x2d, weight, spec.split, spec.ring - ): - return None - return k3_ctm_gemv_long( - x2d, - weight, - sig_col0=spec.sig_col0, - split=spec.split, - ring=spec.ring, - trigger_early=True, - push=spec.push, - ) - - -def _capturing() -> bool: - return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() - - -def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: - """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" - if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: - return x2d.clone(memory_format=torch.contiguous_format) - return x2d - - -def _head_takes_module(lm_head: nn.Module) -> bool: - """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along - the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" - weight = getattr(lm_head, "weight", None) - mapping = getattr(lm_head, "mapping", None) - return ( - isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.dtype == torch.bfloat16 - and weight.is_contiguous() - and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" - and getattr(lm_head, "gather_output", False) - and getattr(lm_head, "gather_output_sizes", None) is None - and getattr(lm_head, "padding_size", None) == 0 - and getattr(lm_head, "bias", None) is None - and not getattr(lm_head, "has_any_quant", True) - and mapping is not None - and not mapping.enable_attention_dp - ) - - -def _plain_rmsnorm(norm: nn.Module) -> bool: - """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for - bit. An unknown module fails closed.""" - weight = getattr(norm, "weight", None) - return ( - IS_FLASHINFER_AVAILABLE - and isinstance(weight, torch.Tensor) - and weight.dtype == torch.bfloat16 - and not getattr(norm, "use_gemma", True) - and not getattr(norm, "is_nvfp4", True) - and not getattr(norm, "use_cuda_tile", True) - and not getattr(norm, "return_hp_output", True) - and getattr(norm, "nvfp4_scale", None) is None - and hasattr(norm, "variance_epsilon") - ) - - -class K3DecodeGemvs: - """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran - eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are - final; owned by the target.""" - - def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: - self.head_workspace = head_workspace - self._ran: Set[tuple] = set() - - @classmethod - def create( - cls, - lm_head: Optional[nn.Module] = None, - sites: Iterable[str] = tuple(SITES), - device: Optional[torch.device] = None, - ) -> "K3DecodeGemvs": - """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: - it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on - zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or - head whose kernel does not take its shape keeps the generic path.""" - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" - ) - head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None - if device is None: - device = ( - head_weight.device - if isinstance(head_weight, torch.Tensor) - else torch.device("cuda", torch.cuda.current_device()) - ) - state = cls() - for site in sites: - spec = SITES[site] - weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) - for rows in (1, 16, 32, 64) if spec.wide else (1,): - if rows > spec.wide_rows: - continue - state._project(site, weight.new_zeros(rows, _width(spec)), weight, warm=True) - del weight - if "dense_gate_up" in sites: - gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) - if _ctm_op.supports_situ_mul(gu): - for linear_beta in (None, 1.0): - k3_situ_mul(gu, 1.0, linear_beta) - state._ran.add(("situ_mul", linear_beta is not None)) - if lm_head is not None and _head_takes_module(lm_head): - x = head_weight.new_zeros(1, head_weight.shape[1]) - if _head_op.supports(x, head_weight): - workspace = K3HeadGemvWorkspace.create( - head_weight.shape[0], head_weight.shape[1], head_weight.device - ) - k3_head_gemv(x, head_weight, workspace) - state.head_workspace = workspace - state._ran.add(("lm_head",)) - torch.cuda.synchronize(device) - return state - - def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: - """``x @ weight.T`` (``[..., N]``: bf16, the site's sigmoid columns through the sigmoid; fp32 at an - ``out_fp32`` site; on a SiLU-and-mul site, ``silu(gate) * up`` of ``x = [gate | up]`` first) for ``site``'s - weight on its decode kernel, or None where none takes the call: more rows than the site's kernels take, - another shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its - module.""" - return self._project(site, x, weight, warm=False) - - def _project( - self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool - ) -> Optional[torch.Tensor]: - spec = SITES[site] - width = _width(spec) - if ( - weight.dtype != torch.bfloat16 - or tuple(weight.shape) != (spec.n, spec.k) - or not weight.is_contiguous() - or x.dtype != torch.bfloat16 - or x.dim() < 1 - or x.shape[-1] != width - ): - return None - rows = x.numel() // width - if 0 < rows <= MAX_ROWS: - kernel = spec.small - elif spec.wide and MAX_ROWS < rows <= spec.wide_rows: - kernel = "wide" - else: - return None - key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) - capturing = _capturing() - if capturing and not warm and key not in self._ran: - return None - y = _run(spec, kernel, _dense_rows(x.reshape(rows, width)), weight) - if y is None: - return None - if not capturing: - self._ran.add(key) - return y.view(*x.shape[:-1], spec.n) - - def lm_head_logits( - self, rows: torch.Tensor, lm_head: nn.Module, gather: bool = True - ) -> Optional[torch.Tensor]: - """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on - ``gemm/k3_head_gemv`` (``gather`` False: the shard alone, ``[M, vocab / TP]``); None where it does not take - the call (more than `MAX_ROWS` rows, another head, ...).""" - workspace = self.head_workspace - if workspace is None or not _head_takes_module(lm_head): - return None - weight = lm_head.weight - if ( - tuple(weight.shape) != (workspace.n_out, workspace.k_in) - or weight.device != workspace.partials.device - or rows.dim() != 2 - or rows.dtype != torch.bfloat16 - or not 0 < rows.shape[0] <= MAX_ROWS - or rows.shape[1] != workspace.k_in - ): - return None - if _capturing() and ("lm_head",) not in self._ran: - return None - group = lm_head.mapping.tp_group - if gather and len(group) > 1 and mpi_disabled(): - return None - x = _dense_rows(rows) - if not _head_op.supports(x, weight): - return None - local = k3_head_gemv(x, weight, workspace) - if not gather or len(group) == 1: - return local - gathered = allgather(local, None, group) - return torch.cat(torch.split(gathered, rows.shape[0], dim=0), dim=-1) - - def dense_mlp( - self, - x: torch.Tensor, - gate_up_weight: torch.Tensor, - down_weight: torch.Tensor, - beta: float, - linear_beta: Optional[float], - ) -> Optional[torch.Tensor]: - """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on - ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on - ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then - runs the module.""" - gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] - rows = x.shape[0] if x.dim() == 2 else 0 - if ( - not 0 < rows <= MAX_ROWS - or linear_beta == 0.0 - or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) - or tuple(down_weight.shape) != (down.n, down.k) - or gate_up.n != 2 * down.k - ): - return None - if ( - _capturing() - and not { - ("dense_gate_up", "long", 0), - ("dense_down", "long", 0), - ("situ_mul", linear_beta is not None), - } - <= self._ran - ): - return None - gu = self.project("dense_gate_up", x, gate_up_weight) - if gu is None or not _ctm_op.supports_situ_mul(gu): - return None - return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) - - def embed_norm( - self, - input_ids: torch.Tensor, - table: torch.Tensor, - norm: nn.Module, - bank: torch.Tensor, - ) -> Optional[torch.Tensor]: - """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one - ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not - apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose - kernel has not run eagerly).""" - if not _plain_rmsnorm(norm): - return None - ids = input_ids.reshape(-1) - if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): - return None - key = ("embed_norm", ids.numel(), ids.dtype) - capturing = _capturing() - if capturing and key not in self._ran: - return None - normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) - if not capturing: - self._ran.add(key) - return normed - - -class _K3Head: - """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" - - __slots__ = ("_gemvs", "_lm_head") - - def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: - self._gemvs = gemvs - self._lm_head = lm_head - - def __call__(self, rows: torch.Tensor) -> torch.Tensor: - logits = self._gemvs.lm_head_logits(rows, self._lm_head) - return self._lm_head(rows) if logits is None else logits - - -class K3LogitsProcessor(nn.Module): - """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. - - It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock - processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. - """ - - def __init__(self, stock: nn.Module) -> None: - super().__init__() - self.stock = stock - self.gemvs: Optional[K3DecodeGemvs] = None - - def forward( - self, - hidden_states: torch.Tensor, - lm_head: nn.Module, - attn_metadata, - return_context_logits: bool = False, - ) -> torch.Tensor: - head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) - return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) - - def lm_head_shard(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: - """This rank's bf16 vocabulary shard of ``lm_head(rows)`` on ``gemm/k3_head_gemv``, without the gather, or - None where ``gemvs`` does not take the rows. A speculative worker that keeps its logits vocabulary-sharded - computes them here, so they hold the values of the logits this processor gathers.""" - if self.gemvs is None: - return None - return self.gemvs.lm_head_logits(rows, lm_head, gather=False) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index 082f8798ae75..e2a3f736e8f6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -1,3595 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""ModelingV2 target: Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4. +"""Kimi K3 MXFP4, SM100, attention TP16 and MoE TP4 x EP4.""" -Kimi K3's language model: 93 layers, hidden 7168. Every fourth layer (3, 7, ..., 91) is MLA attention with 96 -query heads; the others are Kimi Delta Attention (KDA), a gated linear-attention recurrence with per-request state. -Layer 0's MLP is dense. Every other layer's MLP is a latent MoE: 896 routed experts, top-16, expert width 3584. -Attention residuals mix each sublayer's input from a bank of earlier outputs. The routed experts are MXFP4 (the -checkpoint's compressed-tensors default, read as W4A8_MXFP4_MXFP8); everything else is bf16, and so is the KV pool. - -`tp16_moetp4ep4` is `tensor_parallel_size: 16` with `moe_tensor_parallel_size: 4` and -`moe_expert_parallel_size: 4`, no attention data parallelism, all 16 GPUs in one NVLink domain. Attention is -head-split (6 MLA query heads per rank), and every rank holds a quarter of the width of a quarter of the experts. - -**Each step is classified once, on the host, from its shape** (`decode_step`, a `DecodeStep` or None), and the -classification decides which kernels each module runs: - -* **small**: at most 8 tokens (`DECODE_MAX_TOKENS`, one token tile), context requests included. The token-count - kernels take it: decode GEMVs, the MoE front and routed experts, the sandwiches, the embedding and residual - epilogues. -* **decode**: a pure decode step of R <= 8 generation requests with the same T <= 8 tokens each (one token without - speculation, 1 + 7 drafts with DSpark) and no context request, so at most 64 tokens. The request-aware kernels - take it: MLA attention and its KV store, the KDA verify, the drafter's attention. -* **wide**: a decode step of more than one token tile (DSpark verify of several requests). It keeps the decode - layout's MoE head and tail, on M-general ops. - -Every other step (prefill, mixed steps, decode steps above those bounds) runs the **generic path**: this target's -text model (`KimiLinearModel` below: decoder layers, attention residuals, the MLA / KDA / MoE runtimes), computed as -the built-in Kimi K3 text model computes it, on stock modules and ops that have no catalog entries yet. One weight -differs: where the MoE decode path takes a layer, its latent norm's weight is folded into the latent up projection -(the same function, rounded differently). And the routed experts' top-16 is computed outside the TRTLLM-Gen kernel -(`KimiK3MoeRoutingMethod`), with the kernel's arithmetic. `UNCERTIFIED_GENERIC_CALLS` names the stock code. - -The text model hands each step's classification to its attention modules, which run a **decode step** on the K3 -decode kernels' catalog entries: - -* KDA (`K3DecodeKDA`): one token per request, the fused input projection and the plain decode in one - `ssm/k3_kda_decode_attn` launch. Verify tokens (DFlash / DSpark at an even verify width up to 8): the cache manager - then keeps the last verify's draft records (per draft the row innovations, beta * k and the decay), and every - verify of the layer, on any step, runs the kernels that write them and replay the accepted drafts from them: - `ssm/k3_kda_attn` for one request of 8 tokens (the projection fused in), else `ssm/k3_kda_verify`. -* MLA (`K3DecodeMLA`): `attention/k3_mla_qkv` (the query path and the step's latent KV rows into the paged cache), - then `attention/k3_mla_attn_vb_out` (the attention, v_b and the output gate in one launch). - -The projections around them run on the decode GEMV sites of `decode_gemv.py` (the [W_a; W_g] and KDA verify-row -projections, `o_proj` on every classified step), as do the LM head, the embedding and layer 0's dense MLP. A -classified step's attention-residual epilogues (the selection and the RMSNorm after it) take the fused kernels up to -one token tile, 32 tokens on a wide decode step. On any step of at most 16 tokens but a wide one, the post-attention -all-reduce carries the residual update (`decode_comm.py`): at most 8 tokens on an attention decode branch, o_proj, the -all-reduce and the update are one `comm/k3_sandwich_oproj` kernel; otherwise the attention's unreduced o_proj output -goes through `comm/mnnvl_allreduce_attn_res`. The stock MNNVL all-reduces send one-shot up to 4 MiB, a wide decode -step's up to the stock 1 MiB. The state those kernels share (the KDA projection's Lamport buffers, the MLA attention -workspace, the decode GEMVs' state, the TP group's MNNVL and sandwich workspaces, the MoE path's) lives in typed -objects this target creates in `post_load_weights`, before any graph capture. - -A MoE layer on a step of at most 8 tokens runs `decode_moe.py`: `moe/k3_moe_front` (this rank's head slice, its -all-gather, the routing, the MXFP8 latent and the shared experts' gate_up + SiTU in one kernel), `moe/k3_moe`, the -latent all-reduce (on a pushing step, `DecodeStep.latent_push`, `moe/k3_moe`'s push form plus `comm/k3_latent_reduce`), -then the row-parallel tail, which the next layer's pre-attention step (the final norm's, after the -last layer) runs with its all-reduce and residual update as one `comm/k3_sandwich_tail` kernel. A wide decode step's -MoE keeps the sharded head and the row-parallel tail, on `moe/k3_route_quant`, `moe/k3_moe` and M-general ops. - -**What this target asserts rather than adapts**: SM 10.0; the topology above; the MXFP4 checkpoint's quantization -(W4A16_MXFP4 with no per-layer declarations, so the routed experts run the W4A8_MXFP4_MXFP8 default and the excluded -modules stay bf16); bf16 weights and a bf16 KV pool; -tokens_per_block 64 (the MLA generation kernels K3's 96 heads reach exist only at 64); the V2 hybrid KV / state -manager, which holds the KDA states, with block reuse off and fp32 recurrent states; an all-reduce strategy of AUTO or -MNNVL. The construction-time ones fail in `__init__`, the per-engine ones on the first forward, each naming the -setting. A layer the decode kernels do not take fails the weight load. - -**Text only.** The checkpoint is the vision-language wrapper. This target builds and loads no vision tower (its -weights are a predicted non-load, `weights.py`), and a step carrying multimodal input raises. - -**Speculative decoding** goes through the stock one-engine shell: DSpark or DFlash with an external drafter -checkpoint, and SA. The DSpark drafter is this target's `K3DSparkDrafter`, the stock GQA drafter with a decode step's -block on the drafter entries (`attention/k3_drafter_attn_qknorm` and the drafter's decode GEMV sites), which this -target puts in place of the shell's stock one. Where every attention all-reduce runs over MNNVL, the drafter also runs -on the TP group's collective state: its context projection's `fc` is split by input feature over the group, -`hidden_norm` applied in the all-reduce of the partial products (`comm/mnnvl_fusion_allreduce`), and a block's -residual adds and RMSNorms run in its all-reduces (`comm/k3_sandwich_plain` with the projection up to 8 rows, else -`comm/mnnvl_fusion_allreduce`). The DFlash / DSpark worker is this target's (`spec_worker.py`): the stock worker with -the decode kernels `trtllm::k3_spec_accept`, `k3_ctx_kv` and `k3_markov` behind its `k3_decode`. -""" - -from __future__ import annotations - -import copy -import math -import os -from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( - k3_drafter_attn_qknorm, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( - k3_mla_attn_vb_out, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( - K3MlaAttnWorkspace, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import k3_mla_qkv -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_attn import k3_kda_attn -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import K3KdaBuffers -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( - k3_kda_decode_attn, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify -from tensorrt_llm._torch.attention.backends import AttentionMetadata - -# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. -from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 -from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.decode_view import k3_mla_decode_view -from tensorrt_llm._torch.distributed import AllReduce from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM -from tensorrt_llm._torch.models.modeling_dspark import ( - GQADSparkForCausalLM, - draft_is_embedded_in_target, -) -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM -from tensorrt_llm._torch.models.modeling_utils import DecoderModel -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention -from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention -from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace -from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel -from tensorrt_llm._torch.modules.rms_norm import RMSNorm -from tensorrt_llm._torch.modules.situ import SituAndMul -from tensorrt_llm._torch.moe.fused_moe import ( - ConfigurableMoE, - SiTuActivation, - TRTLLMGenFusedMoE, - create_moe, -) -from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind -from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph -from tensorrt_llm._torch.pyexecutor.config_utils import is_mla -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 -from tensorrt_llm._torch.utils import AuxStreamType -from tensorrt_llm.functional import AllReduceStrategy -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - -from . import decode_comm as _decode_comm -from . import decode_gemv as _decode_gemv -from . import decode_moe as _decode_moe -from . import spec_worker as _spec_worker -from . import weights as _weights - -if TYPE_CHECKING: - from transformers import PretrainedConfig - - -# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, -# and the certification is per arch. -_SM = (10, 0) - -#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in -#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router -#: and fused-A GEMMs), then the decode kernels'. -REQUIRED_TRTLLM_OPS = ( - "attn_res_fwd", - "attn_res_rmsnorm_fwd", - "attn_res_add_rmsnorm_fwd", - "attn_res_add_rmsnorm_persistent_fwd", - "kda_prefill", - "kda_decode", - "kda_mtp_decode", - "dsv3_router_gemm_op", - "dsv3_fused_a_gemm_op", - "k3_kda_decode_attn", - "k3_kda_attn", - "k3_kda_verify", - "k3_mla_qkv", - "k3_mla_attn_vb_out", - # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). - "k3_decode_gemv", - "k3_ctm_gemv_wide", - "k3_ctm_gemv_long", - "k3_situ_mul", - "k3_head_gemv", - "k3_embed_norm", - "allgather", - # The decode path's collectives (decode_comm.py) and MoE (decode_moe.py). - "mnnvl_allreduce_attn_res", - "k3_sandwich_oproj", - "k3_sandwich_tail", - "k3_moe_front", - "k3_moe", - "k3_latent_reduce", - "k3_route_quant", - "mnnvl_allgather_split", - # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces - # with the residual add and RMSNorm (the split context projection's with hidden_norm). - "k3_ctm_gemv", - "k3_ctm_gemv_swiglu", - "k3_drafter_attn_qknorm", - "k3_sandwich_plain", - "mnnvl_fusion_allreduce", -) - -#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. -#: A renamed field upstream fails here, loudly, instead of reading as a default. -REQUIRED_ENGINE_FIELDS = { - "attn_metadata": ( - "num_contexts", - "num_generations", - "num_seqs", - "num_tokens", - "seq_lens", - "tokens_per_block", - "kv_cache_manager", - "mamba_metadata", - ), - "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), -} - -#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every -#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry -#: leaves the list when a catalog entry replaces it. -UNCERTIFIED_GENERIC_CALLS = ( - # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. - "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", - "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", - "tensorrt_llm._torch.models.modeling_utils.DecoderModel", - # The text model's stock modules. - "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", - "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", - "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", - "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", - "tensorrt_llm._torch.moe.fused_moe.create_moe", - "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", - "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", - "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", - "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", - "tensorrt_llm._torch.modules.situ.SituAndMul", - "tensorrt_llm._torch.modules.rms_norm.RMSNorm", - "tensorrt_llm._torch.distributed.AllReduce", - "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", - # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries - # do not take a block, its context projection without the TP group's collective state, its context k / v, its - # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. - "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", - "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", - "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", - "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", -) - -# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware -# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). -DECODE_MAX_TOKENS = 8 -MAX_REQUESTS = 8 -MAX_TOKENS_PER_REQUEST = 8 - -# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own -# get_model_defaults sets it for the same reason). -_TOKENS_PER_BLOCK = 64 - -_LANG_PREFIX = "language_model." - - -# ---------------------------------------------------------------------------------------------------------------------- -# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and -# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, -# `K3DecodeMLA` below). -# ---------------------------------------------------------------------------------------------------------------------- - -# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections -# instead of the min-latency fused GEMM op (read once at import). -_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( - os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" -) - - -# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table -# holds cos=1/sin=0), but the chunked-context path indexes the table by -# absolute position, so it must cover max_position_embeddings (~512MB per -# backend for the 1M-position checkpoint); a smaller table is read out of -# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context -# deployments. -_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" - - -class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): - """DeepSeek-V3 routing, computed outside the TRTLLM-Gen MoE kernel. - - The kernel's own top-16 of 896 experts is the slow part of a generic MoE call of a few tokens. Outside it, the - scheduler routes with the backend's fused route + MXFP8 quantize (``trtllm::k3_route_quant``) up to 64 tokens and - with ``noaux_tc_op`` above, and the kernel takes the expert ids and weights. Both compute the kernel's arithmetic - and break ties the same way; a bf16 weight can differ by one ulp where it lies within fp32 rounding of a bf16 - rounding boundary. From about a thousand tokens on, ``noaux_tc_op`` takes longer than the kernel's routing. - """ - - @property - def requires_separated_routing(self) -> bool: - return True - - -class KimiK3MoEGate(nn.Module): - """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" - - def __init__( - self, - config: Any, - *, - logits_gemm_dtype: torch.dtype | None = None, - device: torch.device | None = None, - ) -> None: - super().__init__() - self.config = config - self.top_k = config.num_experts_per_token - self.num_experts = config.num_experts - self.routed_scaling_factor = config.routed_scaling_factor - self.moe_router_activation_func = config.moe_router_activation_func - self.num_expert_group = getattr(config, "num_expert_group", 1) - self.topk_group = getattr(config, "topk_group", 1) - self.moe_renormalize = config.moe_renormalize - self.gating_dim = config.hidden_size - - assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( - "K3 MoE gate supports sigmoid or softmax scoring only" - ) - - # The checkpoint stores the gate weight in bf16. Storing it in bf16 - # permits the single bf16xbf16 router GEMM while retaining fp32 output. - weight_dtype = logits_gemm_dtype or torch.float32 - self.weight = nn.Parameter( - torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) - ) - self.e_score_correction_bias = nn.Parameter( - torch.empty(self.num_experts, dtype=torch.float32, device=device) - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" - hidden_2d = hidden_states.reshape(-1, self.gating_dim) - if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: - return torch.ops.trtllm.dsv3_router_gemm_op( - hidden_2d.contiguous(), - self.weight.t(), - bias=None, - out_dtype=torch.float32, - ) - return torch.nn.functional.linear( - hidden_2d.type(torch.float32), - self.weight.type(torch.float32), - None, - ) - - @property - def routing_method(self) -> KimiK3MoeRoutingMethod: - """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" - if self.moe_router_activation_func != "sigmoid": - raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") - if not self.moe_renormalize: - raise ValueError( - "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." - ) - return KimiK3MoeRoutingMethod( - top_k=self.top_k, - n_group=self.num_expert_group, - topk_group=self.topk_group, - routed_scaling_factor=self.routed_scaling_factor, - callable_e_score_correction_bias=lambda: self.e_score_correction_bias, - is_fused=True, - ) - - -class KimiK3RMSNorm(nn.Module): - """RMSNorm matching the Kimi checkpoint implementation's rounding.""" - - def __init__( - self, - hidden_size: int, - eps: float = 1e-6, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - ) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) - self.eps = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states_float = hidden_states.to(torch.float32) - variance = hidden_states_float.pow(2).mean(-1, keepdim=True) - hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) - return self.weight * hidden_states_float.to(input_dtype) - - -def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: - """Return the finite SiTu betas required by the routed-expert kernels.""" - config_situ_beta = getattr(cfg, "activation_situ_beta", None) - situ_beta = 1.0 if config_situ_beta is None else config_situ_beta - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - if situ_linear_beta is None: - raise ValueError( - "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " - "None means an identity linear branch that the fused kernels cannot represent." - ) - if situ_beta <= 0 or situ_linear_beta <= 0: - raise ValueError( - f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." - ) - return float(situ_beta), float(situ_linear_beta) - - -def _get_text_config(pretrained_config: "PretrainedConfig"): - """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" - if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( - not hasattr(pretrained_config, "linear_attn_config") - and hasattr(pretrained_config, "text_config") - ): - return pretrained_config.text_config - return pretrained_config - - -def _is_kda_layer(cfg, layer_idx: int) -> bool: - return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] - - -def _is_mla_layer(cfg, layer_idx: int) -> bool: - return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] - - -KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" - - -_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" - - -KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" - - -_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" - - -KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" - - -KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" - - -_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") - - -# Turning the port on should not require knowing which of the three topologies -# is the right one: "1" selects the measured policy (``split``), so enabling the -# feature and choosing the policy are one action through one variable. -_ATTN_RES_TOPOLOGY_ON = "1" - - -def _read_attn_res_topology() -> str: - """Default ``per_token``: the persistent kernel is opt-in. - - ``1`` is the only accepted on-value and resolves to ``split``. The named - topologies stay available for measurement: ``persistent`` uses the - persistent kernel at every shape it implements, ``per_token`` at none. - """ - raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") - if raw == _ATTN_RES_TOPOLOGY_ON: - return "split" - if raw not in _ATTN_RES_TOPOLOGIES: - # Loudly, for the same reason as the token ceiling below: a mistyped A/B - # arm that silently fell back to the default would measure one side - # twice and report no difference. - raise ValueError( - f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " - f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " - f"(which means 'split'), got {raw!r}" - ) - return raw - - -def _read_fused_attn_res_max_tokens() -> int: - """Resolved after the topology, because its default follows it. - - With the persistent port off -- the default -- the ceiling is 1, the - pre-existing gate: the fused epilogue is taken at the single-token decode - shape and nowhere else. Enabling the port raises it to 32, the top of the - measured range, so that "off" keeps meaning "unchanged". - """ - default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" - raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) - try: - value = int(raw) - except ValueError: - # Failing loudly matters more than usual here: a mistyped A/B arm that - # silently fell back to the default would measure the candidate twice - # and report no difference. - raise ValueError( - f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" - ) from None - if value < 1: - raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") - return value - - -_ATTN_RES_TOPOLOGY = _read_attn_res_topology() - - -_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() - - -def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: - """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. - - No token ceiling: the persistent grid is sized by the SM count, not by the - token count, so prefill is the case it exists for. - """ - del M # deliberately unused; see above - return H == 7168 and 2 <= N <= 9 - - -def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: - """Pick between the two fused kernels for this call site. - - ``persistent`` takes the persistent kernel at every shape it implements; - ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in - for the prefill/decode boundary. Shapes the persistent kernel does not - implement fall through to the caller's existing gate and land on the - unfused path. - """ - if not _persistent_attn_res_applicable(M, H, N): - return False - if _ATTN_RES_TOPOLOGY == "persistent": - return True - return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens - - -def _apply_attn_res_fused( - prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm -) -> Optional[torch.Tensor]: - """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. - - Returns ``None`` when the call falls outside the fused kernel's - contract (dtype/shape/arch) so the caller can use the exact fp32 reference - instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` - layout. Candidate order matches the reference: snapshots first, the - running prefix sum last. - """ - if ( - prefix_sum.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not block_residual.is_cuda - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): - return None - try: - attn_res_op = torch.ops.trtllm.attn_res_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output, _rsigma, _probs, _logits = attn_res_op( - layer_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - ) - return output.reshape(M, H) - - -def _rms_norm_eps(norm: nn.Module) -> float: - if hasattr(norm, "eps"): - return float(norm.eps) - return float(norm.variance_epsilon) - - -def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: - """Report whether the fused path was actually reached, once per shape. - - ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that - the shape gate let the call through, and a rejected call looks exactly like - a disabled one in the logs. Emitted at debug level: it is once per distinct - shape, not once per process, so it is a diagnostic rather than a summary. - """ - logger.debug_once( - f"Kimi K3 attn-res fusion [{site}]: " - f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", - key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", - ) - - -def _apply_attn_res_rmsnorm_fused( - prefix_sum: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Optional[torch.Tensor]: - """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, - default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" - if max_fused_tokens is None: - max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS - if ( - prefix_sum.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not block_residual.is_cuda - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - N = K + 1 - # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, - # which is the measured window; larger token counts have not been measured - # and fall back to the unfused add + attn_res_fwd + RMSNorm path. - if _use_persistent_attn_res(M, H, N, max_fused_tokens): - try: - persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd - except (AttributeError, RuntimeError): - return None - _, output = persistent_op( - prefix_sum.reshape(M, 1, H).contiguous(), - None, - block_residual.reshape(K, M, 1, H).contiguous(), - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - ) - _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) - return output.reshape(M, H) - - if M > max_fused_tokens or H != 7168 or N > 12: - _note_attn_res_fusion("attn_res+norm", False, M, H, N) - return None - try: - attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output = attn_res_rmsnorm_op( - layer_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - early_trigger=True, - ) - _note_attn_res_fusion("attn_res+norm", True, M, H, N) - return output.reshape(M, H) - - -def _apply_attn_res_add_rmsnorm_fused( - prefix_sum: torch.Tensor, - addend: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in - ``_apply_attn_res_rmsnorm_fused``). - - The production residual add produces a BF16 tensor that remains live across - the following MLP. The kernel therefore returns that materialized, - BF16-rounded prefix sum alongside the normalized attention-residual output, - while avoiding a separate add launch and a re-read of the intermediate by - attention-residual selection. - """ - if max_fused_tokens is None: - max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS - if ( - prefix_sum.dtype is not torch.bfloat16 - or addend.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not addend.is_cuda - or not block_residual.is_cuda - or prefix_sum.shape != addend.shape - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - N = K + 1 - # Same measured window as _apply_attn_res_rmsnorm_fused above. - if _use_persistent_attn_res(M, H, N, max_fused_tokens): - try: - persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd - except (AttributeError, RuntimeError): - return None - updated_prefix_sum, output = persistent_op( - prefix_sum.reshape(M, 1, H).contiguous(), - addend.reshape(M, 1, H).contiguous(), - block_residual.reshape(K, M, 1, H).contiguous(), - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - ) - _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) - return updated_prefix_sum.reshape(M, H), output.reshape(M, H) - - if M > max_fused_tokens or H != 7168 or N > 12: - _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) - return None - try: - attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - addend_kernel = addend.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - updated_prefix_sum, output = attn_res_add_rmsnorm_op( - layer_kernel, - addend_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - early_trigger=True, - ) - _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) - return updated_prefix_sum.reshape(M, H), output.reshape(M, H) - - -def _apply_attn_res( - prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm -) -> torch.Tensor: - """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). - - prefix_sum: ``[num_tokens, hidden_size]`` - block_residual: ``[num_snapshots, num_tokens, hidden_size]`` - - Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's - contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. - Only the fallback boundary restores the HF ``[M, K, H]`` layout. - """ - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) - if fused is not None: - return fused - block_residual_hf = block_residual.transpose(0, 1) - v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) - v_float = v.float() - variance = v_float.pow(2).mean(-1, keepdim=True) - k = v_float * torch.rsqrt(variance + norm.eps) - score_weight = norm.weight.float() * proj.weight.squeeze(0).float() - scores = (k * score_weight).sum(-1) - probs = scores.softmax(-1).unsqueeze(1) - hidden_states = torch.matmul(probs, v_float).squeeze(1) - return hidden_states.to(v.dtype) - - -def _apply_attn_res_and_rmsnorm( - prefix_sum: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> torch.Tensor: - """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the - fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_rmsnorm_fused( - prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens - ) - if fused is not None: - return fused - return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) - - -def _apply_attn_res_add_and_rmsnorm( - prefix_sum: torch.Tensor, - addend: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in - ``_apply_attn_res_and_rmsnorm``).""" - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_add_rmsnorm_fused( - prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens - ) - if fused is not None: - return fused - updated_prefix_sum = prefix_sum + addend - return updated_prefix_sum, _apply_attn_res_and_rmsnorm( - updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens - ) - - -# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where -# they are faster than the add -> attn_res -> RMSNorm chain. -_WIDE_ATTN_RES_MAX_TOKENS = 32 - - -def _attn_res_max_tokens(step: DecodeStep) -> int: - """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step - ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" - return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS - - -_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") - - -# The routed experts' module path prefixes. ``exclude_modules`` matches with -# wildcards and walks ancestor prefixes, so an empty prefix would widen what -# matches instead of just missing. -_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") - - -# Routed-expert quantization used when the checkpoint declares nothing per -# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors -# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that -# checkpoint is what this default has always served. -_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 - - -def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: - backend.quant_method.load_packed_mxfp4_expert( - backend, - global_expert_id=expert_idx, - local_slot_id=local_slot_id, - w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), - w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), - w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), - w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), - w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), - w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), - ) - - -class _K3ExpertCkptSpec(NamedTuple): - """How one routed-expert quantization is spelled and loaded.""" - - # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. - kinds: Tuple[str, ...] - loader: Callable[..., None] - # Set of filled slots the loader maintains, checked after the load. - loaded_slots_attr: str - # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). - needs_layer_finalize: bool - - -_K3_EXPERT_CKPT_SPECS = { - QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( - kinds=("weight_packed", "weight_scale"), - loader=_load_packed_mxfp4_expert, - loaded_slots_attr="_packed_mxfp4_loaded_slots", - needs_layer_finalize=False, - ), -} - - -def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: - spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) - if spec is None: - raise NotImplementedError( - f"Kimi K3 routed experts are quantized as {quant_algo}, for which " - "no per-expert checkpoint layout is known. Supported: " - f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." - ) - return spec - - -class KimiK3MoERuntime(nn.Module): - """Kimi K3 latent MoE block backed by ConfigurableMoE.""" - - def __init__( - self, - model_config: ModelConfig, - cfg, - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ): - """Build the routed experts and the shared expert for one MoE layer. - - ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: - the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields - that ``ModelConfig`` does not carry. - - ``aux_stream_dict`` is shared across every layer of the model, so the - streams reached through it are borrowed and must not be synchronized - or reassigned here. - """ - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = cfg.hidden_size - self.num_experts = cfg.num_experts - self.top_k = cfg.num_experts_per_token - self.moe_hidden_size = cfg.routed_expert_hidden_size - # ValueError (not assert): these guard unsupported checkpoint - # configurations and must stay active under ``python -O``. - if self.moe_hidden_size is None: - raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") - if not getattr(cfg, "latent_moe_use_norm", False): - raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") - - situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) - self._situ_betas = (situ_beta, situ_linear_beta) - dtype = torch.bfloat16 - - # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 - # accumulate/output (checkpoint stores the gate weight in bf16; saves a - # per-layer input cast + fp32 splitK pair on the bs1 decode path). - # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. - _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") - _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True - self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) - - routed_moe_model_config = self._routed_moe_model_config(model_config) - routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) - # Resolved here so ``load_weights`` reads the checkpoint layout off the - # module instead of re-deriving it at each of its three call sites. - self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) - routed_moe_kwargs = dict( - routing_method=self.gate.routing_method, - num_experts=self.num_experts, - hidden_size=self.moe_hidden_size, - intermediate_size=cfg.moe_intermediate_size, - dtype=dtype, - # Kimi owns the latent reduction so it can order that collective - # after the shared expert's auxiliary-stream reduction. - reduce_results=False, - model_config=routed_moe_model_config, - override_quant_config=routed_quant_config, - layer_idx=layer_idx, - aux_stream_dict=aux_stream_dict, - # Let CommunicationFactory select the best available strategy. - communication_method=None, - activation=SiTuActivation( - gate_softcap=situ_beta, - linear_softcap=situ_linear_beta, - ), - # A request that silently degraded to CUTLASS would be benchmarked - # as if it were the backend that was asked for, and the decline is - # easy to trigger: MegaMoE has its own token / top-k limits and is - # EP-only, and CuteDSL declines on activation shape, SM version and - # the CuTe DSL dependency. As measured once: a CUTEDSL request - # was turned down on every one of the 92 MoE layers, on all 16 - # ranks, and still produced correct text and a zero exit -- the - # only trace was a warning line per layer. Fail in the resolver - # instead, which reports the rejection trail. - # - # CUTLASS is absent on purpose: it is the fallback target, so - # "degraded to CUTLASS" is not a thing that can happen to it. - allow_backend_degradation=routed_moe_model_config.moe_backend - not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), - ) - self._check_trtllm_situ_quant( - routed_moe_model_config.moe_backend, routed_quant_config.quant_algo - ) - - self.routed_experts = create_moe(**routed_moe_kwargs) - if not isinstance(self.routed_experts, ConfigurableMoE): - raise RuntimeError( - "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." - ) - if self.routed_experts.layer_load_balancer is not None: - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming does not yet support " - "dynamic EPLB or replicated expert slots." - ) - local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) - if local_expert_ids != list( - range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) - ): - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming currently requires a " - "contiguous static expert partition." - ) - self.local_expert_ids = tuple(local_expert_ids) - self.experts_per_rank = len(local_expert_ids) - self.expert_lo = local_expert_ids[0] - self.expert_hi = self.expert_lo + self.experts_per_rank - - shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts - shared_model_config = copy.copy(model_config) - shared_model_config.quant_config = QuantConfig() - # Direct MoE-TP leaves both branches as partials for one concatenated - # all-reduce. - use_shared_tp = model_config.mapping.tp_size > 1 - self._reduce_routed_output = ( - use_shared_tp - and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM - ) - if self._reduce_routed_output and self.routed_experts.all_reduce is None: - raise RuntimeError( - "Kimi K3 direct MoE tensor parallelism requires the " - "ConfigurableMoE all-reduce even when reduce_results=False." - ) - self.shared_experts = GatedMLP( - hidden_size=cfg.hidden_size, - intermediate_size=shared_intermediate, - bias=False, - activation=SituAndMul( - beta=situ_beta, - linear_beta=situ_linear_beta, - use_fused_activation=True, - ), - dtype=dtype, - config=shared_model_config, - reduce_output=use_shared_tp, - layer_idx=layer_idx, - is_shared_expert=True, - ) - # Side stream (+ fork/join events) for overlapping shared-expert - # compute with the routed chain. Only engaged when multi-stream is - # active (CUDA graphs on); otherwise both run in order on the default - # stream. - self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] - self.moe_main_event = torch.cuda.Event() - self.moe_shared_event = torch.cuda.Event() - self.routed_expert_down_proj = nn.Linear( - cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype - ) - self.routed_expert_up_proj = nn.Linear( - self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype - ) - # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer - # fallback is the same fp32-variance eager math as KimiK3RMSNorm). - self.routed_expert_norm = RMSNorm( - hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. - self.decode_moe: Optional[_decode_moe.K3DecodeMoeLayer] = None - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - @staticmethod - def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: - if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): - return projection(hidden_states) - return torch.ops.trtllm.dsv3_fused_a_gemm_op( - hidden_states, projection.weight.t(), None, None - ) - - @staticmethod - def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: - """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit - ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (4 x 4, asserted at - construction).""" - return mapping.moe_tp_size, mapping.moe_ep_size - - @staticmethod - def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: - """Routed-expert quantization for ``layer_idx``: the MXFP4 checkpoint declares nothing per layer (asserted at - construction), so the experts keep the ``W4A8_MXFP4_MXFP8`` default. - - An exclusion outranks the default: ``create_weights`` treats an override - as authoritative over anything ``__post_init__`` wrote, so this return - value stands in for both quantization passes and exclusion is the one - that runs second. It is matched as a pattern, so it is asked only about - real module names. - """ - quant_config = model_config.quant_config - if quant_config is not None and any( - quant_config.is_module_excluded_from_quantization( - f"{prefix}layers.{layer_idx}.{suffix}" - ) - for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES - for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES - ): - logger.debug( - "Kimi K3 layer %d routed experts: excluded from quantization, " - "keeping them unquantized", - layer_idx, - ) - return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) - - logger.debug( - "Kimi K3 layer %d routed experts: no per-layer quant config in the " - "checkpoint, defaulting to %s", - layer_idx, - _K3_DEFAULT_ROUTED_QUANT_ALGO, - ) - return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) - - @staticmethod - def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: - """Reject a routed-expert format trtllm-gen has no fused SiTu cubin for. - - trtllm-gen has fused SiTu FC1 cubins for two input formats and no - standalone SiTu activation kernel, so anything else has to die here - rather than in a cubin lookup deep inside the runner. Checked against - the resolved backend, not the K3 architecture branch, because the - generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also - land on TRTLLM. - - The admitted set is read off the backend rather than restated here, - because restating it is what broke. This guard was written in #17865 - when MXFP4 was the only fused SiTu drop; #17940 then added the NVFP4 - (group-16 ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``) cubins and updated - ``TRTLLMGenFusedMoE``'s set without touching this copy. For the week - in between, an NVFP4 K3 checkpoint could not start at all -- and not - only when TRTLLM was asked for by name, because - ``ModelConfig.resolve_moe_backend`` sends every K3 architecture to - TRTLLM, so the default AUTO configuration hit this raise too. The unit - tests did not catch it: they call ``create_moe`` directly and never - reach this guard, so the kernel path stayed green while the model path - was closed. - - A staticmethod, not an inline block, so that the invariant is - reachable from a test without constructing the whole runtime. - """ - situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() - if moe_backend != "TRTLLM" or quant_algo in situ_supported: - return - supported = ", ".join(sorted(algo.name for algo in situ_supported)) - raise ValueError( - f"Kimi K3 routed experts are quantized as {quant_algo}, which the " - "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " - f"exist only for {supported}. Set moe_config.backend to CUTLASS " - "or MEGAMOE_CUTEDSL." - ) - - @staticmethod - def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: - """Build a private routed-expert mapping without mutating the shared - config, with the split of ``_select_moe_tp_ep``.""" - # Every backend here declares ``ActivationType.SiTu`` in its - # ``activation_support``; the list is not a preference order. CUTEDSL - # joined once its act-fusion kernel grew the SiTU epilogue. - supported_backends = { - "CUTLASS", - "TRTLLM", - "CUTEDSL", - "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - } - if model_config.moe_backend not in supported_backends: - raise ValueError( - "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " - "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " - f"got {model_config.moe_backend!r}." - ) - if model_config.moe_load_balancer is not None: - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming does not yet support " - "EPLB or replicated expert slots." - ) - mapping = model_config.mapping - - moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) - logger.info_once( - f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " - f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", - key="kimi_k3_moe_tp_ep_split", - ) - - mapping_dict = mapping.to_dict() - mapping_dict["moe_cluster_size"] = 1 - mapping_dict["moe_tp_size"] = moe_tp - mapping_dict["moe_ep_size"] = moe_ep - routed_mapping = Mapping.from_dict(mapping_dict) - - routed_model_config = copy.copy(model_config) - routed_model_config._frozen = False - routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) - routed_model_config.mapping = routed_mapping - routed_model_config.moe_backend = model_config.moe_backend - # MegaMoE uses this value as global DP SymmBuffer capacity, then - # divides it by EP size for the per-rank allocation. Other backends - # keep the user-configured value as their MoE chunking bound. - # Preserve an explicitly larger capacity. - if routed_model_config.moe_backend in { - "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - }: - default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size - configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) - if configured_moe_max_num_tokens < default_moe_max_num_tokens: - logger.info_once( - "Kimi K3 MegaMoE raises moe_max_num_tokens from " - f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " - "because the global DP SymmBuffer requires capacity for " - "max_num_tokens * dp_size.", - key=( - "kimi_k3_megamoe_capacity_override_" - f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" - ), - ) - routed_model_config.moe_max_num_tokens = max( - configured_moe_max_num_tokens, - default_moe_max_num_tokens, - ) - routed_model_config._frozen = True - return routed_model_config - - def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: - """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel - share (``partial_tail``): the decode path takes the step with that tail.""" - return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) - - def forward( - self, - hidden_states: torch.Tensor, - all_rank_num_tokens=None, - partial_tail: bool = False, - *, - step: DecodeStep, - ) -> Union[torch.Tensor, _decode_comm.PendingTail]: - """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification - (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where - ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at - most 8 tokens, a tensor on a wide decode step.""" - decode = self.decode_moe - if decode is not None and decode.takes(hidden_states, step, partial_tail): - return decode.forward( - self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push - ) - if partial_tail: - raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") - return self.forward_generic(hidden_states, all_rank_num_tokens) - - def forward_generic( - self, hidden_states: torch.Tensor, all_rank_num_tokens=None - ) -> torch.Tensor: - identity = hidden_states - router_logits = self.gate.compute_logits(hidden_states) - moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None +from tensorrt_llm.quantization.mode import QuantAlgo - def _routed_output(): - # Latent down/up projections via the min-latency fused GEMM op: - # at <=16 tokens (decode graphs) it runs a single pipelined - # bf16 kernel per projection instead of cuBLAS's split-K GEMV + - # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); - # for larger token counts the op falls back to cuBLAS internally. - # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear - # (A/B escape hatch). When the FP8 weight-read conversion has - # replaced the projection module, call it directly: its weight is - # an e4m3 buffer the bf16 dsv3 op must not read, and its forward - # is already a single fused GEMM (fp8_swap_ab_gemm). - routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) - y = self.routed_experts( - routed_in, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - ) - if self._reduce_routed_output: - return y - # Communication-backed paths return a complete routed result. - y = self.routed_expert_norm(y) - return self._routed_projection(y, self.routed_expert_up_proj) +from ..modeling import K3TargetSpec, ModelingV2KimiK3 +from ..modeling import _check_construction as _check_target_construction - # Shared experts depend only on the block input, so overlap their GEMMs - # with the routed dispatch/expert/combine chain. Multi-stream engages - # only under CUDA graphs; otherwise both branches run in order on the - # default stream. The shared GatedMLP includes its output all-reduce on - # the auxiliary stream. The join below must precede the routed - # all-reduce: concurrent collectives on different streams can corrupt - # SYMM_MEM all-reduce state. - routed_out, shared_out = maybe_execute_in_parallel( - _routed_output, - lambda: self.shared_experts(identity), - self.moe_main_event, - self.moe_shared_event, - self.shared_expert_stream, - disable_on_compile=True, - ) - if self._reduce_routed_output: - routed_latent = moe_all_reduce(routed_out) - routed_latent = self.routed_expert_norm(routed_latent) - routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) - return routed_out + shared_out - - -def resolve_attention_quant_config( - config: ModelConfig | None, layer_idx: int, projection: str -) -> QuantConfig: - """Resolve a checkpoint projection, including mixed-precision exclusions.""" - if config is None: - return QuantConfig() - global_config = config.quant_config or QuantConfig() - names = [ - f"{prefix}layers.{layer_idx}.self_attn.{projection}" - for prefix in ("language_model.model.", "model.", "") - ] - if any(global_config.is_module_excluded_from_quantization(name) for name in names): - return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) - declarations = config.quant_config_dict or {} - matches = [declarations[name] for name in names if name in declarations] - if matches: - selected = matches[0] - if any(match.quant_algo != selected.quant_algo for match in matches[1:]): - raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") - elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: - selected = QuantConfig() - else: - selected = global_config - if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): - raise ValueError( - f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " - f"quantization {selected.quant_algo}" - ) - if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): - raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") - result = copy.copy(selected) - result.kv_cache_quant_algo = global_config.kv_cache_quant_algo - return result - - -class KimiMLARuntime(nn.Module): - """Wraps K3 MLA and applies its external TP output reduction.""" - - def __init__( - self, - cfg: "PretrainedConfig", - layer_idx: int, - model_config: ModelConfig, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ) -> None: - super().__init__() - - max_positions = int( - os.environ.get( - _KIMI_K3_MLA_MAX_POSITIONS_ENV, - cfg.max_position_embeddings, - ) - ) - self.layer_idx = layer_idx - # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the - # final output reduction in this wrapper so the output gate remains - # between attention and the row-parallel o_proj. - mapping = model_config.mapping - reduce_output = mapping.tp_size > 1 - self._o_allreduce = ( - AllReduce( - mapping=mapping, - strategy=model_config.allreduce_strategy, - dtype=torch.bfloat16, - ) - if reduce_output - else None - ) - attention_config = copy.copy(model_config) - attention_config._frozen = False - attention_config.quant_config_dict = { - name: resolve_attention_quant_config(model_config, layer_idx, name) - for name in ( - "q_a_proj", - "kv_a_proj_with_mqa", - "q_b_proj", - "kv_b_proj", - "g_proj", - "o_proj", - ) - } - attention_config.quant_config = QuantConfig( - kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo - if model_config.quant_config is not None - else None - ) - attention_config._frozen = model_config._frozen - self.mixer = K3DecodeMLA( - hidden_size=cfg.hidden_size, - num_heads=cfg.num_attention_heads, - q_lora_rank=cfg.q_lora_rank, - kv_lora_rank=cfg.kv_lora_rank, - qk_nope_head_dim=cfg.qk_nope_head_dim, - qk_rope_head_dim=cfg.qk_rope_head_dim, - v_head_dim=cfg.v_head_dim, - rms_norm_eps=cfg.rms_norm_eps, - dtype=torch.bfloat16, - layer_idx=layer_idx, - use_output_gate=cfg.mla_use_output_gate, - max_position_embeddings=max_positions, - model_config=attention_config, - aux_stream_dict=aux_stream_dict, - ) - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" - return self.mixer.will_run_decode_branch(attn_metadata, step) - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - reduce_output: bool = True, - project_output: bool = True, - ) -> torch.Tensor: - """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the - mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" - # MLA.forward takes position_ids first; K3 is NoPE, so pass None. - out = self.mixer( - None, hidden_states, attn_metadata, step=step, project_output=project_output - ) - return self._reduce(out) if project_output and reduce_output else out - - def forward_generic( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - reduce_output: bool = True, - ) -> torch.Tensor: - out = self.mixer.forward_generic(None, hidden_states, attn_metadata) - return self._reduce(out) if reduce_output else out - - def _reduce(self, out: torch.Tensor) -> torch.Tensor: - # Head-sharded TP: sum the row-sharded o_proj partials across - # the head-shard group. - return self._o_allreduce(out) if self._o_allreduce is not None else out - - -class KimiLinearDecoderLayer(nn.Module): - def __init__( - self, - model_config: ModelConfig, - cfg, - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ): - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = cfg.hidden_size - dtype = torch.bfloat16 - - self.is_kda = _is_kda_layer(cfg, layer_idx) - is_mla = _is_mla_layer(cfg, layer_idx) - if self.is_kda == is_mla: - raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") - - if self.is_kda: - projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") - attention_config = copy.copy(model_config) - attention_config._frozen = False - attention_config.quant_config_dict = { - name: resolve_attention_quant_config(model_config, layer_idx, name) - for name in projection_names - } - attention_config._frozen = model_config._frozen - self.linear_attn = K3DecodeKDA( - cfg, - layer_idx, - mapping=model_config.mapping, - allreduce_strategy=model_config.allreduce_strategy, - aux_stream=aux_stream_dict[AuxStreamType.Attention], - model_config=attention_config, - ) - else: - self.self_attn = KimiMLARuntime( - cfg, - layer_idx, - model_config=model_config, - aux_stream_dict=aux_stream_dict, - ) - - self.is_moe = ( - cfg.num_experts is not None - and layer_idx >= cfg.first_k_dense_replace - and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 - ) - if self.is_moe: - self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) - else: - situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) - # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's - # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take - # (decode_gemv.SITES); otherwise it stays within one node. - spans_nodes = self._mnnvl_allreduce() is not None - if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: - self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) - mlp_model_config = copy.copy(model_config) - mlp_model_config.quant_config = QuantConfig() - # K3's dense layer is BF16, so a unit block size gives the same - # subgroup selection as DeepSeek-V3. - self.mlp = GatedMLP( - hidden_size=cfg.hidden_size, - intermediate_size=cfg.intermediate_size, - bias=False, - activation=SituAndMul( - beta=situ_beta, - linear_beta=situ_linear_beta, - use_fused_activation=True, - ), - dtype=dtype, - config=mlp_model_config, - overridden_tp_size=self.mlp_tp_size, - reduce_output=self.mlp_tp_size > 1, - layer_idx=layer_idx, - ) - self._situ = (situ_beta, situ_linear_beta) - # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics - # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied - # after downcast, use_gemma=False). - self.input_layernorm = RMSNorm( - hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.post_attention_layernorm = RMSNorm( - hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - - # Attention residual scheme (always on for K3). The res norms stay - # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by - # _apply_attn_res and the fused attn_res op, never called as - # modules. - self.attn_res_block_size = cfg.attn_res_block_size - assert self.attn_res_block_size is not None, ( - "Kimi K3 runtime expects attn_res_block_size to be set" - ) - self.self_attention_res_norm = KimiK3RMSNorm( - cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) - self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set - # by the target's post_load_weights. - self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None - self.sandwich_oproj = False - - def forward( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - attn_metadata: AttentionMetadata, - step: DecodeStep, - capture: Optional[Tuple[Any, int]] = None, - prenormed: bool = False, - pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, - defer_moe_tail: bool = False, - ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: - """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). - - ``block_residual`` is a preallocated snapshot bank in kernel-native - ``[K_max, M, H]`` layout. Returns the running prefix sum and the - number of valid bank rows. - - ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux - stream for the layer BEFORE this one: the aggregated stream for layer j - is by definition what its next consumer sees, so the mixture computed - below already is it. Reading it here beats recomputing it, and is only - possible because K3 asserts pp_size == 1 -- layer j+1 is always local. - PP support would need a recompute at the rank boundary. - - ``step`` is the step's classification (``decode_step``), handed to the - attention module. - - ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is - this layer's input norm, and the layer's input, the step's embedding, - already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). - - The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens - runs the attention's all-reduce and the residual update in one - collective once the target has built its ``decode_comm``: with o_proj - (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, - the step and the attention's decode branch, else on the attention's - unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). - - ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the - previous layer's MoE output, unreduced (``defer_moe_tail``); then - ``hidden_states`` is the prefix sum without it, and this layer's - pre-attention step reduces and adds it: a ``PendingTail`` in one - ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its - all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer - the ``sandwich_tail`` call also stores the tap into the capture slot - where the speculative metadata exposes it (``capture_view``); other - tapped layers take the tap after the step. - - ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` - instead, ``partial`` this layer's MoE output unreduced, for the next - consumer to reduce. - """ - prefix_sum = hidden_states - valid_block_residual = block_residual[:num_snapshots] - attn_res_max_tokens = _attn_res_max_tokens(step) - tail = ( - pending_moe_partial - if isinstance(pending_moe_partial, _decode_comm.PendingTail) - else None - ) - # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum - # straight into the bank row it snapshots. - snapshot_row = None - if tail is not None and self.layer_idx % self.attn_res_block_size == 0: - snapshot_row = block_residual[num_snapshots] - # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the - # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). - tap_view = None - if tail is not None and capture is not None: - tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) - - if prenormed: - assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 - assert pending_moe_partial is None and capture is None - elif tail is not None: - hidden_states, prefix_sum = self.decode_comm.sandwich_tail( - tail, - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - updated_out=snapshot_row, - tap=tap_view, - tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, - ) - elif pending_moe_partial is not None: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - attn_res_max_tokens, - ) - elif capture is not None: - # The mixture tap needs the PRE-norm value, which the fused - # attn-res + RMSNorm kernel does not expose. Keep the two steps - # split on captured layers only and fuse everywhere else. - if num_snapshots > 0: - hidden_states = _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers - # acceptance, silently. hidden_states is the pre-norm attn_res mixture; - # prefix_only wants the running prefix, already in hand as prefix_sum. - tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - hidden_states = self.input_layernorm(hidden_states) - elif num_snapshots > 0: - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - attn_res_max_tokens, - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - if capture is not None and pending_moe_partial is not None and tap_view is None: - # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output - # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. - tapped = ( - _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - if _AUX_ATTN_RES_STREAM_ENABLED - else prefix_sum - ) - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - - if self.layer_idx % self.attn_res_block_size == 0: - if not prenormed and snapshot_row is None: - block_residual[num_snapshots].copy_(prefix_sum) - num_snapshots += 1 - valid_block_residual = block_residual[:num_snapshots] - prefix_sum = None - attention = self.linear_attn if self.is_kda else self.self_attn - comm = self.decode_comm - if comm is not None and comm.takes_post_attention(hidden_states, step): - if ( - self.sandwich_oproj - and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS - and attention.will_run_decode_branch(attn_metadata, step) - ): - core = attention(hidden_states, attn_metadata, step=step, project_output=False) - hidden_states, prefix_sum = comm.sandwich_oproj( - core, - self._o_proj().weight, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) - hidden_states, prefix_sum = comm.allreduce_attn_res( - partial, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - if comm is not None and step.wide: - partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) - attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) - else: - attention_output = attention(hidden_states, attn_metadata, step=step) - if prefix_sum is None: - prefix_sum = attention_output - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - attn_res_max_tokens, - ) - else: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - attention_output, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - attn_res_max_tokens, - ) - all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) - if self.is_moe and defer_moe_tail: - partial = self.block_sparse_moe( - hidden_states, all_rank_num_tokens, partial_tail=True, step=step - ) - return prefix_sum, num_snapshots, partial - if self.is_moe: - hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) - else: - hidden_states = self._dense_mlp(hidden_states, step) - - prefix_sum = prefix_sum + hidden_states - return prefix_sum, num_snapshots - - def forward_generic( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - attn_metadata: AttentionMetadata, - capture: Optional[Tuple[Any, int]] = None, - ) -> Tuple[torch.Tensor, int]: - prefix_sum = hidden_states - valid_block_residual = block_residual[:num_snapshots] - if capture is not None: - if num_snapshots > 0: - hidden_states = _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - hidden_states = self.input_layernorm(hidden_states) - elif num_snapshots > 0: - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - if self.layer_idx % self.attn_res_block_size == 0: - block_residual[num_snapshots].copy_(prefix_sum) - num_snapshots += 1 - valid_block_residual = block_residual[:num_snapshots] - prefix_sum = None - attention = self.linear_attn if self.is_kda else self.self_attn - comm = self.decode_comm - if comm is not None and comm.takes_post_attention(hidden_states, None): - partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) - hidden_states, prefix_sum = comm.allreduce_attn_res( - partial, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - elif prefix_sum is None: - prefix_sum = attention.forward_generic(hidden_states, attn_metadata) - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - attention.forward_generic(hidden_states, attn_metadata), - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - if self.is_moe: - all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) - hidden_states = self.block_sparse_moe.forward_generic( - hidden_states, all_rank_num_tokens - ) - else: - hidden_states = self.mlp(hidden_states) - return prefix_sum + hidden_states, num_snapshots - - def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: - """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode - kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" - gemvs = self.decode_gemvs - if gemvs is not None and step.small: - out = gemvs.dense_mlp( - hidden_states, - self.mlp.gate_up_proj.weight, - self.mlp.down_proj.weight, - *self._situ, - ) - if out is not None: - return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out - return self.mlp(hidden_states) - - def _mnnvl_allreduce(self): - """The MNNVL all-reduce of this layer's attention output, or None.""" - attention = self.linear_attn if self.is_kda else self.self_attn - return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) - - def _o_proj(self) -> nn.Module: - """This layer's attention output projection (row parallel).""" - return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj - - def _o_allreduce(self) -> AllReduce: - """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" - return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce - - def accepts_moe_partial(self, num_snapshots: int) -> bool: - """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target - built its collective state, and the snapshot bank is not empty.""" - return num_snapshots > 0 and self.decode_comm is not None - - def skip_forward( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. - - ``modeling_utils.skip_forward()`` only drops a module's weights when it - finds this attribute, so without it the layer-wise benchmarks would - allocate all 93 layers instead of the profiled slice. - """ - return hidden_states, block_residual - - -class KimiLinearModel(DecoderModel): - def __init__(self, model_config: ModelConfig): - super().__init__(model_config) - cfg = _get_text_config(model_config.pretrained_config) - self._text_cfg = cfg - dtype = torch.bfloat16 - - # Attention and MoE phases are sequential, so their branch-overlap - # roles share one stream; MoE-internal overlap roles remain separate. - aux_stream_list = [torch.cuda.Stream() for _ in range(4)] - self.aux_stream_dict = { - AuxStreamType.Attention: aux_stream_list[0], - AuxStreamType.MoeShared: aux_stream_list[0], - AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], - AuxStreamType.MoeBalancer: aux_stream_list[2], - AuxStreamType.MoeOutputMemset: aux_stream_list[3], - } - - self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) - self.layers = nn.ModuleList( - [ - KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) - for layer_idx in range(cfg.num_hidden_layers) - ] - ) - self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) - - # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) - # by _apply_attn_res and the fused attn_res op. - self.output_attn_res_norm = KimiK3RMSNorm( - cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - self.num_attn_res_snapshots = ( - cfg.num_hidden_layers + cfg.attn_res_block_size - 1 - ) // cfg.attn_res_block_size - # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the - # weights are final. None: every step embeds on the generic path. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - # Which convention the drafter tap is on is not recoverable from the - # served output -- a mismatch only lowers acceptance -- so state it once - # at construction rather than leaving it to be inferred from an AL. - logger.info_once( - "Kimi K3 aux hidden capture: mode=" - f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " - f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", - key="kimi_k3_aux_capture_mode", - ) - - @property - def kda_draft_records(self) -> bool: - """Whether the hybrid cache manager keeps the last verify's draft records (``kda_state_tok``), from which - ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn`` replay the accepted drafts. The engine reads it once the weights - are loaded, to build the manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer - taking the K3 kernels. Otherwise the KDA verify replays the accepted drafts from the replay caches (the - built-in verify).""" - spec_config = getattr(self.model_config, "spec_config", None) - return bool( - spec_config is not None - and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) - and spec_config.tokens_per_gen_step in (2, 4, 6, 8) - and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) - ) - - def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, - from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from - load-time state, so every rank and every replay of the graph decides alike.""" - if not torch.cuda.is_current_stream_capturing(): - return False - kda = [layer.linear_attn for layer in self.layers if layer.is_kda] - mla = [layer.self_attn for layer in self.layers if not layer.is_kda] - return latent_push( - step, - capturing=True, - breakable=is_in_breakable_cuda_graph(), - kda_draft_records=self.kda_draft_records, - kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), - mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), - ) - - def _defer_moe_tail( - self, - i: int, - hidden_states: torch.Tensor, - num_snapshots: int, - spec_metadata, - capture_set, - step: DecodeStep, - ) -> bool: - """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE - takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after - the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps - its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated - tail.""" - layer = self.layers[i] - if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): - return False - if spec_metadata is not None and ( - capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) - ): - return False - return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - spec_metadata=None, - **kwargs, - ) -> torch.Tensor: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] - step = decode_step(attn_metadata, num_tokens) - if step is None: - return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) - if self._latent_push(attn_metadata, step): - step = replace(step, latent_push=True) - # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. - prenormed = None - if ( - inputs_embeds is None - and self.decode_gemvs is not None - and len(self.layers) > 0 - and self.num_attn_res_snapshots > 0 - ): - table = self.embed_tokens.weight - block_residual = table.new_empty( - self.num_attn_res_snapshots, num_tokens, table.shape[1] - ) - prenormed = self.decode_gemvs.embed_norm( - input_ids, table, self.layers[0].input_layernorm, block_residual - ) - if prenormed is not None: - hidden_states = prenormed - else: - hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - block_residual = hidden_states.new_empty( - self.num_attn_res_snapshots, - hidden_states.shape[0], - hidden_states.shape[1], - ) - num_snapshots = 0 - capture_set = getattr(spec_metadata, "_capture_layer_set", None) - # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) - # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. - pending_moe_partial = None - for i, layer in enumerate(self.layers): - defer_moe_tail = self._defer_moe_tail( - i, hidden_states, num_snapshots, spec_metadata, capture_set, step - ) - outputs = layer( - hidden_states, - block_residual, - num_snapshots, - attn_metadata, - step=step, - capture=self._capture(spec_metadata, capture_set, i), - prenormed=i == 0 and prenormed is not None, - pending_moe_partial=pending_moe_partial, - defer_moe_tail=defer_moe_tail, - ) - if defer_moe_tail: - hidden_states, num_snapshots, pending_moe_partial = outputs - else: - (hidden_states, num_snapshots), pending_moe_partial = outputs, None - - if isinstance(pending_moe_partial, _decode_comm.PendingTail): - normed, _ = self.layers[-1].decode_comm.sandwich_tail( - pending_moe_partial, - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - ) - return normed - if pending_moe_partial is not None: - _, normed = _apply_attn_res_add_and_rmsnorm( - hidden_states, - _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - _attn_res_max_tokens(step), - ) - return normed - return self._final_norm( - hidden_states, - block_residual, - num_snapshots, - spec_metadata, - capture_set, - _attn_res_max_tokens(step), - ) - - def _generic_forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor], - inputs_embeds: Optional[torch.FloatTensor], - spec_metadata, - ) -> torch.Tensor: - hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - block_residual = hidden_states.new_empty( - self.num_attn_res_snapshots, - hidden_states.shape[0], - hidden_states.shape[1], - ) - num_snapshots = 0 - capture_set = getattr(spec_metadata, "_capture_layer_set", None) - for i, layer in enumerate(self.layers): - hidden_states, num_snapshots = layer.forward_generic( - hidden_states, - block_residual, - num_snapshots, - attn_metadata, - capture=self._capture(spec_metadata, capture_set, i), - ) - return self._final_norm( - hidden_states, block_residual, num_snapshots, spec_metadata, capture_set - ) - - def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: - # DFlash/DSpark hidden-state capture. The drafter is distilled on - # the aggregated stream value -- the pre-norm softmax mixture its - # next consumer sees -- not on the raw prefix sum a layer returns, - # which is SGLang's fallback for models without the - # attention-residual scheme. Capturing the prefix sum costs 4.5pt - # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). - # The tap fires inside layer i+1, which computes that tensor - # anyway; see its forward docstring. Ground truth: SGLang - # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 - # aggregate_stream_torch. - if ( - spec_metadata is not None - and i > 0 - and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) - ): - return (spec_metadata, self.layers[i - 1].layer_idx) - return None - - def _final_norm( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - spec_metadata, - capture_set, - max_fused_tokens: Optional[int] = None, - ) -> torch.Tensor: - # The last layer has no successor, so this one recompute is - # unavoidable -- output-side score weights, matching SGLang's - # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set - # against 93 layers; kept so a set that does include the final layer - # gets the right tensor rather than the raw prefix sum. - if spec_metadata is not None and len(self.layers) > 0: - last = self.layers[-1] - if capture_set is None or last.layer_idx in capture_set: - tail = ( - _apply_attn_res( - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - ) - if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED - else hidden_states - ) - spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) - - return _apply_attn_res_and_rmsnorm( - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - max_fused_tokens, - ) - - -# ---------------------------------------------------------------------------------------------------------------------- -# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. -# ---------------------------------------------------------------------------------------------------------------------- - - -class K3DecodeKDA(KimiKDALinearAttention): - """Kimi K3's KDA attention: the built-in module, with the decode kernels on the steps they take. - - * A decode step of one token per request runs the fused input projection and the plain decode in one - ``ssm/k3_kda_decode_attn`` launch. - * With the cache manager's per-token states (``KimiLinearModel.kda_draft_records``), every verify of the layer, on - any step, runs the kernels that keep them: ``ssm/k3_kda_attn`` for one request of 8 tokens (the projection fused - in), else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it takes - them). The built-in verify replays drafts from caches these kernels do not fill, so the two never run on one - manager. - * On every step ``decode_step`` classifies, ``o_proj`` runs on the ``o_proj`` decode GEMV site where it takes the - rows, then the module's all-reduce. - - Every other step runs the built-in module. The kernels read one ``[q | k | v | g | f_a | b]`` weight, built at the - checkpoint load from the module's own, and the device's ``K3KdaBuffers`` and decode GEMVs' state, which the - target sets in ``post_load_weights``. - """ - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - # The fused [q | k | v | g | f_a | b | pad] projection weight; the six projections' weights and the built-in - # fused [q | k | v | g] and [f_a | b | pad] ones are views of it. - self.k3_proj_weight: Optional[torch.Tensor] = None - # The fused projection's Lamport buffers: one set per device, shared by every KDA layer. - self.k3_buffers: Optional[K3KdaBuffers] = None - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - @property - def takes_k3_kernels(self) -> bool: - """Whether the decode kernels can run this layer: its fused projection weight is built.""" - return self.k3_proj_weight is not None - - def finalize_decode_weights(self) -> None: - """The built-in fused weights, then one ``[q | k | v | g | f_a | b | pad]`` weight of both.""" - super().finalize_decode_weights() - assert ( - self.use_full_rank_gate - and self.gate_lower_bound is not None - and self._qkvg_proj_weight is not None - and self._bfa_proj_weight is not None - and self._qkvg_proj_weight.dtype == self._bfa_proj_weight.dtype == torch.bfloat16 - ), ( - f"Kimi K3 KDA layer {self.layer_idx}: the decode kernels read the bf16 fused projections the built-in " - "module builds on CUDA at head dim 128, with a full-rank output gate and a gate lower bound" - ) - rows = self._qkvg_proj_weight.shape[0] - with torch.no_grad(): - fused = self._merge_projection_weights( - (self.q_proj, self.k_proj, self.v_proj, self.g_proj, self.f_a_proj, self.b_proj), - pad_rows_to=8, - ) - self.k3_proj_weight = fused - self._qkvg_proj_weight, self._bfa_proj_weight = fused[:rows], fused[rows:] - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether ``forward`` runs ``step`` on the decode branch: outside a breakable CUDA graph.""" - return not is_in_breakable_cuda_graph() - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - reduce_output: bool = True, - project_output: bool = True, - ) -> torch.Tensor: - """``forward_generic`` where ``will_run_decode_branch`` does not hold. On the decode branch: the plain decode - on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the built-in dispatch; then - ``o_proj`` on its decode GEMV site and the TP all-reduce. - - ``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce), and, on the decode branch only, - ``project_output=False`` the post-o_norm core ``[N, H * 128]`` (no ``o_proj``).""" - if not self.will_run_decode_branch(attn_metadata, step): - return self.forward_generic(hidden_states, attn_metadata, reduce_output) - if ( - step.decode - and step.tokens_per_request == 1 - and self.takes_k3_kernels - and self.k3_buffers is not None - ): - core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) - else: - core = self._forward_impl(hidden_states, attn_metadata) - if not project_output: - return core.reshape(-1, self.proj_size) - return self._k3_project_output(core, reduce_output) - - def forward_generic( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - reduce_output: bool = True, - ) -> torch.Tensor: - if reduce_output: - return super().forward(hidden_states, attn_metadata) - return self.o_proj( - self._builtin_core(hidden_states, attn_metadata).reshape(-1, self.proj_size) - ) - - def _builtin_core( - self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata - ) -> torch.Tensor: - """The built-in forward's core ``[N, H, 128]``: inside a breakable CUDA graph from the eager - ``maybe_bcg_kda_core_inplace``, else from ``_forward_impl``.""" - if self.register_to_config and is_in_breakable_cuda_graph(): - core = hidden_states.new_empty( - (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 - ) - maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) - return core - return self._forward_impl(hidden_states, attn_metadata) - - def _k3_project_output(self, core: torch.Tensor, reduce_output: bool = True) -> torch.Tensor: - """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then, with - ``reduce_output``, the TP all-reduce.""" - core2d = core.reshape(-1, self.proj_size) - out = None - if self.decode_gemvs is not None: - out = self.decode_gemvs.project("o_proj", core2d, self.o_proj.weight) - if out is None: - if reduce_output: - return self._project_output(core) - out = self.o_proj(core2d) - if reduce_output and self._o_allreduce is not None: - out = self._o_allreduce(out) - return out - - def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: - """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; - each slot's conv window and state advance in place.""" - mamba_metadata = attn_metadata.mamba_metadata - slots = getattr(mamba_metadata, "generation_state_indices", None) - if slots is None: - slots = mamba_metadata.state_indices[: x.shape[0]] - layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) - w_q, w_k, w_v = self._get_mtp_conv_weights() - core = k3_kda_decode_attn( - x.contiguous(), - self.k3_proj_weight, - self.f_b_proj.weight, - w_q, - w_k, - w_v, - self._A_log_f32, - self._dt_bias_f32, - self._onorm_w_f32, - layer_cache.conv, - layer_cache.temporal, - slots, - self.k3_buffers, - float(self.gate_lower_bound), - self.head_k_dim**-0.5, - float(self.o_norm.eps), - ) - # Speculative decoding's replay caches keep their committed conv window in step with the pool's. - self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) - return core - - def forward_verify( - self, - x2d, - num_steps, - layer_cache, - conv_pool, - ssm_pool, - slot_indices, - output: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """The built-in verify, unless the cache manager keeps the per-token states; then the decode kernels'.""" - if not (layer_cache.has_kda_replay_caches and layer_cache.kda_state_tok is not None): - return super().forward_verify( - x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices, output=output - ) - assert self.takes_k3_kernels and self.k3_buffers is not None, ( - f"Kimi K3 KDA layer {self.layer_idx}: the cache manager keeps the per-token verify states, which only " - "the decode kernels write, and this layer has no fused projection weight or buffers" - ) - core = self._k3_verify(x2d, num_steps, layer_cache, ssm_pool, slot_indices) - return self._store_core(core, output) - - def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch.Tensor: - """The core output ``[N T, H, 128]`` of N requests of T verify tokens: ``ssm/k3_kda_attn`` for one request of - 8 tokens, else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it - takes them). Each slot's state after the golden token, its drafts' states and its conv window are written in - place.""" - num_requests = x.shape[0] // num_steps - slots = slot_indices[:num_requests] - w_q, w_k, w_v = self._get_mtp_conv_weights() - constants = (float(self.gate_lower_bound), self.head_k_dim**-0.5, float(self.o_norm.eps)) - if num_requests == 1 and num_steps == 8: - out = k3_kda_attn( - x.contiguous(), - self.k3_proj_weight, - self.f_b_proj.weight, - w_q, - w_k, - w_v, - self._A_log_f32, - self._dt_bias_f32, - self._onorm_w_f32, - layer_cache.kda_conv_q, - layer_cache.kda_conv_k, - layer_cache.kda_conv_v, - ssm_pool, - layer_cache.kda_state_tok, - slots, - layer_cache.prev_num_accepted_tokens, - self.k3_buffers, - num_steps - 1, - *constants, - ) - else: - rows = None - if self.decode_gemvs is not None: - rows = self.decode_gemvs.project("kda_proj", x, self.k3_proj_weight) - if rows is None: - rows = torch.nn.functional.linear(x, self.k3_proj_weight) - out = k3_kda_verify( - rows, - self.f_b_proj.weight, - w_q, - w_k, - w_v, - self._A_log_f32, - self._dt_bias_f32, - self._onorm_w_f32, - layer_cache.kda_conv_q, - layer_cache.kda_conv_k, - layer_cache.kda_conv_v, - ssm_pool, - layer_cache.kda_state_tok, - slots, - layer_cache.prev_num_accepted_tokens, - num_steps - 1, - *constants, - ) - return out.view(-1, self.num_heads, self.head_dim) - - -class K3DecodeMLA(KimiK3MLAAttention): - """Kimi K3's MLA attention: the built-in module, with a decode step's attention on the decode kernels. - - A decode step runs ``x [W_a; W_g]^T`` with the gate columns through a sigmoid (the ``mla_ag`` decode GEMV site - where it takes the rows, else one GEMM), then ``attention/k3_mla_qkv`` (the q_a / kv_a RMSNorms, q_b and the k_b - absorption into the fused query, the step's latent rows stored into the paged cache) and - ``attention/k3_mla_attn_vb_out`` (the attention over the paged cache, v_b and the output gate in one launch), then - ``o_proj`` (the ``o_proj`` site where it takes the rows, else the module). Every other step, and a decode step whose - cache the kernels do not read (``k3_mla_decode_view`` says why), runs the built-in module. - - ``[W_a; W_g]`` is built at load from the module's weights, which become views of it. The attention workspace is - the device's ``K3MlaAttnWorkspace``; it and the decode GEMVs' state are set by the target in - ``post_load_weights``. - """ - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - # [W_a; W_g]: the fused q_a / kv_a projection's rows, then the output gate's. - self.k3_ag_weight: Optional[torch.Tensor] = None - # The decode attention's workspace: one per device, shared by every MLA layer. - self.k3_workspace: Optional[K3MlaAttnWorkspace] = None - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - def post_load_weights(self) -> None: - """The built-in post-load, then ``[W_a; W_g]`` (once: CUDA graphs captured since read it).""" - super().post_load_weights() - if self.k3_ag_weight is not None: - return - gaps = self._k3_layout_gaps() - assert not gaps, ( - f"Kimi K3 MLA layer {self.layer_idx}: the decode kernels do not take {'; '.join(gaps)}" - ) - qkv_a, gate = self.kv_a_proj_with_mqa, self.g_proj - rows = qkv_a.weight.shape[0] - with torch.no_grad(): - fused = torch.cat([qkv_a.weight, gate.weight]) - qkv_a.weight = nn.Parameter(fused[:rows], requires_grad=False) - gate.weight = nn.Parameter(fused[rows:], requires_grad=False) - self.k3_ag_weight = fused - - def _k3_layout_gaps(self) -> list: - """What of this layer the decode kernels do not take (empty when they take all of it).""" - if not (self.use_output_gate and self.fuse_qkv_a_proj and not self.is_lite): - return ["a layer without the output gate or the fused q_a / kv_a projection"] - linears = (self.kv_a_proj_with_mqa, self.g_proj, self.q_b_proj, self.o_proj) - checks = ( - (not self.mapping.has_cp_helix(), "helix context parallelism"), - (not self.apply_rotary_emb and not self.llama_4_scaling, "RoPE or llama-4 scaling"), - (self.sparse_attn_hooks is None, "sparse attention"), - ( - self.kv_cache_dtype != "fp8_ds_mla" - and not getattr(self.mqa, "has_fp8_kv_cache", False) - and not getattr(self.mqa, "has_fp4_kv_cache", False), - "a quantized KV cache", - ), - ( - all(m.weight.dtype == torch.bfloat16 and m.bias is None for m in linears) - and self.k_b_proj_trans.dtype == self.v_b_proj.dtype == torch.bfloat16, - "projections other than bf16 and unbiased", - ), - ( - not getattr(self.q_a_layernorm, "is_nvfp4", False) - and not getattr(self.kv_a_layernorm, "use_gemma", False), - "an NVFP4 q_a norm or a Gemma kv_a norm", - ), - ( - self.num_heads_tp % 6 == 0 - and self.kv_lora_rank == 512 - and self.qk_rope_head_dim == 64 - and self.q_lora_rank == 1536, - f"{self.num_heads_tp} heads, latent {self.kv_lora_rank}, rope {self.qk_rope_head_dim}, " - f"q_lora {self.q_lora_rank}", - ), - ) - return [why for ok, why in checks if not ok] - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether ``forward`` runs ``step`` on the decode kernels (see ``_k3_step_view``). Only there does it return - the gated attention output before ``o_proj``; the KV writes forbid running the attention twice, so a caller - that needs it asks first.""" - return self._k3_step_view(attn_metadata, step) is not None - - def _k3_step_view(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> Optional[dict]: - """The paged-cache view the decode kernels read on ``step``: a decode step, this layer's ``[W_a; W_g]`` and - workspace built, outside a breakable CUDA graph, and a cache ``k3_mla_decode_view`` takes. Else None.""" - if not ( - step.decode - and self.k3_ag_weight is not None - and self.k3_workspace is not None - and not is_in_breakable_cuda_graph() - ): - return None - return self._k3_decode_view(attn_metadata, step.num_tokens) - - def forward( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - project_output: bool = True, - ) -> torch.Tensor: - """``forward_generic``, except on a decode step whose cache the decode kernels read - (``will_run_decode_branch``). There only, ``project_output=False`` returns the gated attention output - ``[M, H * 128]``, the input of ``o_proj``.""" - view = self._k3_step_view(attn_metadata, step) - if view is None: - if not project_output: - raise ValueError("project_output=False needs a step will_run_decode_branch takes") - return self.forward_generic(position_ids, hidden_states, attn_metadata) - x = hidden_states[: step.num_tokens].contiguous() - rows = self.kv_a_proj_with_mqa.weight.shape[0] - gemvs = self.decode_gemvs - ag = None if gemvs is None else gemvs.project("mla_ag", x, self.k3_ag_weight) - if ag is None: - ag = torch.nn.functional.linear(x, self.k3_ag_weight) - ag[:, rows:].sigmoid_() - fused_q = k3_mla_qkv( - ag, - self.q_a_layernorm.weight, - float(self.q_a_layernorm.variance_epsilon), - self.q_b_proj.weight, - self.k_b_proj_trans, - self.kv_a_layernorm.weight, - float(self.kv_a_layernorm.variance_epsilon), - view["pool"], - view["row_stride"], - view["page_table"], - view["page_offset"], - view["seq_len"], - ) - attn_output = self.create_output(x, 0) - k3_mla_attn_vb_out( - fused_q, - view["pool"], - view["row_stride"], - view["page_table"], - view["page_offset"], - view["seq_len"], - view["softmax_scale"], - self.v_b_proj, - attn_output, - self.k3_workspace, - gate=ag, - gate_col0=rows, - ) - if not project_output: - return attn_output - out = None if gemvs is None else gemvs.project("o_proj", attn_output, self.o_proj.weight) - if out is None: - out = self._project_output([attn_output], position_ids, attn_metadata, None) - return out - - def forward_generic( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - return super().forward(position_ids, hidden_states, attn_metadata) - - def _k3_decode_view(self, attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[dict]: - """The paged latent cache as the decode kernels read it this step, or None when they do not read it (the - reason is logged once).""" - view = k3_mla_decode_view(self.mqa, attn_metadata, num_tokens) - if isinstance(view, str): - logger.info_once( - f"Kimi K3 MLA: the built-in path for a decode step the decode kernels do not read ({view})", - key=f"k3_mla_decode_view_{view}", - ) - return None - return view - - -# ---------------------------------------------------------------------------------------------------------------------- -# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. -# ---------------------------------------------------------------------------------------------------------------------- - -#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at -#: max_draft_len 7. -DRAFTER_ATTN_SPLITS = frozenset( - {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} -) -# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k -# RMSNorm's epsilon and the NeoX RoPE's base. -_DRAFTER_HEADS = (6, 1) -_DRAFTER_HEAD_DIM = 64 -_DRAFTER_PAGE = 64 -_DRAFTER_EPS = 1e-5 -_DRAFTER_ROPE_BASE = 10000.0 - - -def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: - """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input - feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" - if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: - return None - width = in_features // tp_size - return tp_rank * width, (tp_rank + 1) * width - - -class K3FcSlice(nn.Module): - """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is - ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's - ranks is ``fc``'s output.""" - - def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: - super().__init__() - self.weight = nn.Parameter(weight, requires_grad=False) - self.start = start - self.end = end - - @classmethod - def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: - """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns - (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" - columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) - if columns is None: - return None - start, end = columns - return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the - GEMM reads the column block in place, through the rows' stride.""" - return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) - - -class K3DSparkDrafter(GQADSparkForCausalLM): - """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. - - A block the entries take runs, per layer: - - * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; - * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to - its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the - worker writes the accepted rows' context k / v before a block reads them; - * the output projection on ``drafter_o``, then the module's all-reduce; - * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the - module's all-reduce. - - The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state - (below); a projection whose site does not take its rows runs its module. Every other block runs the stock - ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or - normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under - CUDA-graph capture. - - Where the target hands the drafter the TP group's collective state (``use_decode_comm``): - - * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and - ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; - * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the - post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the - block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the - projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection - after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is - followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers - whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. - - Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the - stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); - the context k / v and the Markov head stay upstream code. - """ - - def __init__( - self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" - ) -> None: - super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. - self._k3_take: Optional[bool] = None - # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes - # only these. - self._k3_attn_ran: set = set() - # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). - self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None - # The zero residual of the split context projection's all-reduce, up to a decode step's rows. - self._k3_zero_rows: Optional[torch.Tensor] = None - # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms - # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. - self._k3_norms_fuse = False - self._k3_sandwich_forms: frozenset = frozenset() - - def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: - """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the - target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are - loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. - - * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they - split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` - sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. - * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and - RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose - shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not - compile is not used.""" - self.decode_comm = comm - self._k3_split_fc() - self._k3_norms_fuse = self._k3_norms_take_comm() - forms = set() - if self._k3_norms_fuse: - layers = self.model.layers - for form, swiglu, weights in ( - ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), - ("down", True, [layer.mlp.down_proj.weight for layer in layers]), - ): - if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( - weights[0], swiglu=swiglu - ): - forms.add(form) - self._k3_sandwich_forms = frozenset(forms) - logger.info( - "Kimi K3 DSpark drafter: residual adds and RMSNorms " - + ( - f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " - "mnnvl_fusion_allreduce)" - if self._k3_norms_fuse - else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" - ) - ) - - def _k3_norms_take_comm(self) -> bool: - """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and - post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no - quantized output), and every output and down projection is row parallel with its all-reduce.""" - hidden = self.config.hidden_size - - def plain(norm: nn.Module) -> bool: - weight = getattr(norm, "weight", None) - return ( - isinstance(weight, torch.Tensor) - and weight.dtype == torch.bfloat16 - and tuple(weight.shape) == (hidden,) - and hasattr(norm, "variance_epsilon") - and not getattr(norm, "use_gemma", True) - and not getattr(norm, "is_nvfp4", True) - and not getattr(norm, "return_hp_output", True) - ) - - def reduces(linear: nn.Module) -> bool: - return ( - getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" - and getattr(linear, "reduce_output", False) - and getattr(linear, "all_reduce", None) is not None - ) - - return plain(self.model.norm) and all( - plain(layer.input_layernorm) - and plain(layer.post_attention_layernorm) - and reduces(layer.self_attn.o_proj) - and reduces(layer.mlp.down_proj) - for layer in self.model.layers - ) - - def _k3_tp_all_reduce(self) -> Optional[nn.Module]: - """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" - o_proj = self.model.layers[0].self_attn.o_proj - if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": - return None - return getattr(o_proj, "all_reduce", None) - - def _k3_split_fc(self) -> None: - """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero - residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is - a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP - all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" - fc = getattr(self, "fc", None) - if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): - return - mapping = self.model_config.mapping - weight = getattr(fc, "weight", None) - sliced = None - if ( - isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.dtype == torch.bfloat16 - and getattr(fc, "bias", None) is None - and self._k3_tp_all_reduce() is not None - ): - sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) - if sliced is None: - logger.info( - "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " - f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" - ) - return - self.fc = sliced - self._k3_zero_rows = weight.new_zeros( - MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] - ) - logger.info( - f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " - f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" - ) - - def load_weights(self, weights, weight_mapper=None, **kwargs): - """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" - result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) - self._k3_split_fc() - return result - - def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. - - With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group - with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode - step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then - ``hidden_norm``. Otherwise the stock projection.""" - fc = self.fc - if not isinstance(fc, K3FcSlice): - return super().project_target_hidden(hidden_states) - partial = fc(hidden_states.to(fc.weight.dtype)) - rows = partial.shape[0] - zeros = self._k3_zero_rows - if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): - normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) - return normed - if rows == 0: - return self.hidden_norm(partial) - return self.hidden_norm(self._k3_tp_all_reduce()(partial)) - - def dflash_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_k_cache: torch.Tensor, - ctx_v_cache: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor] = None, - ctx_page_table: Optional[torch.Tensor] = None, - ctx_rows_start: Optional[int] = None, - ) -> torch.Tensor: - """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class - docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows - ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows - as a view instead of gathering them.""" - keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) - if keys is None: - return super().dflash_forward( - noise_embedding, - query_positions, - num_ctx_per_req, - ctx_k_cache, - ctx_v_cache, - ctx_cache_batch_idx, - ctx_kv_cache, - ctx_page_table, - ) - out = self._k3_block_forward( - noise_embedding, - query_positions, - num_ctx_per_req, - ctx_cache_batch_idx, - ctx_kv_cache, - ctx_page_table, - ctx_rows_start, - ) - if not torch.cuda.is_current_stream_capturing(): - self._k3_attn_ran |= keys - return out - - def _k3_layers_take(self) -> bool: - """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head - layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, - non-causal unwindowed block attention, and no block convolution or post-attention gate.""" - if self._k3_take is None: - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - take = ( - self._use_fused_qk_norm_rope - and not self.has_block_conv - and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate - ) - for layer_idx, layer in enumerate(self.model.layers): - attn, mlp = layer.self_attn, layer.mlp - pos = getattr(attn, "pos_embd_params", None) - rope = getattr(pos, "rope", None) - norms = (attn.q_norm, attn.k_norm) - take = take and ( - (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS - and attn.head_dim == _DRAFTER_HEAD_DIM - and getattr(attn, "is_qk_norm", False) - and not getattr(attn, "use_gemma_rms_norm", True) - and rope is not None - and pos.is_neox - and getattr(pos, "mrope_section", None) is None - and rope.theta == _DRAFTER_ROPE_BASE - and all( - norm.variance_epsilon == _DRAFTER_EPS - and norm.weight.dtype == torch.bfloat16 - and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) - for norm in norms - ) - and attn.qkv_proj.bias is None - and attn.o_proj.bias is None - and mlp.activation is torch.nn.functional.silu - and mlp.swiglu_limit is None - and mlp.swiglu_alpha in (None, 1.0) - and mlp.swiglu_beta in (None, 0.0) - and mlp.gate_up_proj.bias is None - and mlp.down_proj.bias is None - and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) - ) - self._k3_take = bool(take) - logger.info( - f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " - + ( - "on k3_drafter_attn_qknorm and the drafter GEMV sites" - if self._k3_take - else "on the stock block forward (a layer the drafter entries do not take)" - ) - ) - return self._k3_take - - def _k3_block_keys( - self, - noise_embedding: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor], - ctx_page_table: Optional[torch.Tensor], - ) -> Optional[frozenset]: - """The attention's compile keys for this block when the drafter entries take it, else None.""" - if ( - self.dflash_attention_backend != "TRTLLM" - or ctx_kv_cache is None - or ctx_page_table is None - or noise_embedding.dtype != torch.bfloat16 - or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS - or is_in_breakable_cuda_graph() - or not self._k3_layers_take() - ): - return None - num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads - page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) - dense = ( - num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, - _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, - ) - strides = set() - for layer_idx in range(len(self.model.layers)): - cache = ctx_kv_cache[layer_idx] - if ( - cache.dtype != torch.bfloat16 - or cache.dim() != 5 - or tuple(cache.shape[1:]) != page - or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) - or cache.stride(0) % 8 - ): - return None - strides.add(cache.stride(0)) - keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) - if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: - return None - return keys - - def _k3_block_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: torch.Tensor, - ctx_page_table: torch.Tensor, - ctx_rows_start: Optional[int] = None, - ) -> torch.Tensor: - batch, block = noise_embedding.shape[:2] - rows = batch * block - ctx_len = num_ctx_per_req[:batch].to(torch.int32) - if ctx_rows_start is None: - page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) - else: - # The batch's page-table rows are one contiguous run: a view, no gather. - page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] - positions = query_positions.reshape(-1).contiguous() - hidden = noise_embedding.reshape(rows, -1) - layers = self.model.layers - comm = self._k3_fused_comm(hidden) - residual = None - if comm is not None: - # The fused all-reduces read the residual and return the updated one: the block's input serves as the - # first residual, uncopied. - residual = hidden - normed = layers[0].input_layernorm(hidden) - for layer_idx, layer in enumerate(layers): - attn = layer.self_attn - if comm is None: - if residual is None: - residual = hidden.clone() - normed = layer.input_layernorm(hidden) - else: - normed, residual = layer.input_layernorm(hidden, residual) - qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) - out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) - k3_drafter_attn_qknorm( - qkv, - attn.q_norm.weight, - attn.k_norm.weight, - positions, - attn.q_norm.variance_epsilon, - attn.pos_embd_params.rope.theta, - ctx_kv_cache[layer_idx], - page_table, - ctx_len, - attn.num_heads, - attn.num_key_value_heads, - out, - ) - if comm is not None: - # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after - # the last layer the final norm. - next_norm = ( - layers[layer_idx + 1].input_layernorm - if layer_idx + 1 < len(layers) - else self.model.norm - ) - normed, residual = self._k3_project_norm( - comm, out, attn.o_proj, residual, layer.post_attention_layernorm - ) - normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) - continue - hidden = self._k3_project("drafter_o", out, attn.o_proj) - hidden, residual = layer.post_attention_layernorm(hidden, residual) - hidden = self._k3_mlp(layer.mlp, hidden) - if comm is not None: - return normed - out, _ = self.model.norm(hidden, residual) - return out - - def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: - """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its - all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace - holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" - comm = self.decode_comm - if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): - return None - return comm - - def _k3_project_norm( - self, - comm: _decode_comm.K3DecodeComm, - x: torch.Tensor, - linear: nn.Module, - residual: torch.Tensor, - norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, - the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the - projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` - where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" - if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( - x, linear.weight, residual, norm - ): - return comm.sandwich_plain(x, linear.weight, residual, norm) - gemvs = self.decode_gemvs - partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) - if partial is None: - partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) - return comm.allreduce_norm(partial, residual, norm) - - def _k3_mlp_norm( - self, - comm: _decode_comm.K3DecodeComm, - mlp: nn.Module, - x: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down - projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul - form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down - projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; - then ``comm/mnnvl_fusion_allreduce``.""" - gemvs = self.decode_gemvs - gate_up = ( - None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) - ) - partial = None - if gate_up is not None: - if "down" in self._k3_sandwich_forms and comm.takes_plain( - gate_up, mlp.down_proj.weight, residual, norm, swiglu=True - ): - return comm.sandwich_plain( - gate_up, mlp.down_proj.weight, residual, norm, swiglu=True - ) - partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) - if partial is None: - partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) - return comm.allreduce_norm(partial, residual, norm) - - def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: - """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel - projection's all-reduce.""" - y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) - if y is None: - return linear(x) - return _row_parallel_reduce(linear, y) - - def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: - """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on - ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" - gemvs = self.decode_gemvs - gate_up = ( - None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) - ) - down = ( - None - if gate_up is None - else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) - ) - if down is None: - return mlp(x) - return _row_parallel_reduce(mlp.down_proj, down) - - -def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: - """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" - all_reduce = getattr(linear, "all_reduce", None) - if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: - return all_reduce(partial) - return partial - - -# ---------------------------------------------------------------------------------------------------------------------- -# The target: step classification, the construction checks and the class route returns. -# ---------------------------------------------------------------------------------------------------------------------- - - -def _text_model_config(model_config: ModelConfig) -> ModelConfig: - """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. - - The checkpoint names its language-model modules `language_model.`; the text model's are ``, with - `layers.*` under `model.`. - """ - config = model_config.pretrained_config - text = copy.copy(model_config) - text._frozen = False - text.pretrained_config = config.text_config - excluded = text.quant_config.exclude_modules - if excluded: - text.quant_config = copy.copy(text.quant_config) - renamed = [] - for name in excluded: - if name.startswith(_LANG_PREFIX): - name = name[len(_LANG_PREFIX) :] - if name.startswith("layers."): - name = "model." + name - renamed.append(name) - text.quant_config.exclude_modules = renamed - text.skip_create_weights_in_init = True - text._frozen = True - return text - - -@dataclass(frozen=True) -class DecodeStep: - """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` - generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" - - num_tokens: int - num_requests: Optional[int] = None - tokens_per_request: Optional[int] = None - #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` - #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. - latent_push: bool = False - - @property - def small(self) -> bool: - """Whether the step fits one token tile of the token-count kernels.""" - return self.num_tokens <= DECODE_MAX_TOKENS - - @property - def decode(self) -> bool: - """Whether the step is a pure decode step the request-aware kernels take.""" - return self.num_requests is not None - - @property - def wide(self) -> bool: - """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode - layout's MoE head and tail, on M-general ops.""" - return self.decode and not self.small - - -def latent_push( - step: DecodeStep, - *, - capturing: bool, - breakable: bool, - kda_draft_records: bool, - kda_decode_kernels: bool, - mla_decode_branch: bool, -) -> bool: - """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at - most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a - breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels - (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA - layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). - Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel - between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a - graph launch orders every pushing replay behind whatever ran before it.""" - if not (step.decode and step.small) or not capturing or breakable: - return False - if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): - return False - return mla_decode_branch - - -def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: - """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). - - ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers - and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input - here is fixed by that shape, so a captured step and its replays are classified alike. - """ - if num_tokens <= 0: - return None - requests = _decode_requests(attn_metadata, num_tokens) - if requests is not None: - return DecodeStep(num_tokens, requests, num_tokens // requests) - if num_tokens <= DECODE_MAX_TOKENS: - return DecodeStep(num_tokens) - return None - - -def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: - """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" - if attn_metadata.num_contexts != 0: - return None - requests = attn_metadata.num_generations - if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: - return None - tokens = num_tokens // requests - if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: - return None - seq_lens = getattr(attn_metadata, "seq_lens", None) - if seq_lens is not None and seq_lens.device.type == "cpu": - lens = seq_lens[:requests] - if lens.numel() != requests or bool((lens != tokens).any()): - return None - return requests +_TARGET = K3TargetSpec(sm=(10, 0), quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, moe_tp=4, moe_ep=4) def _check_construction(model_config: ModelConfig) -> None: - """The settings this target is built for that are fixed before the first step.""" - capability = torch.cuda.get_device_capability() - assert capability == _SM, ( - f"this target is certified on sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" - ) - mapping = model_config.mapping - topology = ( - mapping.world_size, - mapping.tp_size, - mapping.pp_size, - mapping.moe_tp_size, - mapping.moe_ep_size, - mapping.enable_attention_dp, - ) - assert topology == (16, 16, 1, 4, 4, False), ( - "the tp16_moetp4ep4 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " - "moe_tensor_parallel_size 4, moe_expert_parallel_size 4 and enable_attention_dp false; the engine built " - f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" - ) - assert model_config.torch_dtype == torch.bfloat16, ( - f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" - ) - kv_algo = model_config.quant_config.kv_cache_quant_algo - assert kv_algo is None, ( - f"this target's MLA kernels read a bf16 KV pool; kv_cache_config.dtype resolved to {kv_algo}" - ) - quant_algo = model_config.quant_config.quant_algo - assert quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( - "this target loads the MXFP4 checkpoint, whose compressed-tensors config the model config reads as " - "W4A16_MXFP4 with no per-layer declarations (its routed experts run the W4A8_MXFP4_MXFP8 default); the " - f"engine read {quant_algo} with {len(model_config.quant_config_dict or {})} per-layer declarations" - ) - strategy = model_config.allreduce_strategy - assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( - f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" - ) - - -class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(KimiLinearForCausalLM): - """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. - - It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through - weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model - keeps the built-in one's. - """ - - @classmethod - def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: - """The V2 hybrid manager holds the KDA states; the step contract requires it.""" - return "V2" - - def __init__(self, model_config: ModelConfig): - config = model_config.pretrained_config - assert getattr(config, "text_config", None) is not None, ( - "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " - "text_config" - ) - _check_construction(model_config) - text = _text_model_config(model_config) - spec_config = getattr(text, "spec_config", None) - assert ( - spec_config is None - or spec_config.spec_dec_mode.is_sa() - or spec_config.spec_dec_mode.is_dflash() - or spec_config.spec_dec_mode.is_dspark() - ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" - # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 - # weight-read conversion of the shared / latent MLPs. - self._fp8_weight_read_moe_mlp = False - self.mapping_with_cp = None - self._repurposed_tp_mapping = None - cfg = text.pretrained_config - SpecDecOneEngineForCausalLM.__init__( - self, - KimiLinearModel(text), - text, - hidden_size=cfg.hidden_size, - vocab_size=cfg.vocab_size, - ) - self._own_spec_branch(text) - self._step_checked = False - # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the - # processor the speculative worker's target logits and a parallel drafter's logits call too. - stock_logits_processor = self.logits_processor - self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) - draft_model = getattr(self, "draft_model", None) - if getattr(draft_model, "logits_processor", None) is stock_logits_processor: - draft_model.logits_processor = self.logits_processor - # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which - # must therefore be the text config, as the built-in wrapper leaves it. - model_config._frozen = False - model_config.pretrained_config = self.config - model_config._frozen = True - - def _own_spec_branch(self, model_config: ModelConfig) -> None: - """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a - standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is - `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" - spec_config = model_config.spec_config - if spec_config is None: - return - mode = spec_config.spec_dec_mode - standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) - if ( - standalone_dspark - and self.draft_config is not None - and not is_mla(self.draft_config.pretrained_config) - ): - drafter = K3DSparkDrafter( - self.draft_config, dflash_attention_backend=spec_config.attention_backend - ) - if mode.is_parallel_draft(): - drafter.logits_processor = self.logits_processor - self.epilogue[self.epilogue.index(self.draft_model)] = drafter - self.draft_model = drafter - if mode.is_dflash(): - worker_cls = _spec_worker.KimiK3DFlashWorker - elif standalone_dspark: - worker_cls = _spec_worker.KimiK3DSparkWorker - else: - return - worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) - worker.set_draft_model(self.draft_model) - self.epilogue[self.epilogue.index(self.spec_worker)] = worker - self.spec_worker = worker - - def load_weights(self, weights, *args, **kwargs): - _weights.load(self, weights) - - def cache_derived_state(self) -> None: - """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of - every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: - a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites - are among them, and the drafter gets the state too.""" - super().cache_derived_state() - if self.model.decode_gemvs is not None: - return - gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head) - self.model.decode_gemvs = gemvs - self.logits_processor.gemvs = gemvs - for layer in self.model.layers: - if not layer.is_moe: - layer.decode_gemvs = gemvs - if isinstance(self.draft_model, K3DSparkDrafter): - self.draft_model.decode_gemvs = gemvs - - def post_load_weights(self) -> None: - """The state the decode kernels share, built once per device before any CUDA-graph capture and handed to - every layer of its kind: the KDA projection's Lamport buffers and the MLA decode attention's workspace; the - decode GEMVs' state (built by ``cache_derived_state``) handed to every attention module; and, where every - attention all-reduce runs over MNNVL, the TP group's collective state (``K3DecodeComm``, collective: every - rank builds it here) handed to every layer, with whether its sandwich takes the layer's o_proj, the decode - path's one-shot ceiling on every stock MNNVL all-reduce (``use_decode_one_shot``), and the MoE decode path - (``_build_decode_moe``); then the speculative worker's decode kernels (``_gate_spec_worker_kernels``) and the - DSpark drafter's collectives (``_gate_drafter_comm``).""" - super().post_load_weights() - kda = [layer.linear_attn for layer in self.model.layers if layer.is_kda] - mla = [layer.self_attn.mixer for layer in self.model.layers if not layer.is_kda] - if kda[0].k3_buffers is not None: - return # built by an earlier call; CUDA graphs captured since hold it - device = self.model.embed_tokens.weight.device - buffers = K3KdaBuffers.create(device) - workspace = K3MlaAttnWorkspace.create(device, mla[0].num_heads_tp // 6) - for module in kda: - module.k3_buffers = buffers - for module in mla: - module.k3_workspace = workspace - for module in kda + mla: - module.decode_gemvs = self.model.decode_gemvs - layers = self.model.layers - comm = None - if all(layer._mnnvl_allreduce() is not None for layer in layers): - for layer in layers: - layer.sandwich_oproj = _decode_comm.K3DecodeComm.takes_oproj( - layer._o_proj(), self.model.num_attn_res_snapshots - ) - oproj = next((layer._o_proj().weight for layer in layers if layer.sandwich_oproj), None) - comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, oproj) - for layer in layers: - layer.decode_comm = comm - _decode_comm.use_decode_one_shot(self) - moe_layers = 0 if comm is None else self._build_decode_moe(comm) - logger.info( - "Kimi K3 decode kernels: KDA on k3_kda_decode_attn, k3_kda_attn and k3_kda_verify " - f"({sum(m.takes_k3_kernels for m in kda)} / {len(kda)} layers take them), MLA on k3_mla_qkv and " - f"k3_mla_attn_vb_out ({len(mla)} layers); the post-attention all-reduce of at most " - f"{_decode_comm.AR_ATTN_RES_MAX_TOKENS} tokens " - + ( - "with the residual update: k3_sandwich_oproj with o_proj at most " - f"{_decode_comm.SANDWICH_MAX_TOKENS} decode tokens " - f"({sum(layer.sandwich_oproj for layer in layers)} / {len(layers)} layers), else " - "mnnvl_allreduce_attn_res" - if comm is not None - else "unfused (an attention all-reduce does not run over MNNVL)" - ) - + f"; MoE on k3_moe_front, k3_moe and the row-parallel tail ({moe_layers} layers)" - ) - self._gate_spec_worker_kernels(comm) - self._gate_drafter_comm(comm) - - def _build_decode_moe(self, comm: _decode_comm.K3DecodeComm) -> int: - """The MoE decode path (``decode_moe.py``) on every MoE layer it takes: the shared state (collective: the - front's head workspace), each layer's decode weights with its latent norm folded into the latent up - projection, the decode GEMVs' state; then one call of every kernel of the path (collective: the front and - the sandwich tail). Every rank builds it here. Returns the number of MoE layers it takes.""" - mapping = self.model_config.mapping - moes = [layer.block_sparse_moe for layer in self.model.layers if layer.is_moe] - gaps = { - id(moe): _decode_moe.layout_gaps( - moe, mapping.tp_size, self.model.num_attn_res_snapshots - ) - for moe in moes - } - takes = [moe for moe in moes if not gaps[id(moe)]] - for reason in sorted({gap for moe in moes for gap in gaps[id(moe)]}): - logger.info_once( - f"Kimi K3 MoE decode path off on some layers: {reason}", - key=f"k3_decode_moe_off_{reason}", - ) - if not takes: - return 0 - backend = takes[0].routed_experts.backend - state = _decode_moe.K3DecodeMoe.create( - mapping, - backend.w3_w1_weight.device, - backend.w3_w1_weight.shape[1] // 2, - backend.expert_size_per_partition, - comm.mnnvl, - ) - for moe in takes: - _decode_moe.fold_latent_norm(moe) - moe.decode_moe = _decode_moe.K3DecodeMoeLayer.create( - moe, state, mapping.tp_rank, mapping.tp_size - ) - moe.decode_gemvs = self.model.decode_gemvs - first = takes[0].decode_moe - first.warm_up(takes[0]) - logger.info( - "Kimi K3 MoE decode path: the latent all-reduce at <= 8 tokens on " - + ( - "the routed experts' push form and k3_latent_reduce on pushing steps (a pure decode step captured " - "into a CUDA graph on the decode kernels), else the routed experts' all-reduce" - if state.exchange is not None - else f"the routed experts' all-reduce (no latent exchange at TP {mapping.tp_size})" - ) - ) - comm.compile_tail(takes[0].moe_hidden_size, first.shared_cols, first.tail_weight) - return len(takes) - - def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: - """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, - ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input - that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's - ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the - vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks - each step's own conditions. Returns the setting (False without such a worker).""" - worker = getattr(self, "spec_worker", None) - if not hasattr(worker, "k3_decode"): - return False - gemvs = self.model.decode_gemvs - worker.k3_decode = ( - comm is not None and gemvs is not None and gemvs.head_workspace is not None - ) - logger.info( - f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " - + ( - "on" - if worker.k3_decode - else "off (no MNNVL decode state or no k3_head_gemv LM head)" - ) - ) - return worker.k3_decode + _check_target_construction(model_config, _TARGET) - def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: - """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built - it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context - projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms - in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on - the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate - requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` - workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the - drafter took the state (False without such a drafter).""" - drafter = getattr(self, "draft_model", None) - if comm is None or not isinstance(drafter, K3DSparkDrafter): - return False - drafter.use_decode_comm(comm) - return True - def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: - """First-forward checks of the engine surface and the per-engine settings.""" - objects = { - "attn_metadata": attn_metadata, - "kv_cache_manager": attn_metadata.kv_cache_manager, - } - missing = [ - f"{owner}.{name}" - for owner, names in REQUIRED_ENGINE_FIELDS.items() - for name in names - if not hasattr(objects[owner], name) - ] - assert not missing, f"engine fields this target reads are missing: {missing}" - assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( - f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " - f"{attn_metadata.tokens_per_block}" - ) - manager = attn_metadata.kv_cache_manager - assert isinstance(manager, MambaHybridCacheManagerV2), ( - "this target needs the V2 hybrid KV / state cache manager " - "(kv_cache_config.use_kv_cache_manager_v2); the engine built " - f"{type(manager).__name__}" - ) - assert not manager.enable_block_reuse, ( - "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" - ) - kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) - state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype - assert state_dtype == torch.float32, ( - "this target's KDA kernels keep fp32 recurrent states (kv_cache_config.mamba_ssm_cache_dtype float32 or " - f"auto); the engine built a {state_dtype} state pool" - ) - self._step_checked = True +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(ModelingV2KimiK3): + """Concrete deployment of the shared Kimi K3 text model.""" - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: bool = False, - spec_metadata=None, - resource_manager=None, - **kwargs, - ) -> torch.Tensor: - if kwargs.pop("multimodal_params", None): - raise ValueError( - "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" - ) - if not self._step_checked: - self._check_step_contract(attn_metadata) - return super().forward( - attn_metadata, - input_ids, - position_ids, - inputs_embeds, - return_context_logits, - spec_metadata, - resource_manager, - **kwargs, - ) + target_spec = _TARGET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py deleted file mode 100644 index abc1a1388213..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py +++ /dev/null @@ -1,1343 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -import functools -import inspect -from typing import Optional - -import torch -from torch import nn - -from tensorrt_llm._torch.attention.backends import AttentionMetadata -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager -from tensorrt_llm._torch.speculative.dflash import ( - DFlashSpecMetadata, - DFlashWorker, - dflash_allocated_ctx_limit, - dflash_noise_block_embedding, -) -from tensorrt_llm._torch.speculative.dspark import DSparkWorker -from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase -from tensorrt_llm.logger import logger - - -@functools.lru_cache(maxsize=None) -def _takes_ctx_rows_start_cls(cls) -> bool: - return "ctx_rows_start" in inspect.signature(cls.dflash_forward).parameters - - -def _takes_ctx_rows_start(draft_model) -> bool: - """Whether the draft model's ``dflash_forward`` accepts ``ctx_rows_start``.""" - return _takes_ctx_rows_start_cls(type(draft_model)) - - -def dflash_ctx_rows_kwargs(draft_model, ctx_rows_start: Optional[int]) -> dict: - """The ``ctx_rows_start`` keyword (from ``prepare_1st_drafter_inputs``) for a draft model whose ``dflash_forward`` - takes it; no keyword for another draft model or without a start.""" - if ctx_rows_start is None or not _takes_ctx_rows_start(draft_model): - return {} - return {"ctx_rows_start": ctx_rows_start} - - -def capture_view( - spec_metadata: Optional[SpecMetadata], layer_id: int, num_tokens: int -) -> Optional[torch.Tensor]: - """``layer_id``'s slot of the DFlash capture buffer for ``num_tokens`` rows (a strided view a kernel can write the - tap into), or None when ``spec_metadata`` is not DFlash's or the layer is not captured.""" - if ( - not isinstance(spec_metadata, DFlashSpecMetadata) - or spec_metadata.captured_hidden_states is None - ): - return None - i = spec_metadata._layer_to_idx.get(layer_id) - if i is None: - return None - return spec_metadata.captured_hidden_states[ - :num_tokens, i * spec_metadata.hidden_size : (i + 1) * spec_metadata.hidden_size - ] - - -class KimiK3DFlashWorker(DFlashWorker): - # Set by a Kimi K3 target running its decode kernels: decode steps then run the acceptance and the drafter's - # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context - # K/V with trtllm::k3_ctx_kv. - k3_decode = False - - def _mask_token_id(self, draft_model) -> int: - """The drafter's mask token id, resolved once.""" - if self._resolved_mask_token_id is None: - if ( - hasattr(self.spec_config, "mask_token_id") - and self.spec_config.mask_token_id is not None - ): - self._resolved_mask_token_id = self.spec_config.mask_token_id - elif hasattr(draft_model, "mask_token_id"): - self._resolved_mask_token_id = draft_model.mask_token_id - elif hasattr(draft_model.model, "mask_token_id"): - self._resolved_mask_token_id = draft_model.model.mask_token_id - else: - raise ValueError( - "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " - "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." - ) - return self._resolved_mask_token_id - - @staticmethod - def _trained_mask_embedding(draft_model, mask_token_id: int) -> Optional[torch.Tensor]: - """The drafter's own trained mask row, if it kept one for the mask id in use.""" - if mask_token_id != getattr(draft_model, "mask_token_id", None): - return None - return getattr(draft_model, "mask_token_embedding", None) - - def _k3_accept_applies( - self, logits, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int - ) -> bool: - """Whether this step's acceptance and the drafter's inputs run as ``trtllm::k3_spec_accept``. - - That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no - rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target - logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block - table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. - """ - if not self._k3_accept_step_applies( - attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ): - return False - K = spec_metadata.runtime_draft_len - if ( - logits.dim() != 2 - or logits.dtype != torch.float32 - or logits.shape[0] != num_gens * (K + 1) - ): - return False - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - - embed = draft_model.draft_model_full.model.embed_tokens - return accept_op.supports( - logits.shape[1], num_gens, self._compute_block_size, K, embed.weight.shape[1] - ) - - def _k3_accept_step_applies( - self, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int - ) -> bool: - """Everything ``_k3_accept_applies`` checks but the target logits.""" - if not self.k3_decode: - return False - from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MambaHybridCacheManagerV2, - ) - - if num_contexts != 0 or not 0 < num_gens <= 8 or self.guided_decoder is not None: - return False - if ( - not spec_metadata.is_all_greedy_sample - or self._can_use_rejection_sampling(spec_metadata) - or getattr(spec_metadata, "enable_penalty", False) - or type(self)._reshape_draft_tokens_for_accept - is not SpecWorkerBase._reshape_draft_tokens_for_accept - or type(self)._reshape_logits_for_accept - is not SpecWorkerBase._reshape_logits_for_accept - ): - return False - K = spec_metadata.runtime_draft_len - draft_tokens = spec_metadata.draft_tokens - if ( - K <= 0 - or draft_tokens is None - or draft_tokens.dtype != torch.int32 - or draft_tokens.numel() != num_gens * K - ): - return False - mgr = attn_metadata.kv_cache_manager - mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) - if ( - not isinstance(mgr, MambaHybridCacheManagerV2) - or type(mgr).update_mamba_states is not MambaHybridCacheManagerV2.update_mamba_states - or not mgr.use_kda_replay_update - or getattr(mgr, "prev_num_accepted_tokens", None) is None - or getattr(mgr, "_dummy_request_mask", None) is None - or mamba_metadata is None - or mamba_metadata.state_indices.dtype != torch.int32 - ): - return False - if ( - self._ctx_block_tables is None - or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None - or getattr(attn_metadata, "kv_lens_cuda", None) is None - ): - return False - embed = draft_model.draft_model_full.model.embed_tokens - return getattr(embed, "tp_size", 1) == 1 and embed.weight.dtype == torch.bfloat16 - - def target_logits( - self, hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model - ) -> torch.Tensor: - """This rank's bf16 vocabulary shard of the target logits, [rows, vocab / TP] (``_k3_head_shard``: the head - GEMM without its all-gather and fp32 cast), for a step whose acceptance runs as the sharded - ``trtllm::k3_spec_accept`` (see ``_k3_logits_shard``); the gathered fp32 logits otherwise. At such a step the - acceptance is their only reader: the one-model spec sampler takes the worker's tokens and rejects requests for - logits or log probabilities, so the shard is also what the step returns as ``"logits"``.""" - self._k3_step_shard = None - shard = self._k3_logits_shard( - hidden_states, lm_head, attn_metadata, spec_metadata, draft_model - ) - if shard is None: - return super().target_logits( - hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model - ) - logits = self._k3_head_shard(logits_processor, lm_head, hidden_states) - self._k3_step_shard = (logits, shard) - return logits - - @staticmethod - def _k3_head_shard(logits_processor, lm_head, rows: torch.Tensor) -> torch.Tensor: - """This rank's bf16 vocabulary shard of ``lm_head(rows)``, without the head's all-gather: from the logits - processor's own head kernel where it has one that takes the rows (``lm_head_shard``; the Kimi K3 target's - ``gemm/k3_head_gemv``), so the shard holds the values of the processor's gathered logits; else from the head's - own GEMM.""" - head_shard = getattr(logits_processor, "lm_head_shard", None) - logits = None if head_shard is None else head_shard(rows, lm_head) - if logits is None: - logits = lm_head.apply_linear(rows, lm_head.bias) - return logits - - def _k3_logits_shard(self, hidden_states, lm_head, attn_metadata, spec_metadata, draft_model): - """``(workspace, first column)`` when this step's target logits stay vocabulary-sharded, else None. - - That needs a step ``_k3_accept_step_applies`` takes, plain TP whose all-reduces own an MNNVL workspace for - this mapping, and an unquantized, bias-free, unpadded, evenly split column-parallel bf16 lm_head with shards - the kernel splits. The exchange is collective, so every rank must decide alike: every condition is the - batch's or the configuration's. The workspace is allocated (collectively) and the kernel compiled outside - CUDA-graph capture only, so a graph whose warmup allocated no workspace is captured on the gathered path. - """ - if not self.k3_decode: - return None - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce - - mapping = self.mapping - if getattr(self, "_k3_shard_head", None) is not lm_head: - reason = self._k3_logits_shard_reason(lm_head) - if reason is None and MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None: - # Not kept: re-checked on later steps, since the model's MNNVL workspace may not exist yet. - logger.info_once( - "DFlash: the target logits are all-gathered (the TP all-reduces are not MNNVL)", - key="dflash_k3_logits_no_mnnvl", - ) - return None - if mapping is not None and mapping.tp_size > 1: - if reason is None: - logger.info( - f"DFlash: decode-step target logits stay vocabulary-sharded ({lm_head.weight.shape[0]} " - f"columns per rank); trtllm::k3_spec_accept exchanges the row maxima, and a request's " - f"logits post-processors get them gathered" - ) - else: - logger.info_once( - f"DFlash: the target logits are all-gathered ({reason})", key=reason - ) - self._k3_shard_head, self._k3_shard_ok = lm_head, reason is None - if not self._k3_shard_ok: - return None - num_contexts = attn_metadata.num_contexts - num_gens = attn_metadata.num_seqs - num_contexts - K = spec_metadata.runtime_draft_len - if hidden_states.shape[0] != num_gens * (K + 1) or not self._k3_accept_step_applies( - attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ): - return None - embed = draft_model.draft_model_full.model.embed_tokens - columns = lm_head.weight.shape[0] - if not accept_op.supports( - columns, num_gens, self._compute_block_size, K, embed.weight.shape[1], mapping.tp_size - ): - return None - if torch.cuda.is_current_stream_capturing(): - ws = accept_op.existing_workspace(mapping) - if ws is None: - return None - else: - ws = accept_op.workspace(mapping) - return ws, mapping.tp_rank * columns - - def _k3_logits_shard_reason(self, lm_head): - """Why the target logits cannot stay vocabulary-sharded under this configuration and head (None if they can, - given an MNNVL workspace for the mapping, which ``_k3_logits_shard`` checks).""" - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - from tensorrt_llm._torch.modules.linear import TensorParallelMode - - mapping = self.mapping - if mapping is None or mapping.tp_size <= 1: - return "no TP" - if mapping.enable_attention_dp or mapping.has_cp() or mapping.has_pp(): - return "attention DP, CP or PP" - if ( - getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN - or not getattr(lm_head, "gather_output", False) - or getattr(lm_head, "padding_size", 0) != 0 - or getattr(lm_head, "gather_output_sizes", None) is not None - ): - return "the lm_head is not an evenly split, unpadded column-parallel head" - if ( - lm_head.bias is not None - or lm_head.has_any_quant - or lm_head.weight.dtype != torch.bfloat16 - ): - return "the lm_head has a bias, is quantized or is not bf16" - if not accept_op.supports_columns(lm_head.weight.shape[0], mapping.tp_size): - columns = lm_head.weight.shape[0] - return f"trtllm::k3_spec_accept cannot split {columns} columns over {mapping.tp_size} ranks" - return None - - def _k3_static_arange(self, lo: int, hi: int) -> torch.Tensor: - """``arange(lo, hi)`` (int64), built once outside CUDA-graph capture.""" - cache = self.__dict__.setdefault("_k3_aranges", {}) - t = cache.get((lo, hi)) - if t is None: - t = torch.arange(lo, hi, dtype=torch.long, device="cuda") - if not torch.cuda.is_current_stream_capturing(): - cache[(lo, hi)] = t - return t - - def _k3_mask_row(self, draft_model) -> torch.Tensor: - """The noise block's mask embedding, bf16 [hidden], cached: the drafter's own trained mask row if it kept one - for the mask id in use, else the embedding's row (``dflash_noise_block_embedding`` reads the same row).""" - row = getattr(self, "_k3_mask_embedding", None) - if row is None: - embed = draft_model.draft_model_full.model.embed_tokens - mask_token_id = self._mask_token_id(draft_model) - trained = self._trained_mask_embedding(draft_model, mask_token_id) - if trained is not None: - row = trained.to(embed.weight.dtype).reshape(-1).contiguous() - else: - row = embed.weight[mask_token_id].contiguous() - if torch.cuda.is_current_stream_capturing(): - return row - self._k3_mask_embedding = row - return row - - def _k3_ctx_pool_key(self): - """The identity of the pool the ctx cache is bound to now (KV cache estimation rebinds it to a new one).""" - return ( - tuple(t.data_ptr() for t in self._ctx_kv_buf) if self._ctx_kv_buf is not None else None - ) - - def _k3_ctx_kv_applies(self, draft_model, projected: torch.Tensor) -> bool: - """Whether ``trtllm::k3_ctx_kv`` writes this step's context K/V: the manager-bound paged pool with K and V, - bf16, the fused K/V weight without bias or context input norm, k_norm, NeoX RoPE from flashinfer's fp32 - cache, a context the op can split (up to 64 tokens: B <= 8 requests of K + 1 <= 8, ``ctx_op.pick_split``), one - allocation for every layer's pool. Checked once per token count and - bound pool; the pool view is kept for the bound pool only, so a rebind never leaves the kernel writing the - replaced pool.""" - pool_key = self._k3_ctx_pool_key() - if getattr(self, "_k3_ctx_pool_bound", None) != pool_key: - self._k3_ctx_pool_bound = pool_key - self._k3_ctx_pool = None - self._k3_ctx_kv_ok = {} - cache = self._k3_ctx_kv_ok - n = projected.shape[0] - ok = cache.get(n) - if ok is not None: - return ok - from tensorrt_llm._torch.cute_dsl_kernels.k3_ctx_kv import op as ctx_op - from tensorrt_llm._torch.models import modeling_dflash - - reason = None - if not (self._ctx_paged and self._ctx_block_tables is not None): - reason = "the context pool is not the manager-bound paged pool" - elif projected.dtype != torch.bfloat16 or self._ctx_kv_buf[0].dtype != torch.bfloat16: - reason = "not bf16" - else: - if draft_model._fused_kv_weight is None: - draft_model._build_fused_kv_buffers() - view = ctx_op.pool_view(list(self._ctx_kv_buf)) - if ( - draft_model._fused_kv_bias is not None - or getattr(draft_model, "_input_ln_eps", None) is not None - ): - reason = "K/V bias or context input norm" - elif draft_model._k_norm_stacked is None or not draft_model._is_neox: - reason = "no k_norm or not NeoX RoPE" - elif ( - modeling_dflash._flashinfer_rope is None - or draft_model._get_cos_sin_cache().shape[1] != 64 - ): - reason = "RoPE is not flashinfer's full-rotary fp32 cache" - elif view is None or self._ctx_kv_buf[0].size(1) != 2: - reason = "the layers' pools are not K/V views of one allocation" - elif not ctx_op.pick_split( - draft_model._fused_kv_weight.shape[0], - projected.shape[1], - draft_model._num_kv_heads, - self._draft_tokens_per_req, - n, - projected.device, - ): - reason = f"{n} context tokens (at most 64) or the K/V weight shape" - else: - self._k3_ctx_pool = view - ok = reason is None - if not ok: - logger.info(f"DFlash: context K/V on the Python path ({reason})") - if not torch.cuda.is_current_stream_capturing(): - cache[n] = ok - return ok - - def _k3_ctx_kv( - self, draft_model, projected, ctx_positions, num_accepted, slots, rows - ) -> torch.Tensor: - """``trtllm::k3_ctx_kv``: every drafter layer's context K/V of this step into the paged pool and - ``_ctx_len += num_accepted`` (clamped); returns the context length each request may advertise.""" - if self._k3_ctx_pool is None or self._k3_ctx_pool_bound != self._k3_ctx_pool_key(): - raise RuntimeError("k3_ctx_kv: the context pool was rebound after the path was chosen") - flat, layer_off, page_stride, kv_stride, head_stride = self._k3_ctx_pool - return torch.ops.trtllm.k3_ctx_kv( - projected.contiguous(), - draft_model._fused_kv_weight, - draft_model._k_norm_stacked, - draft_model._get_cos_sin_cache(), - ctx_positions, - num_accepted, - self._ctx_len, - slots, - rows, - self._ctx_block_tables, - self._ctx_block_counts, - flat, - layer_off, - page_stride, - kv_stride, - head_stride, - draft_model._k_norm_eps, - self._max_ctx, - self._ctx_page_size, - self._compute_block_size, - draft_model._num_kv_heads, - ) - - def _on_acceptance( - self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata - ) -> None: - """Called with this step's acceptance when ``k3_spec_accept`` produced it (a hook for drafter families).""" - - def _k3_accept(self, logits, attn_metadata, spec_metadata, draft_model, shard=None): - """``trtllm::k3_spec_accept`` for a step ``_k3_accept_applies`` accepted: the acceptance, the block table, - the KDA replay record, kv_lens + 1 and the drafter's inputs in one launch. ``shard``: the exchange - workspace and first column when ``logits`` are this rank's vocabulary shard (``_k3_logits_shard``). - Returns (accepted tokens, num accepted, {rewind, bonus, qpos, cpos, noise}).""" - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - - num_gens = attn_metadata.num_seqs - K = spec_metadata.runtime_draft_len - mgr = attn_metadata.kv_cache_manager - is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() - force = 0.0 if is_warmup else float(self.force_num_accepted_tokens) - mode, _, _ = accept_op.force_mode(force, K) - if mode == accept_op._kernel_module().FORCE_FRAC: - self._ensure_force_accept_rng_state(logits.device) - pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter - else: - dummy_rng = getattr(self, "_k3_dummy_rng", None) - if dummy_rng is None: - dummy_rng = self._k3_dummy_rng = ( - torch.zeros(4, dtype=torch.float32, device=logits.device), - torch.zeros(2, dtype=torch.int64, device=logits.device), - ) - pool, counter = dummy_rng - embed = draft_model.draft_model_full.model.embed_tokens - accepted, num_acc, rewind, bonus, qpos, cpos, noise = torch.ops.trtllm.k3_spec_accept( - logits, - spec_metadata.draft_tokens.reshape(num_gens, K), - attn_metadata.draft_kv_cache_block_offsets, - self._ctx_pool_idx, - self._ctx_block_divisor, - self._ctx_block_counts, - self._ctx_block_tables, - mgr.prev_num_accepted_tokens, - attn_metadata.mamba_metadata.state_indices, - mgr._dummy_request_mask, - attn_metadata.kv_lens_cuda, - self._batch_to_slot, - self._ctx_len, - self._max_ctx, - embed.weight, - self._k3_mask_row(draft_model), - pool, - counter, - force, - self._compute_block_size, - *self._k3_shard_args(shard), - ) - self._on_acceptance(accepted, num_acc, attn_metadata, spec_metadata) - return ( - accepted, - num_acc, - dict(rewind=rewind, bonus=bonus, qpos=qpos, cpos=cpos, noise=noise), - ) - - @staticmethod - def _k3_shard_args(shard) -> tuple: - """``trtllm::k3_spec_accept``'s exchange arguments for a vocabulary-sharded step (none otherwise).""" - if shard is None: - return () - ws, first_column = shard - return ( - ws["uc"], - ws["mc"], - ws["flags"], - ws["rank"], - ws["slots"], - ws["push_copies"], - first_column, - ) - - def _forward_impl( - self, - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - resource_manager=None, - ): - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - # Set when target_logits kept this step's logits vocabulary-sharded (only k3_spec_accept reads them). - step_shard, self._k3_step_shard = getattr(self, "_k3_step_shard", None), None - shard = None - if step_shard is not None: - if step_shard[0] is not logits: - raise RuntimeError( - "DFlash: the target logits are not the vocabulary shard target_logits returned" - ) - shard = step_shard[1] - - raw_logits = logits - K = spec_metadata.runtime_draft_len - - if K == 0: - return self.skip_drafting( - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - ) - - # Lazy init buffers and attach worker reference for prepare() - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - self._lazy_init_ctx_buffers( - draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager - ) - spec_metadata._dflash_worker = self - # A decode step of the supported mode runs the acceptance and the drafter's inputs as one kernel - # (trtllm::k3_spec_accept), which also decodes the block table below. - if shard is None: - fused_accept = self._k3_accept_applies( - logits, attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ) - elif self._k3_accept_step_applies( - attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ): - fused_accept = True - else: - raise RuntimeError( - "DFlash: the target logits are vocabulary-sharded, but the fused acceptance no longer applies" - ) - # Before any store: prefill and decode both address pages through it. - # Returning False here means an empty batch -- the missing-offsets case - # raises inside, with the metadata type in the message. - if not fused_accept: - self._refresh_ctx_block_tables(attn_metadata, batch_size) - - # Save context lengths so both warmup and a failed forward can roll - # back the in-place _ctx_len updates made during drafting. - is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() - if not torch.cuda.is_current_stream_capturing(): - # Never allocate the snapshot while capturing a CUDA graph: the - # clone would live in the graph memory pool, and its replay-time - # writes could alias blocks reused by later captures. Rollback is - # only meaningful for eager/warmup forwards anyway; a failure - # during capture aborts the graph itself, and captured ops do not - # mutate _ctx_len until replay. - self._saved_ctx_len = self._ctx_len.clone() - self._saved_ctx_len_host = list(self._ctx_len_host) - self._saved_req_ctx_pos = dict(self._req_ctx_pos) - self._ctx_len_restore_pending = True - - self._execute_guided_decoder_if_present(logits) - - prep = None - if fused_accept: - # Saved before the kernel updates kv_lens (warmup also snapshots kv_lens_cuda here). - self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) - accepted_tokens, num_accepted_tokens, prep = self._k3_accept( - logits, attn_metadata, spec_metadata, draft_model, shard - ) - else: - accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( - logits, attn_metadata, spec_metadata - ) - - # Opt-in acceptance recording (env-gated; eager-mode measurement - # runs only). Skipped for CUDA-graph batches (capture/replay/warmup - # use synthetic requests and forbid the host sync). - if ( - self._accept_stats is not None - and num_gens > 0 - and not spec_metadata.is_cuda_graph - and not torch.cuda.is_current_stream_capturing() - and spec_metadata.request_ids is not None - ): - self._accept_stats.on_accept( - spec_metadata.request_ids[num_contexts:batch_size], - num_accepted_tokens[num_contexts:batch_size].tolist(), - ) - - if fused_accept: - # The kernel recorded the KDA replay acceptance and advanced kv_lens_cuda. - self._kv_rewind_amount = prep["rewind"] - self._kv_rewind_nc = num_contexts - self._kv_rewind_bs = batch_size - self._kv_rewind_pending = True - attn_metadata.update_for_spec_dec() - else: - # Update GDN/Mamba recurrent states to the accepted token's state. - if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): - attn_metadata.kv_cache_manager.update_mamba_states( - attn_metadata=attn_metadata, - num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices, - ) - - self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) - self._prepare_kv_for_draft_forward( - attn_metadata, num_accepted_tokens, num_contexts, batch_size - ) - - # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. - # The draft model uses standard 1D RoPE, so only scalar positions are needed. - if position_ids.ndim == 3: - position_ids = position_ids[0, 0] - else: - position_ids = position_ids.squeeze(0) - - # Get total tokens processed by target model (for hidden state extraction) - total_target_tokens = input_ids.shape[0] - - # Capture prefill (context) hidden states for future gen steps. - # This gives the draft model the full prompt context, not just gen tokens. - if num_contexts > 0: - self._store_prefill_context( - draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens - ) - # Rebuild batch_to_slot after prefill assigns new slots - if self._ctx_buf_inited and spec_metadata.request_ids: - self._rebuild_batch_to_slot(spec_metadata.request_ids) - - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - total_target_tokens=total_target_tokens, - prep=prep, - ) - - if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states_out = draft_model.dflash_forward( - noise_embedding=inputs["noise_embedding"], - query_positions=inputs["query_positions"], - num_ctx_per_req=inputs["num_ctx_per_req"], - ctx_k_cache=inputs["ctx_k_cache"], - ctx_v_cache=inputs["ctx_v_cache"], - ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], - ctx_kv_cache=inputs["ctx_kv_cache"], - ctx_page_table=inputs["ctx_page_table"], - **dflash_ctx_rows_kwargs(draft_model, inputs["ctx_rows_start"]), - ) - - # Gather K logits per gen request from the block outputs. - # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. - # Which block slots carry them is a drafter-family convention, - # resolved through _draft_slot_ids. - block_size = self._compute_block_size - gen_hidden_states = self._draft_block_hidden_states( - draft_model, hidden_states_out, num_gens, block_size, K - ) - gen_logits = self._draft_block_logits( - draft_model, gen_hidden_states, attn_metadata, spec_metadata - ) - - vocab_size = gen_logits.shape[-1] - gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - - gen_logits = self._refine_block_logits( - draft_model, gen_logits, inputs, spec_metadata - ) - - # DFlash 2: replace the independent per-position picks with one - # coherent path through the block (absent for plain DFlash). - if getattr(draft_model, "has_candidate_selector", False): - gen_logits = self._apply_dflash2_selector( - draft_model, - gen_logits, - gen_hidden_states.reshape(num_gens, K, -1), - inputs["first_prev_tokens"], - spec_metadata, - ) - vocab_size = gen_logits.shape[-1] - - gen_draft_tokens = self.sample_draft_tokens( - gen_logits, - spec_metadata, - batch_size, - num_contexts=num_contexts, - ) - - # Opt-in confidence-calibration recording (env-gated). The - # confidence values come from an optional provider on the - # recorder (None here -> no calibration rows collected; the - # DSpark confidence-scheduled verification MR supplies the - # real provider, see accept_stats.py). Guarded off for - # CUDA-graph batches and d2t vocab-mapped drafters. - if ( - self._accept_stats is not None - and self._accept_stats.confidence_provider is not None - and not spec_metadata.is_cuda_graph - and not torch.cuda.is_current_stream_capturing() - and spec_metadata.request_ids is not None - and self._d2t is None - ): - self._accept_stats.record_draft_confidence( - spec_metadata.request_ids[num_contexts:batch_size], - draft_model, - gen_hidden_states.reshape(num_gens, K, -1), - inputs["first_prev_tokens"], - gen_draft_tokens, - ) - - else: - gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") - - # Context requests are not drafted by the block worker (zero placeholder - # token); fill their draft-prob slot rows with a one-hot placeholder so - # they are a legal distribution when they become gen requests next iter. - gen_vocab = vocab_size if num_gens > 0 else None - self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) - - if num_contexts > 0 and num_gens > 0: - ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) - elif num_contexts > 0: - next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - else: - next_draft_tokens = gen_draft_tokens - - self._restore_attn_metadata_from_spec_dec(attn_metadata) - self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) - - self._rollback_guided_decoder_after_verify(num_accepted_tokens) - - next_new_tokens = self._prepare_next_new_tokens( - accepted_tokens, - next_draft_tokens, - spec_metadata.batch_indices_cuda, - batch_size, - num_accepted_tokens, - ) - - # Restore context lengths after warmup; real runs keep the updates. - if is_warmup: - self._ctx_len.copy_(self._saved_ctx_len) - self._restore_ctx_len_host() - self._ctx_len_restore_pending = False - - outputs = { - "logits": raw_logits, - "new_tokens": accepted_tokens, - "new_tokens_lens": num_accepted_tokens, - "next_draft_tokens": next_draft_tokens, - "next_new_tokens": next_new_tokens, - } - if shard is not None: - # "logits" is this rank's vocabulary shard: the engine gathers it for any logits post-processor. - outputs["logits_vocab_shard"] = True - return outputs - - def _draft_block_hidden_states( - self, - draft_model, - hidden_states_out: torch.Tensor, - num_gens: int, - block_size: int, - num_draft_tokens: int, - ) -> torch.Tensor: - """The block-output rows that produce the K draft logits per gen request (``_draft_slot_ids``). - - The slot ids depend only on the shapes, so they are built once per shape, outside CUDA-graph - capture, and cached. When they are one contiguous run of rows (a single request, or the - shift_label convention with a block of K slots, whose requests' slots follow each other) the - rows are returned as a view, without a gather; otherwise they are gathered. - """ - rows = hidden_states_out.shape[0] - key = (num_gens, block_size, num_draft_tokens, rows) - cache = self.__dict__.setdefault("_draft_block_rows", {}) - entry = cache.get(key) - if entry is None: - ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) - # Shields only the last request: at block_size == K with - # shift_label off, slots run 1..K, so every request reads the - # next one's slot 0 and the last overruns. Degrades, never raises. - ids = ids.clamp(max=rows - 1) - if torch.cuda.is_current_stream_capturing(): - return hidden_states_out[ids] - host = ids.tolist() - if host and host == list(range(host[0], host[0] + len(host))): - entry = host[0] - else: - entry = ids - cache[key] = entry - if isinstance(entry, int): - return hidden_states_out[entry : entry + num_gens * num_draft_tokens] - return hidden_states_out[entry] - - def _draft_block_logits( - self, - draft_model, - gen_hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - spec_metadata, - ) -> torch.Tensor: - """fp32 logits of the block positions over the full vocabulary (the TP lm_head all-gathers its - shards). A drafter family whose refinement and draft sampler reduce across ranks themselves may - return this rank's vocabulary shard instead.""" - return draft_model.logits_processor( - gen_hidden_states, draft_model.lm_head, attn_metadata, True - ) - - def _ctx_rows_start(self, num_contexts: int, num_gens: int) -> Optional[int]: - """The first row of the gen requests' page-table rows where those rows are one contiguous run, else None. - - The manager's block table is keyed by batch position, so the gen requests' rows are [num_contexts, - num_contexts + num_gens): a drafter can view them instead of gathering ``ctx_cache_batch_idx``'s rows. The - private arena's table is keyed by slot, which need not be contiguous. - """ - if num_gens > 0 and self._ctx_block_tables is not None: - return num_contexts - return None - - def prepare_1st_drafter_inputs( - self, - input_ids: torch.LongTensor, - position_ids: torch.LongTensor, - hidden_states: torch.Tensor, - accepted_tokens: torch.Tensor, - num_accepted_tokens: torch.Tensor, - attn_metadata: AttentionMetadata, - spec_metadata: DFlashSpecMetadata, - draft_model: nn.Module, - total_target_tokens: int = 0, - prep: Optional[dict] = None, - ): - """Prepare inputs for DFlash's draft forward. - - ``prep``: ``trtllm::k3_spec_accept``'s bonus tokens, positions and noise embedding for a decode step it - took (see ``_k3_accept``); computed here otherwise. - - For gen requests, builds: - - noise_embedding: token embeddings for [accepted + mask] tokens - - query_positions: position IDs for the query tokens - - num_ctx_per_req: per-request context length in the pool - - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed - views of the persistent per-layer K/V pool. - - ctx_rows_start: where ctx_cache_batch_idx is the contiguous rows - [ctx_rows_start, ctx_rows_start + num_gens) of ctx_page_table, their - start (``_ctx_rows_start``), else None. - """ - num_contexts = attn_metadata.num_contexts - batch_size = attn_metadata.num_seqs - num_gens = batch_size - num_contexts - - mask_token_id = self._mask_token_id(draft_model) - - # Get the embed_tokens layer from the draft model - embed_tokens = draft_model.draft_model_full.model.embed_tokens - hidden_dim = ( - spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] - ) - - if num_gens > 0: - gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] - gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] - - total_tokens_per_req = self._draft_tokens_per_req # K+1 - K = spec_metadata.runtime_draft_len - - # Get captured multi-layer hidden states from spec_metadata - captured_hs = spec_metadata.get_hidden_states(total_target_tokens) - has_target_features = ( - captured_hs is not None - and hasattr(draft_model, "fc") - and hasattr(draft_model, "hidden_norm") - ) - - # Use cached block_size (resolved once on first call) - block_size = self._compute_block_size - query_tokens_per_req = block_size - - # Get slots for gen requests from pre-computed mapping - slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] - K_plus_1 = K + 1 - - if prep is not None: - gen_rows_out = self._k3_static_arange(num_contexts, num_contexts + num_gens) - offsets_kp1 = self._k3_static_arange(0, K_plus_1) - bonus = prep["bonus"] - query_position_ids = prep["qpos"] - ctx_position_ids = prep["cpos"] - noise_embed_2d = prep["noise"] - else: - gen_rows_out = torch.arange( - num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" - ) - - bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) - bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() - - ctx_len_gen = self._ctx_len[slots] - j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") - offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - - # _ctx_len is clamped to _max_ctx only AFTER this step's accepted - # tokens are folded in (see the update below), so the running length - # used here has to be clamped on its own -- otherwise a request that - # already sits at the ceiling indexes num_accepted positions past - # any position the sequence can legitimately reach. - ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) - query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) - ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) - - noise_embed_2d = dflash_noise_block_embedding( - embed_tokens, - bonus, - mask_token_id, - query_tokens_per_req, - self._trained_mask_embedding(draft_model, mask_token_id), - ) - - # Accumulate new accepted features into context buffers - fused_ctx_kv = False - if has_target_features: - gen_start = attn_metadata.num_ctx_tokens - # Target now processes exactly K+1 tokens per gen req, so the - # captured slice is already the full set we need to project. - gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] - gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) - projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) - fused_ctx_kv = prep is not None and self._k3_ctx_kv_applies( - draft_model, projected_to_store - ) - if fused_ctx_kv: - num_ctx_fused = self._k3_ctx_kv( - draft_model, - projected_to_store, - ctx_position_ids, - gen_num_accepted, - slots, - gen_rows_out, - ) - # k3_ctx_kv adds the accepted tokens in the kernel; empty the dummy slot after - # it, as _advance_ctx_len does after the add on the torch path. - self._ctx_len[self._dummy_slot].zero_() - if has_target_features and not fused_ctx_kv: - gen_num_accepted_long = gen_num_accepted.long() - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) - write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - if self._ctx_block_tables is not None: - # Bound by what the manager handed THIS request: the masked - # entries below are still written (fixed-size, for graph - # safety) and a placeholder page resolves to another's. - gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size - col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - else: - col_idx = col_idx.clamp(max=self._max_ctx - 1) - - # Fixed-size writes for CUDA graph compatibility: - # Write ALL entries but zero out invalid ones. Invalid - # writes land at clamped column indices beyond valid - # _ctx_len range, so they're harmless. - slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) - col_flat = col_idx.reshape(-1) - proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] - pos_flat = ctx_position_ids.long().reshape(-1) - mask_1d = write_mask.reshape(-1) - - # Fast path: store the pre-projected/pre-RoPE'd K/V. - # dflash_forward reads these directly via cache_batch_idx. - cache_dtype = ( - self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype - ) - k_new, v_new = draft_model.precompute_context_kv( - proj_flat.to(cache_dtype), pos_flat - ) - mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) - k_new.mul_(mask_bc) - if v_new is not None: - v_new.mul_(mask_bc) - slot_long = slot_flat.long() - col_long = col_flat.long() - if self._ctx_paged: - if self._ctx_block_tables is not None: - # Batch positions of the gen requests, matching the - # per-request block table's row order. - rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) - else: - rows_long = slot_long - self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - else: # VANILLA DFlash backend (FlashAttention) - self._ctx_k_buf[slot_long, :, col_long] = k_new - if v_new is not None: - self._ctx_v_buf[slot_long, :, col_long] = v_new - - self._advance_ctx_len(slots, gen_num_accepted_long) - - if fused_ctx_kv: - num_ctx_per_req_t = num_ctx_fused - else: - num_ctx_per_req_t = self._ctx_len[slots] - if not fused_ctx_kv and self._ctx_block_tables is not None: - # The write above clamps columns to the request's allocation, so - # a context that outruns it has its tail written over the last - # valid slot. Truncate what is advertised to match, or the read - # attends that clobbered value -- or another request's page. - allocated = dflash_allocated_ctx_limit( - self._ctx_block_counts, self._ctx_page_size, block_size - ) - num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) - noise_embedding = noise_embed_2d - query_positions = query_position_ids.long() - - # Update seq_lens for gen requests to K+1 (the number of tokens - # the target forward actually processed). - attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( - total_tokens_per_req - ) - attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req - else: - noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) - query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") - num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") - slots = torch.empty(0, dtype=torch.long, device="cuda") - gen_rows_out = slots - bonus = torch.empty(0, dtype=torch.long, device="cuda") - - return { - "noise_embedding": noise_embedding, - "query_positions": query_positions, - "num_ctx_per_req": num_ctx_per_req_t, - "ctx_k_cache": self._ctx_k_buf, - "ctx_v_cache": self._ctx_v_buf, - # Slots index the private arena; the manager's block table is keyed - # by batch position, so the drafter reads its own rows there. - "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, - # Anchor token per gen request (block slot 0): last accepted - # token. The dspark Markov chain conditions its first step on it. - "first_prev_tokens": bonus, - "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": ( - self._ctx_block_tables - if self._ctx_block_tables is not None - else self._ctx_page_table - ), - "ctx_rows_start": self._ctx_rows_start(num_contexts, num_gens), - } - - -class KimiK3DSparkWorker(KimiK3DFlashWorker, DSparkWorker): - # Whether this step's block logits are the Kimi K3 decode path's vocabulary shard (see ``_draft_block_logits``). - _k3_sharded_block_logits = False - - def _apply_dspark_markov_bias( - self, - draft_model, - gen_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - spec_metadata, - ) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): - step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head - logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} - is the greedy token from step i-1's biased logits. Greedy per-position - argmax of the returned logits therefore reproduces the reference - sampled chain; the rejection-sampling path samples from the same - biased distributions (proposal conditioned on the greedy chain). - - Block logits the Kimi K3 decode path kept vocab-sharded (see ``_draft_block_logits``) run the - whole chain as ``trtllm::k3_markov``, which also returns the greedy draft tokens and - next_new_tokens. Other TP-sharded logits slice markov_w2's rows to this rank's contiguous - shard and chain through the TP-aware global argmax. - """ - # The d2t guard lives in set_draft_model: it is model-static, so raising - # it here would surface a load-time config error per decode step. - # Unlike the d2t guard this one cannot move to set_draft_model: it - # keys on the runtime logits width, and reproducing that at init would - # duplicate the draft head's sharding rules. A standalone drafter - # borrows the target lm_head, whose gather_output defaults to True, so - # the logits normally arrive full-vocab and this branch is skipped. - self._k3_markov = None - k3_sharded, self._k3_sharded_block_logits = self._k3_sharded_block_logits, False - full_vocab = draft_model.markov_w2.shape[0] - shard = gen_logits.shape[-1] - vocab_slice = None - if shard != full_vocab: - mapping = self.mapping - if ( - mapping is None - or getattr(mapping, "enable_attention_dp", False) - or shard * mapping.tp_size != full_vocab - ): - raise NotImplementedError( - f"DSpark Markov head: draft logits width {shard} does not " - f"match the drafter vocab {full_vocab} and is not a plain " - "TP column shard of it." - ) - vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) - if k3_sharded: - return self._k3_markov_chain( - draft_model, gen_logits, first_prev_tokens, vocab_slice - ) - - def argmax_fn(step_logits): - # Full-vocab token ids (TP-aware when sharded); tokens stay in - # draft-vocab space, which is what markov_w1 indexes. - return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() - - return draft_model.apply_markov_chain_logits( - gen_logits, - first_prev_tokens, - argmax_fn=argmax_fn, - vocab_slice=vocab_slice, - ) - - def _keep_draft_logits_sharded(self, draft_model, spec_metadata, num_gens: int) -> bool: - """Keep the draft logits vocab-sharded: ``trtllm::k3_markov`` reduces across the ranks itself. - - The chain's global argmaxes are the greedy draft tokens, so the lm_head's all-gather of the - full logits and the draft sampler's gather are both skipped. Needs a Kimi K3 target running its - decode kernels (``k3_decode``), plain TP, greedy drafting (rejection sampling reads - full-vocabulary probabilities), an unquantized bias-free column-parallel bf16 lm_head whose - shards tile the Markov vocabulary, bf16 Markov weights of the kernel's rank, MNNVL (the model's - all-reduces own an MNNVL workspace for this mapping), and a block, shard and batch of - ``num_gens`` requests the kernel can split. - """ - if not self.k3_decode: - return False - from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op - from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce - from tensorrt_llm._torch.modules.linear import TensorParallelMode - - mapping = self.mapping - lm_head = getattr(draft_model, "lm_head", None) - if ( - lm_head is None - or not getattr(draft_model, "has_markov_head", False) - or mapping is None - or mapping.tp_size <= 1 - or mapping.enable_attention_dp - or spec_metadata.wants_advanced_draft_sampling - or getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN - or not getattr(lm_head, "gather_output", False) - or getattr(lm_head, "bias", None) is not None - or lm_head.weight.dtype != torch.bfloat16 - or draft_model.markov_w1.dtype != torch.bfloat16 - or draft_model.markov_w2.dtype != torch.bfloat16 - or lm_head.weight.shape[0] * mapping.tp_size != draft_model.markov_w2.shape[0] - or MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None - ): - return False - shard = lm_head.weight.shape[0] - block = spec_metadata.runtime_draft_len - rank = k3_markov_op._kernel_module().MARKOV_RANK - if ( - not 0 < block <= k3_markov_op.WORKSPACE_MAX_BLOCK - or draft_model.markov_w1.dim() != 2 - or draft_model.markov_w1.shape[1] != rank - or tuple(draft_model.markov_w2.shape[1:]) != (rank,) - ): - return False - if k3_markov_op.pick_grid(shard, block, num_gens) == 0: - logger.warning_once( - f"DSpark Markov head: trtllm::k3_markov cannot split a {shard}-row vocab shard for " - f"{num_gens} requests; the draft logits are all-gathered and the chain runs unfused.", - key=f"dspark_k3_markov_shard_{num_gens}", - ) - return False - return True - - def _draft_block_logits( - self, - draft_model, - gen_hidden_states: torch.Tensor, - attn_metadata, - spec_metadata, - ) -> torch.Tensor: - """This rank's bf16 shard of the block logits when ``k3_markov`` takes them (it converts them to fp32 - exactly, as ``.float()`` would), from the drafter's logits processor's head kernel where it takes the rows, - else the head's own GEMM (``_k3_head_shard``); otherwise the base class' fp32 logits.""" - num_gens = gen_hidden_states.shape[0] // max(spec_metadata.runtime_draft_len, 1) - self._k3_sharded_block_logits = self._keep_draft_logits_sharded( - draft_model, spec_metadata, num_gens - ) - if self._k3_sharded_block_logits: - return self._k3_head_shard( - draft_model.logits_processor, draft_model.lm_head, gen_hidden_states - ) - return super()._draft_block_logits( - draft_model, gen_hidden_states, attn_metadata, spec_metadata - ) - - def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): - """The base acceptance; under ``k3_decode`` its outputs are kept for ``k3_markov``'s next_new_tokens.""" - accepted_tokens, num_accepted_tokens = super().sample_and_accept_draft_tokens( - logits, attn_metadata, spec_metadata - ) - if self.k3_decode: - self._on_acceptance(accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata) - return accepted_tokens, num_accepted_tokens - - def _on_acceptance( - self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata - ) -> None: - """Keeps this step's acceptance for ``k3_markov``'s next_new_tokens (and the metadata whose KV lengths it - rewinds).""" - self._k3_acceptance = ( - accepted_tokens, num_accepted_tokens, attn_metadata.num_contexts, spec_metadata, attn_metadata - ) # fmt: skip - - def _k3_markov_chain( - self, - draft_model, - gen_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - vocab_slice: slice, - ) -> torch.Tensor: - """The Markov chain as ``trtllm::k3_markov`` on this rank's shard of the block logits. - - Returns the corrected logits (fp32); keeps the kernel's greedy tokens for - ``sample_draft_tokens`` and its next_new_tokens (built from this step's acceptance) for - ``_prepare_next_new_tokens``. - """ - from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op - - acceptance = getattr(self, "_k3_acceptance", None) - if acceptance is None: - raise RuntimeError("DSpark Markov head: k3_markov needs this step's acceptance first") - accepted_tokens, num_accepted_tokens, num_contexts, spec_metadata, attn_metadata = ( - acceptance - ) - num_gens = gen_logits.shape[0] - # The draft forward's pending KV-length rewind (see _apply_kv_rewind_after_draft) runs in the kernel: every - # reader of kv_lens_cuda in the draft forward precedes it. In warmup the restore of kv_lens_cuda that - # follows overwrites it, as it would the Python rewind. - rewind = getattr(self, "_kv_rewind_amount", None) - fold = ( - getattr(self, "_kv_rewind_pending", False) - and rewind is not None - and getattr(attn_metadata, "kv_lens_cuda", None) is not None - and rewind.dtype == torch.int32 - and self._kv_rewind_bs - self._kv_rewind_nc == rewind.numel() <= num_gens - ) - corrected, tokens, next_new = k3_markov_op.markov_chain( - self.mapping, - gen_logits.contiguous(), - first_prev_tokens.long(), - draft_model.markov_w1, - draft_model.markov_w2[vocab_slice], - vocab_slice.start, - accepted_tokens, - num_accepted_tokens[num_contexts : num_contexts + num_gens], - spec_metadata.batch_indices_cuda[num_contexts : num_contexts + num_gens], - kv_lens=attn_metadata.kv_lens_cuda if fold else None, - rewind=rewind if fold else None, - rewind_first=self._kv_rewind_nc if fold else 0, - ) - if fold: - self._kv_rewind_amount = None - self._kv_rewind_pending = False - self._k3_markov = (corrected, tokens, next_new, accepted_tokens, num_accepted_tokens) - return corrected - - def sample_draft_tokens( - self, - logits, - spec_metadata, - batch_size, - *, - num_contexts=0, - draft_step=None, - mapping_lm_head_tp=None, - ): - """Greedy block drafts straight from ``k3_markov``. - - When ``logits`` are the corrected logits the kernel just returned, its per-position global - argmax (first maximum, lowest vocabulary index among equal values) is the token the - TP-gathered greedy sampler would pick. Anything else goes to the base sampler. - """ - chain = getattr(self, "_k3_markov", None) - if ( - chain is not None - and chain[0] is logits - and mapping_lm_head_tp is None - and not spec_metadata.wants_advanced_draft_sampling - ): - self._k3_markov_next = (chain[1], chain[2], chain[3], chain[4]) - return chain[1] - self._k3_markov_next = None - return super().sample_draft_tokens( - logits, - spec_metadata, - batch_size, - num_contexts=num_contexts, - draft_step=draft_step, - mapping_lm_head_tp=mapping_lm_head_tp, - ) - - def _prepare_next_new_tokens( - self, - accepted_tokens, - next_draft_tokens, - batch_indices_cuda, - batch_size, - num_accepted_tokens, - ): - """``k3_markov``'s next_new_tokens when the drafts are its tokens for the whole batch (no context - requests); otherwise the base assembly.""" - chain = getattr(self, "_k3_markov_next", None) - # The step ends here: release its acceptance (which holds the step's attention and spec metadata) and the - # kernel's outputs, so they do not keep a captured graph's pool or metadata alive after the step. - self._k3_markov_next = None - self._k3_acceptance = None - self._k3_markov = None - if ( - chain is not None - and chain[0] is next_draft_tokens - and chain[2] is accepted_tokens - and chain[3] is num_accepted_tokens - and next_draft_tokens.shape[0] == batch_size - ): - return chain[1] - return super()._prepare_next_new_tokens( - accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py index b4177f698d8b..24233d194709 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py @@ -1,42 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Weight loading: Kimi K3 (MXFP4) / sm_100 / tp16_moetp4ep4. +"""The shared Kimi K3 text checkpoint loader.""" -The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and -`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: - -* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the - routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under - moe_tp 4 x moe_ep 4), and checks its own key coverage. -* The vision tower and the projector are a predicted non-load: listed here, never read. -* Any other key fails the load, naming it, rather than being dropped. -""" - -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.models.modeling_utils import filter_weights - -_LANG_PREFIX = "language_model." - -# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. -PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") - - -def load(model, weights) -> None: - """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" - unknown = sorted( - k - for k in weights.keys() - if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) - ) - assert not unknown, ( - f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " - f"first {unknown[:5]}" - ) - lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) - assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" - checkpoint_dir = getattr(weights, "checkpoint_dir", None) - if checkpoint_dir is not None: - lm_weights.checkpoint_dir = checkpoint_dir - lm_weights.checkpoint_prefix = _LANG_PREFIX - KimiLinearForCausalLM.load_weights(model, lm_weights) +from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/__init__.py new file mode 100644 index 000000000000..b0abcb98c94f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MXFP4 / sm_103 / tp16_moetp16ep1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/modeling.py new file mode 100644 index 000000000000..6f5c9ccdb06b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/modeling.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MXFP4, SM103, attention TP16 and MoE TP16 x EP1.""" + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..modeling import K3TargetSpec, ModelingV2KimiK3 +from ..modeling import _check_construction as _check_target_construction + +_TARGET = K3TargetSpec(sm=(10, 3), quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, moe_tp=16, moe_ep=1) + + +def _check_construction(model_config: ModelConfig) -> None: + _check_target_construction(model_config, _TARGET) + + +class ModelingV2KimiK3Mxfp4Sm103Tp16Moetp16ep1(ModelingV2KimiK3): + """Concrete deployment of the shared Kimi K3 text model.""" + + target_spec = _TARGET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/weights.py new file mode 100644 index 000000000000..24233d194709 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp16ep1/weights.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The shared Kimi K3 text checkpoint loader.""" + +from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/__init__.py new file mode 100644 index 000000000000..d644c4a9429b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MXFP4 / sm_103 / tp16_moetp4ep4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/modeling.py new file mode 100644 index 000000000000..bbaa0852d9d5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/modeling.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MXFP4, SM103, attention TP16 and MoE TP4 x EP4.""" + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..modeling import K3TargetSpec, ModelingV2KimiK3 +from ..modeling import _check_construction as _check_target_construction + +_TARGET = K3TargetSpec(sm=(10, 3), quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, moe_tp=4, moe_ep=4) + + +def _check_construction(model_config: ModelConfig) -> None: + _check_target_construction(model_config, _TARGET) + + +class ModelingV2KimiK3Mxfp4Sm103Tp16Moetp4ep4(ModelingV2KimiK3): + """Concrete deployment of the shared Kimi K3 text model.""" + + target_spec = _TARGET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/weights.py new file mode 100644 index 000000000000..24233d194709 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_103__tp16_moetp4ep4/weights.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The shared Kimi K3 text checkpoint loader.""" + +from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/__init__.py new file mode 100644 index 000000000000..bddcc85ec105 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4 / sm_103 / tp16_moetp4ep4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py new file mode 100644 index 000000000000..107683470e2d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4, SM103, attention TP16 and MoE TP4 x EP4.""" + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..modeling import K3TargetSpec, ModelingV2KimiK3 +from ..modeling import _check_construction as _check_target_construction + +_TARGET = K3TargetSpec(sm=(10, 3), quant_algo=QuantAlgo.NVFP4, moe_tp=4, moe_ep=4) + + +def _check_construction(model_config: ModelConfig) -> None: + _check_target_construction(model_config, _TARGET) + + +class ModelingV2KimiK3Nvfp4Sm103Tp16Moetp4ep4(ModelingV2KimiK3): + """Concrete deployment of the shared Kimi K3 text model.""" + + target_spec = _TARGET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py new file mode 100644 index 000000000000..24233d194709 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The shared Kimi K3 text checkpoint loader.""" + +from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py index c43a02997371..f48ed5103338 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 (NVFP4) / sm_107 / tp16 attention, routed experts moe_tp 4 x moe_ep 4.""" +"""Kimi K3 NVFP4 / sm_107 / tp16_moetp4ep4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py deleted file mode 100644 index f5327daf9474..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_comm.py +++ /dev/null @@ -1,402 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode path's residual-update collectives on the catalog's Kimi K3 MNNVL and sandwich entries. - -A layer's post-attention step adds the attention output to the running prefix sum, selects the attention residual -and applies the post-attention RMSNorm. Under TP the attention output is a sum over the TP group, so the step is that -all-reduce followed by the residual update. On a step of at most `AR_ATTN_RES_MAX_TOKENS` tokens (wide decode steps -aside) one collective runs both: - -* `K3DecodeComm.sandwich_oproj`, `comm/k3_sandwich_oproj`: o_proj, its all-reduce and the residual update in one - kernel, at most `SANDWICH_MAX_TOKENS` tokens of an o_proj of the TP16 per-rank shape [7168, 768]. The attention - hands over its gated o_proj input. -* `K3DecodeComm.allreduce_attn_res`, `comm/mnnvl_allreduce_attn_res`, everywhere else: the one-shot MNNVL all-reduce - of the attention's unreduced o_proj partial, with the residual update as its epilogue. - -The pre-attention step after a MoE layer whose row-parallel tail was handed on (a `PendingTail`, `decode_moe.py`) is -the same kind of collective: `K3DecodeComm.sandwich_tail`, `comm/k3_sandwich_tail`, runs the tail GEMV, its -all-reduce and the next layer's residual update (the final norm's, after the last layer) in one kernel. - -The DSpark drafter (`K3DSparkDrafter`) runs its collectives with a plain residual add and RMSNorm on the same state: - -* `K3DecodeComm.sandwich_plain`, `comm/k3_sandwich_plain`: a row-parallel projection of at most `SANDWICH_MAX_TOKENS` - tokens (a layer's attention output projection, or its MLP's SiLU-and-mul and down projection), its all-reduce, the - residual add and the RMSNorm in one kernel. -* `K3DecodeComm.allreduce_norm`, `comm/mnnvl_fusion_allreduce`: the all-reduce of an unreduced projection output of - any token count the MNNVL workspace holds, with the residual add and the RMSNorm in its epilogue (the context - projection's with a zero residual). - -The state is one `MnnvlWorkspace` and one `K3SandwichWorkspace` of the TP group (`K3DecodeComm.create`): collective -over the group and eager, built by the target in `post_load_weights` before any CUDA-graph capture. Every rank must -make the same calls on each in the same order. Which call a step takes is decided from its token count and kind and -from the load-time layout alone, which every rank of the group shares. - -The plain MNNVL all-reduces the decode path keeps (the stock modules') send one-shot up to -`DECODE_AR_ONE_SHOT_MAX_BYTES` (`use_decode_one_shot`), except a wide decode step's (`wide_all_reduce`). -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import NamedTuple, Optional, Tuple - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_oproj import ( - K3SandwichWorkspace, - k3_sandwich_oproj, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_plain import ( - k3_sandwich_plain, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_tail import ( - k3_sandwich_tail, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allreduce_attn_res import ( - mnnvl_allreduce_attn_res, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_fusion_allreduce import ( - mnnvl_fusion_allreduce, - required_buffer_bytes, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( - MnnvlWorkspace, -) - -# The kernel's support predicate: metadata reads only. -from tensorrt_llm._torch.cute_dsl_kernels.k3_sandwich import op as _sandwich_op -from tensorrt_llm._torch.distributed import AllReduceParams - -# The most tokens of a step whose post-attention all-reduce runs one-shot with the residual update in its epilogue. -AR_ATTN_RES_MAX_TOKENS = 16 - -# The most tokens the sandwich kernel takes (one token tile). -SANDWICH_MAX_TOKENS = _sandwich_op.MAX_TOKENS - -# The MNNVL workspace's buffer size: a one-shot call pushes T x 7168 bf16 from each of 16 ranks, 3.5 MiB at T = 16. -MNNVL_BUFFER_BYTES = 4 << 20 - -# The one-shot ceiling of the stock MNNVL all-reduces on the decode path: 8 tokens x 7168 x 16 ranks x 2 B is -# 1.75 MiB, which the stock 1 MiB ceiling would send two-shot. -DECODE_AR_ONE_SHOT_MAX_BYTES = 4 << 20 - -# A wide decode step's ceiling: the stock 1 MiB. At 16 ranks two-shot is faster for every wide step's rows. -WIDE_AR_ONE_SHOT_MAX_BYTES = 1 << 20 - - -def use_decode_one_shot(model: nn.Module) -> None: - """Every stock MNNVL all-reduce of ``model`` sends one-shot up to `DECODE_AR_ONE_SHOT_MAX_BYTES`. Each grows its - workspace on its first eager call of a larger size, before the capture of that size.""" - for module in model.modules(): - mnnvl = getattr(module, "mnnvl_allreduce", None) - if mnnvl is not None: - mnnvl.one_shot_max_bytes = DECODE_AR_ONE_SHOT_MAX_BYTES - mnnvl.early_trigger = True - - -def skip_all_reduce() -> AllReduceParams: - """All-reduce parameters under which a row-parallel module returns this rank's unreduced output (its projection - without its all-reduce): a ``Linear``'s ``all_reduce_params``, a ``GatedMLP``'s ``final_all_reduce_params``.""" - return AllReduceParams(enable_allreduce=False) - - -def wide_all_reduce(all_reduce: nn.Module, x: torch.Tensor) -> torch.Tensor: - """``all_reduce(x)`` of a wide decode step (a stock ``AllReduce`` module, no fusion): its MNNVL all-reduce with - the `WIDE_AR_ONE_SHOT_MAX_BYTES` ceiling, else the module itself.""" - mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) - if mnnvl is not None: - out = mnnvl( - x.contiguous(), AllReduceParams(), one_shot_max_bytes=WIDE_AR_ONE_SHOT_MAX_BYTES - ) - if out is not None: - return out - return all_reduce(x) - - -class PendingTail(NamedTuple): - """A MoE layer's row-parallel tail left to its consumer's fused pre-attention step (``K3DecodeComm.sandwich_tail``): - the reduced latent ``[T, 3584]``, the shared experts' activation, the tail weight ``[latent up columns | padding | - shared down]``, this rank's first latent column and the latent norm's epsilon.""" - - latent: torch.Tensor - act: torch.Tensor - weight: torch.Tensor - lo: int - lat_eps: float - - -def _eps(norm: nn.Module) -> float: - """The epsilon of a KimiK3RMSNorm or ``torch.nn.RMSNorm`` (``eps``; for torch's None, its default: the machine - epsilon of the weight's dtype) or a stock RMSNorm (``variance_epsilon``).""" - eps = norm.eps if hasattr(norm, "eps") else norm.variance_epsilon - return float(torch.finfo(norm.weight.dtype).eps if eps is None else eps) - - -def _res_args(res_proj: nn.Module, res_norm: nn.Module, out_norm: nn.Module) -> tuple: - """The residual update's weights and epsilons in the entries' order.""" - return ( - res_proj.weight.reshape(-1), - res_norm.weight, - out_norm.weight, - _eps(res_norm), - _eps(out_norm), - ) - - -@dataclass(eq=False) -class K3DecodeComm: - """The decode path's collective state for the TP group: the `MnnvlWorkspace` and the `K3SandwichWorkspace` the - post-attention steps run on. Built by `create`; owned by the target and shared by every layer.""" - - mnnvl: MnnvlWorkspace - sandwich: K3SandwichWorkspace - - @classmethod - def create(cls, mapping, oproj_weight: Optional[torch.Tensor] = None) -> "K3DecodeComm": - """The state for ``mapping``'s TP group. Collective: every rank of the group calls it at the same point, - eagerly, before any CUDA-graph capture; each workspace fails on every rank or on none. - - ``oproj_weight``: an o_proj weight ``takes_oproj`` holds for. The sandwich kernel compiles here for its - shape, with one call on a zero row of a zero weight, so no capture compiles it; the call advances the - sandwich workspace on every rank alike.""" - state = cls( - MnnvlWorkspace.create(mapping, MNNVL_BUFFER_BYTES), - K3SandwichWorkspace.create(mapping), - ) - if oproj_weight is not None: - weight = torch.zeros_like(oproj_weight) - hidden = weight.shape[0] - ones = weight.new_ones(hidden) - k3_sandwich_oproj( - weight.new_zeros(1, weight.shape[1]), - weight, - None, - weight.new_zeros(0, 1, hidden), - weight.new_zeros(hidden), - ones, - ones, - 1e-6, - 1e-6, - state.sandwich, - ) - torch.cuda.synchronize(weight.device) - return state - - def compile_tail(self, latent_size: int, act_size: int, tail_weight: torch.Tensor) -> None: - """Compile the sandwich tail kernel for a MoE tail of ``latent_size`` latent and ``act_size`` shared columns - and ``tail_weight``'s shape, with one call on a zero row of a zero weight, before any capture. Collective: every - rank of the group makes the call; it advances the sandwich workspace on every rank alike.""" - weight = torch.zeros_like(tail_weight) - hidden = weight.shape[0] - ones = weight.new_ones(hidden) - k3_sandwich_tail( - weight.new_zeros(1, latent_size), - weight.new_zeros(1, act_size), - weight, - 0, - 1e-6, - None, - weight.new_zeros(0, 1, hidden), - weight.new_zeros(hidden), - ones, - ones, - 1e-6, - 1e-6, - self.sandwich, - ) - torch.cuda.synchronize(weight.device) - - def takes_post_attention(self, hidden_states: torch.Tensor, step) -> bool: - """Whether one of this state's collectives runs the post-attention step of a layer whose attention input is - ``hidden_states``: at most `AR_ATTN_RES_MAX_TOKENS` bf16 rows of a hidden size the MNNVL entry takes, on any - step but a wide decode step (``step.wide``), whose post-attention update stays the all-reduce and the fused - add + attn_res + RMSNorm.""" - if step is not None and step.wide: - return False - if hidden_states.dim() != 2 or hidden_states.dtype != torch.bfloat16: - return False - rows, hidden = hidden_states.shape - return ( - hidden % 1024 == 0 - and hidden <= 8192 - and 0 < rows <= min(AR_ATTN_RES_MAX_TOKENS, self.mnnvl.max_one_shot_tokens(hidden)) - ) - - @staticmethod - def takes_oproj(o_proj: nn.Module, max_snapshots: int) -> bool: - """Whether ``sandwich_oproj`` takes the layer of output projection ``o_proj`` on a step of at most - `SANDWICH_MAX_TOKENS` tokens, decided once the weights are final: a bias-free, contiguous bf16 weight of the - kernel's shape (the TP16 per-rank [7168, 768]) and a snapshot bank of at most ``max_snapshots`` rows the - kernel's candidate count holds.""" - weight = getattr(o_proj, "weight", None) - return ( - getattr(o_proj, "bias", None) is None - and isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.is_cuda - and _sandwich_op.supports(weight.new_empty((1, weight.shape[1])), weight, max_snapshots) - ) - - @staticmethod - def takes_tail( - latent: torch.Tensor, act: torch.Tensor, tail_weight: torch.Tensor, max_snapshots: int - ) -> bool: - """Whether ``sandwich_tail`` takes a MoE tail of these tensors' shapes (rows of ``latent`` and ``act``, at most - `SANDWICH_MAX_TOKENS`; ``tail_weight``) with a snapshot bank of at most ``max_snapshots`` rows: the TP16 - per-rank shapes.""" - return _sandwich_op.supports_tail(latent, act, tail_weight, max_snapshots) - - def allreduce_attn_res( - self, - partial: torch.Tensor, - prefix_sum: Optional[torch.Tensor], - block_residual: torch.Tensor, - res_proj: nn.Module, - res_norm: nn.Module, - out_norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` in one ``comm/mnnvl_allreduce_attn_res`` call: ``updated = prefix_sum + - allreduce(partial)`` (the sum alone without ``prefix_sum``), ``normed = out_norm(attn_res(block_residual..., - updated))``, the attention residual selected with ``res_proj`` and ``res_norm``. ``block_residual`` holds the - valid snapshots, ``[S, T, H]``.""" - return mnnvl_allreduce_attn_res( - partial.contiguous(), - prefix_sum, - block_residual, - *_res_args(res_proj, res_norm, out_norm), - self.mnnvl, - ) - - def sandwich_tail( - self, - pending: PendingTail, - prefix_sum: Optional[torch.Tensor], - block_residual: torch.Tensor, - res_proj: nn.Module, - res_norm: nn.Module, - out_norm: nn.Module, - updated_out: Optional[torch.Tensor] = None, - tap: Optional[torch.Tensor] = None, - tap_updated: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` as ``allreduce_attn_res`` of a MoE layer's row-parallel tail ``pending`` - (``[RMSNorm(latent)[:, lo:lo + 224] | act] @ weight.T``), in one ``comm/k3_sandwich_tail`` call. - ``updated_out``: a bf16 ``[T, H]`` tensor the call stores ``updated`` into (the consumer's snapshot bank row), - returned as ``updated``. ``tap``: a bf16 ``[T, H]`` view (a speculative capture slot) the call also stores - the pre-norm attention-residual mixture into, or ``updated`` with ``tap_updated``.""" - return k3_sandwich_tail( - pending.latent, - pending.act, - pending.weight, - pending.lo, - pending.lat_eps, - prefix_sum, - block_residual, - *_res_args(res_proj, res_norm, out_norm), - self.sandwich, - tap=tap, - tap_updated=tap_updated, - updated_out=updated_out, - ) - - def sandwich_oproj( - self, - core: torch.Tensor, - o_weight: torch.Tensor, - prefix_sum: Optional[torch.Tensor], - block_residual: torch.Tensor, - res_proj: nn.Module, - res_norm: nn.Module, - out_norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` as ``allreduce_attn_res`` of ``core @ o_weight.T``, in one ``comm/k3_sandwich_oproj`` - call: ``core`` is this rank's gated o_proj input ``[T, 768]``, ``o_weight`` its o_proj slice. Bit for bit - o_proj followed by ``allreduce_attn_res``.""" - return k3_sandwich_oproj( - core.contiguous(), - o_weight, - prefix_sum, - block_residual, - *_res_args(res_proj, res_norm, out_norm), - self.sandwich, - ) - - def compile_plain(self, weight: torch.Tensor, swiglu: bool = False) -> bool: - """Compile the plain sandwich (``sandwich_plain``) for a projection of ``weight``'s shape (with ``swiglu``, a - down projection after the SiLU-and-mul) with one call on a zero row of a zero weight, before any capture; - False, with nothing launched, where the kernel does not take that shape. Collective: every rank of the group - makes the call; it advances the sandwich workspace on every rank alike.""" - zero = torch.zeros_like(weight) - hidden, k_in = zero.shape - x = zero.new_zeros(1, 2 * k_in if swiglu else k_in) - residual = zero.new_zeros(1, hidden) - ones = zero.new_ones(hidden) - if not _sandwich_op.supports_plain(x, zero, residual, ones, swiglu): - return False - k3_sandwich_plain(x, zero, residual, ones, 1e-6, self.sandwich, swiglu=swiglu) - torch.cuda.synchronize(zero.device) - return True - - @staticmethod - def takes_plain( - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - swiglu: bool = False, - ) -> bool: - """Whether ``sandwich_plain`` takes the call: at most `SANDWICH_MAX_TOKENS` contiguous bf16 rows ``x`` of a - row-parallel slice ``weight`` [7168, K] (K a multiple of 128 up to 896; with ``swiglu`` K 896 and ``x`` the - ``[gate | up]`` rows, 2 K wide), ``residual`` [rows, 7168], and ``norm``'s bf16 [7168] weight.""" - return _sandwich_op.supports_plain(x, weight, residual, norm.weight, swiglu) - - def sandwich_plain( - self, - x: torch.Tensor, - weight: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - swiglu: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` in one ``comm/k3_sandwich_plain`` call: ``updated = residual + allreduce(x @ - weight.T)`` (with ``swiglu``, of ``silu_and_mul(x) @ weight.T``), ``normed = norm(updated)`` for a plain - RMSNorm ``norm``. Bit for bit, by the kernel's statement, the projection on ``k3_ctm_gemv`` at split 1 (with - ``swiglu``, ``k3_ctm_gemv_swiglu`` at split 2) followed by ``allreduce_norm``'s one-shot call.""" - return k3_sandwich_plain( - x.contiguous(), - weight, - residual.contiguous(), - norm.weight, - _eps(norm), - self.sandwich, - swiglu=swiglu, - ) - - def takes_allreduce_norm(self, rows: int, hidden: int) -> bool: - """Whether ``allreduce_norm`` takes ``rows`` bf16 rows of ``hidden`` columns: the MNNVL workspace holds the - call at the decode path's one-shot ceiling (`DECODE_AR_ONE_SHOT_MAX_BYTES`, two-shot above it).""" - if rows <= 0 or hidden <= 0 or hidden % 8: - return False - world, buffer_bytes = self.mnnvl.world_size, self.mnnvl.buffer_bytes - need = required_buffer_bytes( - rows, hidden, world, torch.bfloat16, DECODE_AR_ONE_SHOT_MAX_BYTES - ) - two_shot = rows * hidden * world * 2 > DECODE_AR_ONE_SHOT_MAX_BYTES - return need <= buffer_bytes and not (two_shot and buffer_bytes % 32) - - def allreduce_norm( - self, partial: torch.Tensor, residual: torch.Tensor, norm: nn.Module - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(normed, updated)`` in one ``comm/mnnvl_fusion_allreduce`` call: ``updated = residual + - allreduce(partial)``, ``normed = norm(updated)`` for a plain RMSNorm ``norm``, sent one-shot up to - `DECODE_AR_ONE_SHOT_MAX_BYTES`. ``partial`` is this rank's unreduced ``[rows, hidden]`` bf16 output of a - row-parallel projection, of a shape ``takes_allreduce_norm`` holds for.""" - return mnnvl_fusion_allreduce( - partial.contiguous(), - self.mnnvl, - DECODE_AR_ONE_SHOT_MAX_BYTES, - residual.contiguous(), - norm.weight, - _eps(norm), - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py deleted file mode 100644 index 0729fe05b91e..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/decode_gemv.py +++ /dev/null @@ -1,467 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. - -* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest - at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, - `gemm/k3_ctm_gemv_wide`, `gemm/k3_ctm_gemv_long`, `gemm/k3_ctm_gemv` or `gemm/k3_ctm_gemv_swiglu`, and, where the - site lists it, more rows (up to `WIDE_ROWS`) on `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused - projections (MLA's [W_a; W_g] with the gate rows through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention - output projection, the built-in MLA path's q_a / kv_a, q_b and gate projections, the MoE decode path's projections - (`decode_moe.py`), two of them with fp32 outputs, and the DSpark drafter's q / k / v, output, gate / up and down - projections (`K3DSparkDrafter`). -* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on - `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the - stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the - drafter's logits on the same head take it too. -* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual - bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. -* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: - gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then - runs the down projection's all-reduce. - -Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. - -A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. -`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The -embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call -whose kernel has not run eagerly is refused. - -The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one -stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import Dict, Iterable, Optional, Set - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul -from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( - k3_ctm_gemv_long, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( - k3_ctm_gemv_swiglu, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( - k3_ctm_gemv_wide, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv -from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( - K3HeadGemvWorkspace, - k3_head_gemv, -) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.concat import concat -from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.split import split - -# The kernels' support predicates: metadata reads only. -from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op -from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op -from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op -from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op -from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE -from tensorrt_llm._utils import mpi_disabled - -# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 -# requests of 8 tokens). -MAX_ROWS = 8 -WIDE_ROWS = 64 - - -@dataclass(frozen=True) -class Site: - """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows - ("decode", "wide", "long", "ctm" for k3_ctm_gemv, or "swiglu" for k3_ctm_gemv_swiglu, whose rows are the - ``[gate | up]`` input of the SiLU-and-mul, 2 ``k`` wide), and k3_ctm_gemv_wide at MAX_ROWS+1..``wide_rows`` rows - where ``wide``. Output columns from ``sig_col0`` on are stored through a sigmoid; ``out_fp32``: an fp32 output - (k3_ctm_gemv_wide only). ``split`` / ``ring`` / ``push``: the CTM kernels' CTAs per 128-row weight tile, - k3_ctm_gemv_long's weight-ring stages, and whether the partial sums are pushed to each row's owner.""" - - n: int - k: int - small: str - wide: bool = False - sig_col0: int = -1 - split: int = 0 - ring: int = 0 - push: bool = False - out_fp32: bool = False - wide_rows: int = WIDE_ROWS - - -SITES: Dict[str, Site] = { - # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), - # the gate rows through a sigmoid. - "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), - # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). - "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), - # The attention output projection (row parallel; 6 heads x 128 in). - "o_proj": Site(7168, 768, "decode", wide=True), - # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. - "kv_a": Site(2112, 7168, "decode"), - "q_b": Site(1152, 1536, "wide"), - "g_proj": Site(768, 7168, "wide"), - # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. - "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), - "dense_down": Site(7168, 2112, "long", split=2, ring=6), - # The MoE decode path (decode_moe.py): this rank's head slice [latent down 224 | router 56] with an fp32 output, - # the shared experts' gate_up, the row-parallel tail [latent up 224 | padding 32 | shared down 384] and the - # replicated tail's latent up projection with an fp32 output. The wide kernel takes the head slice and the tail up - # to 32 rows, where it is faster than the stock GEMM. - "moe_head": Site(280, 7168, "wide", wide=True, out_fp32=True, wide_rows=32), - "moe_shared_gate_up": Site(768, 7168, "wide", wide=True), - "moe_tail": Site(7168, 640, "wide", wide=True, wide_rows=32), - "moe_up": Site(7168, 3584, "wide", out_fp32=True), - # The DSpark drafter's projections (K3DSparkDrafter; 6 query heads, 1 KV head of 64, MLP width 896): the fused - # q / k / v, the output projection (row parallel), gate_up [gate 896 | up 896], and the down projection with the - # SiLU-and-mul (row parallel). - "drafter_qkv": Site(512, 7168, "long", split=8, ring=6, push=True), - "drafter_o": Site(7168, 384, "ctm", split=1, push=True), - "drafter_gate_up": Site(1792, 7168, "long", split=8, ring=6, push=True), - "drafter_down": Site(7168, 896, "swiglu", split=2, push=True), -} - - -def _width(spec: Site) -> int: - """The row width a site's kernel reads: ``k``, or the ``[gate | up]`` input of a SiLU-and-mul site, 2 ``k``.""" - return 2 * spec.k if spec.small == "swiglu" else spec.k - - -def _wide_tile(rows: int) -> int: - from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel - - return k3_ctm_gemv_kernel.wide_tile(rows) - - -def _run( - spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor -) -> Optional[torch.Tensor]: - """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" - if kernel == "decode": - if spec.sig_col0 >= 0 or spec.out_fp32 or not _decode_op.supports(x2d, weight): - return None - return k3_decode_gemv(x2d, weight) - if kernel == "wide": - if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, spec.out_fp32): - return None - return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0, out_fp32=spec.out_fp32) - if spec.out_fp32: - return None - if kernel == "ctm": - if spec.sig_col0 >= 0 or not _ctm_op.supports(x2d, weight, spec.split): - return None - return k3_ctm_gemv(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) - if kernel == "swiglu": - if spec.sig_col0 >= 0 or not _ctm_op.supports_swiglu(x2d, weight, spec.split): - return None - return k3_ctm_gemv_swiglu(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) - # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. - sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count - if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( - x2d, weight, spec.split, spec.ring - ): - return None - return k3_ctm_gemv_long( - x2d, - weight, - sig_col0=spec.sig_col0, - split=spec.split, - ring=spec.ring, - trigger_early=True, - push=spec.push, - ) - - -def _capturing() -> bool: - return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() - - -def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: - """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" - if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: - return x2d.clone(memory_format=torch.contiguous_format) - return x2d - - -def _head_takes_module(lm_head: nn.Module) -> bool: - """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along - the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" - weight = getattr(lm_head, "weight", None) - mapping = getattr(lm_head, "mapping", None) - return ( - isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.dtype == torch.bfloat16 - and weight.is_contiguous() - and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" - and getattr(lm_head, "gather_output", False) - and getattr(lm_head, "gather_output_sizes", None) is None - and getattr(lm_head, "padding_size", None) == 0 - and getattr(lm_head, "bias", None) is None - and not getattr(lm_head, "has_any_quant", True) - and mapping is not None - and not mapping.enable_attention_dp - ) - - -def _plain_rmsnorm(norm: nn.Module) -> bool: - """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for - bit. An unknown module fails closed.""" - weight = getattr(norm, "weight", None) - return ( - IS_FLASHINFER_AVAILABLE - and isinstance(weight, torch.Tensor) - and weight.dtype == torch.bfloat16 - and not getattr(norm, "use_gemma", True) - and not getattr(norm, "is_nvfp4", True) - and not getattr(norm, "use_cuda_tile", True) - and not getattr(norm, "return_hp_output", True) - and getattr(norm, "nvfp4_scale", None) is None - and hasattr(norm, "variance_epsilon") - ) - - -class K3DecodeGemvs: - """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran - eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are - final; owned by the target.""" - - def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: - self.head_workspace = head_workspace - self._ran: Set[tuple] = set() - - @classmethod - def create( - cls, - lm_head: Optional[nn.Module] = None, - sites: Iterable[str] = tuple(SITES), - device: Optional[torch.device] = None, - ) -> "K3DecodeGemvs": - """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: - it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on - zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or - head whose kernel does not take its shape keeps the generic path.""" - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" - ) - head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None - if device is None: - device = ( - head_weight.device - if isinstance(head_weight, torch.Tensor) - else torch.device("cuda", torch.cuda.current_device()) - ) - state = cls() - for site in sites: - spec = SITES[site] - weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) - for rows in (1, 16, 32, 64) if spec.wide else (1,): - if rows > spec.wide_rows: - continue - state._project(site, weight.new_zeros(rows, _width(spec)), weight, warm=True) - del weight - if "dense_gate_up" in sites: - gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) - if _ctm_op.supports_situ_mul(gu): - for linear_beta in (None, 1.0): - k3_situ_mul(gu, 1.0, linear_beta) - state._ran.add(("situ_mul", linear_beta is not None)) - if lm_head is not None and _head_takes_module(lm_head): - x = head_weight.new_zeros(1, head_weight.shape[1]) - if _head_op.supports(x, head_weight): - workspace = K3HeadGemvWorkspace.create( - head_weight.shape[0], head_weight.shape[1], head_weight.device - ) - k3_head_gemv(x, head_weight, workspace) - state.head_workspace = workspace - state._ran.add(("lm_head",)) - torch.cuda.synchronize(device) - return state - - def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: - """``x @ weight.T`` (``[..., N]``: bf16, the site's sigmoid columns through the sigmoid; fp32 at an - ``out_fp32`` site; on a SiLU-and-mul site, ``silu(gate) * up`` of ``x = [gate | up]`` first) for ``site``'s - weight on its decode kernel, or None where none takes the call: more rows than the site's kernels take, - another shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its - module.""" - return self._project(site, x, weight, warm=False) - - def _project( - self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool - ) -> Optional[torch.Tensor]: - spec = SITES[site] - width = _width(spec) - if ( - weight.dtype != torch.bfloat16 - or tuple(weight.shape) != (spec.n, spec.k) - or not weight.is_contiguous() - or x.dtype != torch.bfloat16 - or x.dim() < 1 - or x.shape[-1] != width - ): - return None - rows = x.numel() // width - if 0 < rows <= MAX_ROWS: - kernel = spec.small - elif spec.wide and MAX_ROWS < rows <= spec.wide_rows: - kernel = "wide" - else: - return None - key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) - capturing = _capturing() - if capturing and not warm and key not in self._ran: - return None - y = _run(spec, kernel, _dense_rows(x.reshape(rows, width)), weight) - if y is None: - return None - if not capturing: - self._ran.add(key) - return y.view(*x.shape[:-1], spec.n) - - def lm_head_logits( - self, rows: torch.Tensor, lm_head: nn.Module, gather: bool = True - ) -> Optional[torch.Tensor]: - """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on - ``gemm/k3_head_gemv`` (``gather`` False: the shard alone, ``[M, vocab / TP]``); None where it does not take - the call (more than `MAX_ROWS` rows, another head, ...).""" - workspace = self.head_workspace - if workspace is None or not _head_takes_module(lm_head): - return None - weight = lm_head.weight - if ( - tuple(weight.shape) != (workspace.n_out, workspace.k_in) - or weight.device != workspace.partials.device - or rows.dim() != 2 - or rows.dtype != torch.bfloat16 - or not 0 < rows.shape[0] <= MAX_ROWS - or rows.shape[1] != workspace.k_in - ): - return None - if _capturing() and ("lm_head",) not in self._ran: - return None - group = lm_head.mapping.tp_group - if gather and len(group) > 1 and mpi_disabled(): - return None - x = _dense_rows(rows) - if not _head_op.supports(x, weight): - return None - local = k3_head_gemv(x, weight, workspace) - if not gather or len(group) == 1: - return local - gathered = allgather(local, None, group) - return concat(list(split(gathered, rows.shape[0], dim=0)), dim=-1) - - def dense_mlp( - self, - x: torch.Tensor, - gate_up_weight: torch.Tensor, - down_weight: torch.Tensor, - beta: float, - linear_beta: Optional[float], - ) -> Optional[torch.Tensor]: - """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on - ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on - ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then - runs the module.""" - gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] - rows = x.shape[0] if x.dim() == 2 else 0 - if ( - not 0 < rows <= MAX_ROWS - or linear_beta == 0.0 - or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) - or tuple(down_weight.shape) != (down.n, down.k) - or gate_up.n != 2 * down.k - ): - return None - if ( - _capturing() - and not { - ("dense_gate_up", "long", 0), - ("dense_down", "long", 0), - ("situ_mul", linear_beta is not None), - } - <= self._ran - ): - return None - gu = self.project("dense_gate_up", x, gate_up_weight) - if gu is None or not _ctm_op.supports_situ_mul(gu): - return None - return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) - - def embed_norm( - self, - input_ids: torch.Tensor, - table: torch.Tensor, - norm: nn.Module, - bank: torch.Tensor, - ) -> Optional[torch.Tensor]: - """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one - ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not - apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose - kernel has not run eagerly).""" - if not _plain_rmsnorm(norm): - return None - ids = input_ids.reshape(-1) - if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): - return None - key = ("embed_norm", ids.numel(), ids.dtype) - capturing = _capturing() - if capturing and key not in self._ran: - return None - normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) - if not capturing: - self._ran.add(key) - return normed - - -class _K3Head: - """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" - - __slots__ = ("_gemvs", "_lm_head") - - def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: - self._gemvs = gemvs - self._lm_head = lm_head - - def __call__(self, rows: torch.Tensor) -> torch.Tensor: - logits = self._gemvs.lm_head_logits(rows, self._lm_head) - return self._lm_head(rows) if logits is None else logits - - -class K3LogitsProcessor(nn.Module): - """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. - - It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock - processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. - """ - - def __init__(self, stock: nn.Module) -> None: - super().__init__() - self.stock = stock - self.gemvs: Optional[K3DecodeGemvs] = None - - def forward( - self, - hidden_states: torch.Tensor, - lm_head: nn.Module, - attn_metadata, - return_context_logits: bool = False, - ) -> torch.Tensor: - head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) - return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) - - def lm_head_shard(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: - """This rank's bf16 vocabulary shard of ``lm_head(rows)`` on ``gemm/k3_head_gemv``, without the gather, or - None where ``gemvs`` does not take the rows. A speculative worker that keeps its logits vocabulary-sharded - computes them here, so they hold the values of the logits this processor gathers.""" - if self.gemvs is None: - return None - return self.gemvs.lm_head_logits(rows, lm_head, gather=False) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py index 883f1a0ea418..55d2c599ebd3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py @@ -1,3012 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 NVFP4 on SM107, TP16 attention and MoE TP4 x EP4. +"""Kimi K3 NVFP4, SM107, attention TP16 and MoE TP4 x EP4.""" -The target reads the ModelOpt mixed-precision checkpoint without converting its -NVFP4 experts or FP8 attention weights to another format. Quantized attention and -routed experts use their stock implementations. The BF16 embedding, dense MLP, -LM head, attention-residual collectives and standalone GQA DSpark drafter retain -the PR's decode specializations. MLA DSpark uses its stock drafter and worker. -This target has no SM107 accuracy or performance certification until its -corresponding GPU tests and model gates have run. - -Unlike the MXFP4 target, this path permits the stock attention implementation's -quantized KV cache and BF16 recurrent states. It does not allocate the MXFP4 -fused MoE state or the BF16 fused attention projection workspaces. -""" - -from __future__ import annotations - -import copy -import math -import os -from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union - -import torch -from torch import nn - -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( - k3_drafter_attn_qknorm, -) -from tensorrt_llm._torch.attention.backends import AttentionMetadata - -# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. -from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 -from tensorrt_llm._torch.distributed import AllReduce from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM -from tensorrt_llm._torch.models.modeling_dspark import ( - GQADSparkForCausalLM, - draft_is_embedded_in_target, -) -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM -from tensorrt_llm._torch.models.modeling_utils import DecoderModel -from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention -from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention -from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace -from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel -from tensorrt_llm._torch.modules.rms_norm import RMSNorm -from tensorrt_llm._torch.modules.situ import SituAndMul -from tensorrt_llm._torch.moe.fused_moe import ( - ConfigurableMoE, - SiTuActivation, - TRTLLMGenFusedMoE, - create_moe, -) -from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind -from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod -from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph -from tensorrt_llm._torch.pyexecutor.config_utils import is_mla -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 -from tensorrt_llm._torch.utils import AuxStreamType -from tensorrt_llm.functional import AllReduceStrategy -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - -from . import decode_comm as _decode_comm -from . import decode_gemv as _decode_gemv -from . import spec_worker as _spec_worker -from . import weights as _weights - -if TYPE_CHECKING: - from transformers import PretrainedConfig - - -# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, -# and the certification is per arch. -_SM = (10, 7) - -#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in -#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router -#: and fused-A GEMMs), then the decode kernels'. -REQUIRED_TRTLLM_OPS = ( - "attn_res_fwd", - "attn_res_rmsnorm_fwd", - "attn_res_add_rmsnorm_fwd", - "attn_res_add_rmsnorm_persistent_fwd", - "kda_prefill", - "kda_decode", - "kda_mtp_decode", - "dsv3_router_gemm_op", - "dsv3_fused_a_gemm_op", - # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). - "k3_decode_gemv", - "k3_ctm_gemv_wide", - "k3_ctm_gemv_long", - "k3_situ_mul", - "k3_head_gemv", - "k3_embed_norm", - "allgather", - # The decode path's collectives (decode_comm.py). - "mnnvl_allreduce_attn_res", - "k3_sandwich_oproj", - "k3_sandwich_tail", - "mnnvl_allgather_split", - # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces - # with the residual add and RMSNorm (the split context projection's with hidden_norm). - "k3_ctm_gemv", - "k3_ctm_gemv_swiglu", - "k3_drafter_attn_qknorm", - "k3_sandwich_plain", - "mnnvl_fusion_allreduce", -) - -#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. -#: A renamed field upstream fails here, loudly, instead of reading as a default. -REQUIRED_ENGINE_FIELDS = { - "attn_metadata": ( - "num_contexts", - "num_generations", - "num_seqs", - "num_tokens", - "seq_lens", - "tokens_per_block", - "kv_cache_manager", - "mamba_metadata", - ), - "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), -} - -#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every -#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry -#: leaves the list when a catalog entry replaces it. -UNCERTIFIED_GENERIC_CALLS = ( - # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. - "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", - "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", - "tensorrt_llm._torch.models.modeling_utils.DecoderModel", - # The text model's stock modules. - "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", - "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", - "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", - "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", - "tensorrt_llm._torch.moe.fused_moe.create_moe", - "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", - "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", - "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", - "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", - "tensorrt_llm._torch.modules.situ.SituAndMul", - "tensorrt_llm._torch.modules.rms_norm.RMSNorm", - "tensorrt_llm._torch.distributed.AllReduce", - "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", - # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries - # do not take a block, its context projection without the TP group's collective state, its context k / v, its - # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. - "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", - "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", - "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", - "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", -) - -# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware -# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). -DECODE_MAX_TOKENS = 8 -MAX_REQUESTS = 8 -MAX_TOKENS_PER_REQUEST = 8 - -_LANG_PREFIX = "language_model." - - -# ---------------------------------------------------------------------------------------------------------------------- -# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and -# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, -# `K3DecodeMLA` below). -# ---------------------------------------------------------------------------------------------------------------------- - -# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections -# instead of the min-latency fused GEMM op (read once at import). -_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( - os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" -) - - -# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table -# holds cos=1/sin=0), but the chunked-context path indexes the table by -# absolute position, so it must cover max_position_embeddings (~512MB per -# backend for the 1M-position checkpoint); a smaller table is read out of -# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context -# deployments. -_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" - - -class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): - """Compute expert choices separately from the quantized MoE GEMMs. - - Preserve the PR's routing arithmetic while the backend applies the - checkpoint's NVFP4 activation and weight scales. - """ - - @property - def requires_separated_routing(self) -> bool: - return True - - -class KimiK3MoEGate(nn.Module): - """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" - - def __init__( - self, - config: Any, - *, - logits_gemm_dtype: torch.dtype | None = None, - device: torch.device | None = None, - ) -> None: - super().__init__() - self.config = config - self.top_k = config.num_experts_per_token - self.num_experts = config.num_experts - self.routed_scaling_factor = config.routed_scaling_factor - self.moe_router_activation_func = config.moe_router_activation_func - self.num_expert_group = getattr(config, "num_expert_group", 1) - self.topk_group = getattr(config, "topk_group", 1) - self.moe_renormalize = config.moe_renormalize - self.gating_dim = config.hidden_size - - assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( - "K3 MoE gate supports sigmoid or softmax scoring only" - ) - - # The checkpoint stores the gate weight in bf16. Storing it in bf16 - # permits the single bf16xbf16 router GEMM while retaining fp32 output. - weight_dtype = logits_gemm_dtype or torch.float32 - self.weight = nn.Parameter( - torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) - ) - self.e_score_correction_bias = nn.Parameter( - torch.empty(self.num_experts, dtype=torch.float32, device=device) - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" - hidden_2d = hidden_states.reshape(-1, self.gating_dim) - if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: - return torch.ops.trtllm.dsv3_router_gemm_op( - hidden_2d.contiguous(), - self.weight.t(), - bias=None, - out_dtype=torch.float32, - ) - return torch.nn.functional.linear( - hidden_2d.type(torch.float32), - self.weight.type(torch.float32), - None, - ) - - @property - def routing_method(self) -> KimiK3MoeRoutingMethod: - """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" - if self.moe_router_activation_func != "sigmoid": - raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") - if not self.moe_renormalize: - raise ValueError( - "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." - ) - return KimiK3MoeRoutingMethod( - top_k=self.top_k, - n_group=self.num_expert_group, - topk_group=self.topk_group, - routed_scaling_factor=self.routed_scaling_factor, - callable_e_score_correction_bias=lambda: self.e_score_correction_bias, - is_fused=True, - ) - - -class KimiK3RMSNorm(nn.Module): - """RMSNorm matching the Kimi checkpoint implementation's rounding.""" - - def __init__( - self, - hidden_size: int, - eps: float = 1e-6, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - ) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) - self.eps = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states_float = hidden_states.to(torch.float32) - variance = hidden_states_float.pow(2).mean(-1, keepdim=True) - hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) - return self.weight * hidden_states_float.to(input_dtype) - - -def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: - """Return the finite SiTu betas required by the routed-expert kernels.""" - config_situ_beta = getattr(cfg, "activation_situ_beta", None) - situ_beta = 1.0 if config_situ_beta is None else config_situ_beta - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - if situ_linear_beta is None: - raise ValueError( - "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " - "None means an identity linear branch that the fused kernels cannot represent." - ) - if situ_beta <= 0 or situ_linear_beta <= 0: - raise ValueError( - f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." - ) - return float(situ_beta), float(situ_linear_beta) - - -def _get_text_config(pretrained_config: "PretrainedConfig"): - """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" - if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( - not hasattr(pretrained_config, "linear_attn_config") - and hasattr(pretrained_config, "text_config") - ): - return pretrained_config.text_config - return pretrained_config - - -def _is_kda_layer(cfg, layer_idx: int) -> bool: - return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] - - -def _is_mla_layer(cfg, layer_idx: int) -> bool: - return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] - - -KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" - - -_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" - - -KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" - - -_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" - - -KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" - - -KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" - - -_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") - - -# Turning the port on should not require knowing which of the three topologies -# is the right one: "1" selects the measured policy (``split``), so enabling the -# feature and choosing the policy are one action through one variable. -_ATTN_RES_TOPOLOGY_ON = "1" - - -def _read_attn_res_topology() -> str: - """Default ``per_token``: the persistent kernel is opt-in. - - ``1`` is the only accepted on-value and resolves to ``split``. The named - topologies stay available for measurement: ``persistent`` uses the - persistent kernel at every shape it implements, ``per_token`` at none. - """ - raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") - if raw == _ATTN_RES_TOPOLOGY_ON: - return "split" - if raw not in _ATTN_RES_TOPOLOGIES: - # Loudly, for the same reason as the token ceiling below: a mistyped A/B - # arm that silently fell back to the default would measure one side - # twice and report no difference. - raise ValueError( - f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " - f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " - f"(which means 'split'), got {raw!r}" - ) - return raw - - -def _read_fused_attn_res_max_tokens() -> int: - """Resolved after the topology, because its default follows it. - - With the persistent port off -- the default -- the ceiling is 1, the - pre-existing gate: the fused epilogue is taken at the single-token decode - shape and nowhere else. Enabling the port raises it to 32, the top of the - measured range, so that "off" keeps meaning "unchanged". - """ - default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" - raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) - try: - value = int(raw) - except ValueError: - # Failing loudly matters more than usual here: a mistyped A/B arm that - # silently fell back to the default would measure the candidate twice - # and report no difference. - raise ValueError( - f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" - ) from None - if value < 1: - raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") - return value - - -_ATTN_RES_TOPOLOGY = _read_attn_res_topology() - - -_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() - - -def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: - """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. - - No token ceiling: the persistent grid is sized by the SM count, not by the - token count, so prefill is the case it exists for. - """ - del M # deliberately unused; see above - return H == 7168 and 2 <= N <= 9 - - -def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: - """Pick between the two fused kernels for this call site. - - ``persistent`` takes the persistent kernel at every shape it implements; - ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in - for the prefill/decode boundary. Shapes the persistent kernel does not - implement fall through to the caller's existing gate and land on the - unfused path. - """ - if not _persistent_attn_res_applicable(M, H, N): - return False - if _ATTN_RES_TOPOLOGY == "persistent": - return True - return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens - - -def _apply_attn_res_fused( - prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm -) -> Optional[torch.Tensor]: - """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. - - Returns ``None`` when the call falls outside the fused kernel's - contract (dtype/shape/arch) so the caller can use the exact fp32 reference - instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` - layout. Candidate order matches the reference: snapshots first, the - running prefix sum last. - """ - if ( - prefix_sum.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not block_residual.is_cuda - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): - return None - try: - attn_res_op = torch.ops.trtllm.attn_res_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output, _rsigma, _probs, _logits = attn_res_op( - layer_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - ) - return output.reshape(M, H) - - -def _rms_norm_eps(norm: nn.Module) -> float: - if hasattr(norm, "eps"): - return float(norm.eps) - return float(norm.variance_epsilon) - - -def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: - """Report whether the fused path was actually reached, once per shape. - - ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that - the shape gate let the call through, and a rejected call looks exactly like - a disabled one in the logs. Emitted at debug level: it is once per distinct - shape, not once per process, so it is a diagnostic rather than a summary. - """ - logger.debug_once( - f"Kimi K3 attn-res fusion [{site}]: " - f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", - key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", - ) - - -def _apply_attn_res_rmsnorm_fused( - prefix_sum: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Optional[torch.Tensor]: - """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, - default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" - if max_fused_tokens is None: - max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS - if ( - prefix_sum.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not block_residual.is_cuda - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - N = K + 1 - # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, - # which is the measured window; larger token counts have not been measured - # and fall back to the unfused add + attn_res_fwd + RMSNorm path. - if _use_persistent_attn_res(M, H, N, max_fused_tokens): - try: - persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd - except (AttributeError, RuntimeError): - return None - _, output = persistent_op( - prefix_sum.reshape(M, 1, H).contiguous(), - None, - block_residual.reshape(K, M, 1, H).contiguous(), - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - ) - _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) - return output.reshape(M, H) - - if M > max_fused_tokens or H != 7168 or N > 12: - _note_attn_res_fusion("attn_res+norm", False, M, H, N) - return None - try: - attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - output = attn_res_rmsnorm_op( - layer_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - early_trigger=True, - ) - _note_attn_res_fusion("attn_res+norm", True, M, H, N) - return output.reshape(M, H) - - -def _apply_attn_res_add_rmsnorm_fused( - prefix_sum: torch.Tensor, - addend: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: - """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in - ``_apply_attn_res_rmsnorm_fused``). - - The production residual add produces a BF16 tensor that remains live across - the following MLP. The kernel therefore returns that materialized, - BF16-rounded prefix sum alongside the normalized attention-residual output, - while avoiding a separate add launch and a re-read of the intermediate by - attention-residual selection. - """ - if max_fused_tokens is None: - max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS - if ( - prefix_sum.dtype is not torch.bfloat16 - or addend.dtype is not torch.bfloat16 - or not prefix_sum.is_cuda - or not addend.is_cuda - or not block_residual.is_cuda - or prefix_sum.shape != addend.shape - ): - return None - M, H = prefix_sum.shape - K = int(block_residual.shape[0]) - N = K + 1 - # Same measured window as _apply_attn_res_rmsnorm_fused above. - if _use_persistent_attn_res(M, H, N, max_fused_tokens): - try: - persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd - except (AttributeError, RuntimeError): - return None - updated_prefix_sum, output = persistent_op( - prefix_sum.reshape(M, 1, H).contiguous(), - addend.reshape(M, 1, H).contiguous(), - block_residual.reshape(K, M, 1, H).contiguous(), - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - ) - _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) - return updated_prefix_sum.reshape(M, H), output.reshape(M, H) - - if M > max_fused_tokens or H != 7168 or N > 12: - _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) - return None - try: - attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd - except (AttributeError, RuntimeError): - return None - layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() - addend_kernel = addend.reshape(M, 1, H).contiguous() - block_kernel = block_residual.reshape(K, M, 1, H).contiguous() - updated_prefix_sum, output = attn_res_add_rmsnorm_op( - layer_kernel, - addend_kernel, - block_kernel, - proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), - norm.weight.to(torch.bfloat16).contiguous(), - output_norm.weight.to(torch.bfloat16).contiguous(), - float(norm.eps), - _rms_norm_eps(output_norm), - early_trigger=True, - ) - _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) - return updated_prefix_sum.reshape(M, H), output.reshape(M, H) - - -def _apply_attn_res( - prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm -) -> torch.Tensor: - """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). - - prefix_sum: ``[num_tokens, hidden_size]`` - block_residual: ``[num_snapshots, num_tokens, hidden_size]`` - - Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's - contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. - Only the fallback boundary restores the HF ``[M, K, H]`` layout. - """ - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) - if fused is not None: - return fused - block_residual_hf = block_residual.transpose(0, 1) - v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) - v_float = v.float() - variance = v_float.pow(2).mean(-1, keepdim=True) - k = v_float * torch.rsqrt(variance + norm.eps) - score_weight = norm.weight.float() * proj.weight.squeeze(0).float() - scores = (k * score_weight).sum(-1) - probs = scores.softmax(-1).unsqueeze(1) - hidden_states = torch.matmul(probs, v_float).squeeze(1) - return hidden_states.to(v.dtype) - - -def _apply_attn_res_and_rmsnorm( - prefix_sum: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> torch.Tensor: - """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the - fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_rmsnorm_fused( - prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens - ) - if fused is not None: - return fused - return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) - - -def _apply_attn_res_add_and_rmsnorm( - prefix_sum: torch.Tensor, - addend: torch.Tensor, - block_residual: torch.Tensor, - proj: nn.Linear, - norm: KimiK3RMSNorm, - output_norm: nn.Module, - max_fused_tokens: Optional[int] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in - ``_apply_attn_res_and_rmsnorm``).""" - if _FUSED_ATTN_RES_ENABLED: - fused = _apply_attn_res_add_rmsnorm_fused( - prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens - ) - if fused is not None: - return fused - updated_prefix_sum = prefix_sum + addend - return updated_prefix_sum, _apply_attn_res_and_rmsnorm( - updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens - ) - - -# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where -# they are faster than the add -> attn_res -> RMSNorm chain. -_WIDE_ATTN_RES_MAX_TOKENS = 32 - - -def _attn_res_max_tokens(step: DecodeStep) -> int: - """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step - ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" - return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS - - -_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") - - -# The routed experts' module path prefixes. ``exclude_modules`` matches with -# wildcards and walks ancestor prefixes, so an empty prefix would widen what -# matches instead of just missing. -_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") - - -def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: - backend.quant_method.load_streaming_nvfp4_expert( - backend, - global_expert_id=expert_idx, - local_slot_id=local_slot_id, - **{ - f"{w}_{kind}": get_tensor(f"{base}.{expert_idx}.{w}.{kind}") - for w in ("w1", "w2", "w3") - for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") - }, - ) - - -class _K3ExpertCkptSpec(NamedTuple): - """How one routed-expert quantization is spelled and loaded.""" - - # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. - kinds: Tuple[str, ...] - loader: Callable[..., None] - # Set of filled slots the loader maintains, checked after the load. - loaded_slots_attr: str - # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). - needs_layer_finalize: bool - - -_K3_EXPERT_CKPT_SPECS = { - QuantAlgo.NVFP4: _K3ExpertCkptSpec( - kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"), - loader=_load_nvfp4_expert, - loaded_slots_attr="_streamed_expert_slots", - needs_layer_finalize=True, - ), -} - - -def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: - spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) - if spec is None: - raise NotImplementedError( - f"Kimi K3 routed experts are quantized as {quant_algo}, for which " - "no per-expert checkpoint layout is known. Supported: " - f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." - ) - return spec - - -class KimiK3MoERuntime(nn.Module): - """Kimi K3 latent MoE block backed by ConfigurableMoE.""" - - def __init__( - self, - model_config: ModelConfig, - cfg, - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ): - """Build the routed experts and the shared expert for one MoE layer. - - ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: - the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields - that ``ModelConfig`` does not carry. - - ``aux_stream_dict`` is shared across every layer of the model, so the - streams reached through it are borrowed and must not be synchronized - or reassigned here. - """ - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = cfg.hidden_size - self.num_experts = cfg.num_experts - self.top_k = cfg.num_experts_per_token - self.moe_hidden_size = cfg.routed_expert_hidden_size - # ValueError (not assert): these guard unsupported checkpoint - # configurations and must stay active under ``python -O``. - if self.moe_hidden_size is None: - raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") - if not getattr(cfg, "latent_moe_use_norm", False): - raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") - - situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) - self._situ_betas = (situ_beta, situ_linear_beta) - dtype = torch.bfloat16 - - # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 - # accumulate/output (checkpoint stores the gate weight in bf16; saves a - # per-layer input cast + fp32 splitK pair on the bs1 decode path). - # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. - _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") - _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True - self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) - - routed_moe_model_config = self._routed_moe_model_config(model_config) - routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) - # Resolved here so ``load_weights`` reads the checkpoint layout off the - # module instead of re-deriving it at each of its three call sites. - self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) - routed_moe_kwargs = dict( - routing_method=self.gate.routing_method, - num_experts=self.num_experts, - hidden_size=self.moe_hidden_size, - intermediate_size=cfg.moe_intermediate_size, - dtype=dtype, - # Kimi owns the latent reduction so it can order that collective - # after the shared expert's auxiliary-stream reduction. - reduce_results=False, - model_config=routed_moe_model_config, - override_quant_config=routed_quant_config, - layer_idx=layer_idx, - aux_stream_dict=aux_stream_dict, - # Let CommunicationFactory select the best available strategy. - communication_method=None, - activation=SiTuActivation( - gate_softcap=situ_beta, - linear_softcap=situ_linear_beta, - ), - # A request that silently degraded to CUTLASS would be benchmarked - # as if it were the backend that was asked for, and the decline is - # easy to trigger: MegaMoE has its own token / top-k limits and is - # EP-only, and CuteDSL declines on activation shape, SM version and - # the CuTe DSL dependency. As measured once: a CUTEDSL request - # was turned down on every one of the 92 MoE layers, on all 16 - # ranks, and still produced correct text and a zero exit -- the - # only trace was a warning line per layer. Fail in the resolver - # instead, which reports the rejection trail. - # - # CUTLASS is absent on purpose: it is the fallback target, so - # "degraded to CUTLASS" is not a thing that can happen to it. - allow_backend_degradation=routed_moe_model_config.moe_backend - not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), - ) - self._check_trtllm_situ_quant( - routed_moe_model_config.moe_backend, routed_quant_config.quant_algo - ) - - self.routed_experts = create_moe(**routed_moe_kwargs) - if not isinstance(self.routed_experts, ConfigurableMoE): - raise RuntimeError( - "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." - ) - if self.routed_experts.layer_load_balancer is not None: - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming does not yet support " - "dynamic EPLB or replicated expert slots." - ) - local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) - if local_expert_ids != list( - range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) - ): - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming currently requires a " - "contiguous static expert partition." - ) - self.local_expert_ids = tuple(local_expert_ids) - self.experts_per_rank = len(local_expert_ids) - self.expert_lo = local_expert_ids[0] - self.expert_hi = self.expert_lo + self.experts_per_rank - - shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts - shared_model_config = copy.copy(model_config) - shared_model_config.quant_config = QuantConfig() - # Direct MoE-TP leaves both branches as partials for one concatenated - # all-reduce. - use_shared_tp = model_config.mapping.tp_size > 1 - self._reduce_routed_output = ( - use_shared_tp - and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM - ) - if self._reduce_routed_output and self.routed_experts.all_reduce is None: - raise RuntimeError( - "Kimi K3 direct MoE tensor parallelism requires the " - "ConfigurableMoE all-reduce even when reduce_results=False." - ) - self.shared_experts = GatedMLP( - hidden_size=cfg.hidden_size, - intermediate_size=shared_intermediate, - bias=False, - activation=SituAndMul( - beta=situ_beta, - linear_beta=situ_linear_beta, - use_fused_activation=True, - ), - dtype=dtype, - config=shared_model_config, - reduce_output=use_shared_tp, - layer_idx=layer_idx, - is_shared_expert=True, - ) - # Side stream (+ fork/join events) for overlapping shared-expert - # compute with the routed chain. Only engaged when multi-stream is - # active (CUDA graphs on); otherwise both run in order on the default - # stream. - self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] - self.moe_main_event = torch.cuda.Event() - self.moe_shared_event = torch.cuda.Event() - self.routed_expert_down_proj = nn.Linear( - cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype - ) - self.routed_expert_up_proj = nn.Linear( - self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype - ) - # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer - # fallback is the same fp32-variance eager math as KimiK3RMSNorm). - self.routed_expert_norm = RMSNorm( - hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. - self.decode_moe = None - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - @staticmethod - def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: - if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): - return projection(hidden_states) - return torch.ops.trtllm.dsv3_fused_a_gemm_op( - hidden_states, projection.weight.t(), None, None - ) - - @staticmethod - def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: - """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit - ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (4 x 4, asserted at - construction).""" - return mapping.moe_tp_size, mapping.moe_ep_size - - @staticmethod - def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: - """Read the NVFP4 expert format, preserving all per-layer scale metadata.""" - global_quant = model_config.quant_config - if global_quant is None: - raise ValueError("The NVFP4 target requires checkpoint quantization metadata") - if any( - global_quant.is_module_excluded_from_quantization( - f"{prefix}layers.{layer_idx}.{suffix}" - ) - for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES - for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES - ): - raise ValueError(f"NVFP4 experts at layer {layer_idx} are excluded from quantization") - declarations = model_config.quant_config_dict or {} - matches = [ - declarations[name] - for prefix in ("language_model.model.", "model.", "") - for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES - if (name := f"{prefix}layers.{layer_idx}.{suffix}") in declarations - ] - if matches: - if any( - quant.quant_algo != matches[0].quant_algo - or quant.group_size != matches[0].group_size - for quant in matches[1:] - ): - raise ValueError(f"Conflicting expert quantization aliases at layer {layer_idx}") - return matches[0] - if global_quant.quant_algo == QuantAlgo.NVFP4: - return global_quant - raise ValueError(f"Missing NVFP4 expert declaration at layer {layer_idx}") - - @staticmethod - def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: - """Reject formats without a fused SiTU cubin in the resolved TRTLLM backend. - - Read the supported formats from the backend so model construction and - kernel selection agree, including when AUTO resolves to TRTLLM. - """ - situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() - if moe_backend != "TRTLLM" or quant_algo in situ_supported: - return - supported = ", ".join(sorted(algo.name for algo in situ_supported)) - raise ValueError( - f"Kimi K3 routed experts are quantized as {quant_algo}, which the " - "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " - f"exist only for {supported}. Set moe_config.backend to CUTLASS " - "or MEGAMOE_CUTEDSL." - ) - - @staticmethod - def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: - """Build a private routed-expert mapping without mutating the shared - config, with the split of ``_select_moe_tp_ep``.""" - # Every backend here declares ``ActivationType.SiTu`` in its - # ``activation_support``; the list is not a preference order. CUTEDSL - # joined once its act-fusion kernel grew the SiTU epilogue. - supported_backends = { - "CUTLASS", - "TRTLLM", - "CUTEDSL", - "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - } - if model_config.moe_backend not in supported_backends: - raise ValueError( - "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " - "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " - f"got {model_config.moe_backend!r}." - ) - if model_config.moe_load_balancer is not None: - raise NotImplementedError( - "Kimi K3 packed-checkpoint streaming does not yet support " - "EPLB or replicated expert slots." - ) - mapping = model_config.mapping - - moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) - logger.info_once( - f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " - f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", - key="kimi_k3_moe_tp_ep_split", - ) - - mapping_dict = mapping.to_dict() - mapping_dict["moe_cluster_size"] = 1 - mapping_dict["moe_tp_size"] = moe_tp - mapping_dict["moe_ep_size"] = moe_ep - routed_mapping = Mapping.from_dict(mapping_dict) - - routed_model_config = copy.copy(model_config) - routed_model_config._frozen = False - routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) - routed_model_config.mapping = routed_mapping - routed_model_config.moe_backend = model_config.moe_backend - # MegaMoE uses this value as global DP SymmBuffer capacity, then - # divides it by EP size for the per-rank allocation. Other backends - # keep the user-configured value as their MoE chunking bound. - # Preserve an explicitly larger capacity. - if routed_model_config.moe_backend in { - "MEGAMOE_DEEPGEMM", - "MEGAMOE_CUTEDSL", - }: - default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size - configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) - if configured_moe_max_num_tokens < default_moe_max_num_tokens: - logger.info_once( - "Kimi K3 MegaMoE raises moe_max_num_tokens from " - f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " - "because the global DP SymmBuffer requires capacity for " - "max_num_tokens * dp_size.", - key=( - "kimi_k3_megamoe_capacity_override_" - f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" - ), - ) - routed_model_config.moe_max_num_tokens = max( - configured_moe_max_num_tokens, - default_moe_max_num_tokens, - ) - routed_model_config._frozen = True - return routed_model_config - - def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: - """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel - share (``partial_tail``): the decode path takes the step with that tail.""" - return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) - - def forward( - self, - hidden_states: torch.Tensor, - all_rank_num_tokens=None, - partial_tail: bool = False, - *, - step: DecodeStep, - ) -> Union[torch.Tensor, _decode_comm.PendingTail]: - """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification - (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where - ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at - most 8 tokens, a tensor on a wide decode step.""" - decode = self.decode_moe - if decode is not None and decode.takes(hidden_states, step, partial_tail): - return decode.forward( - self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push - ) - if partial_tail: - raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") - return self.forward_generic(hidden_states, all_rank_num_tokens) - - def forward_generic( - self, hidden_states: torch.Tensor, all_rank_num_tokens=None - ) -> torch.Tensor: - identity = hidden_states - router_logits = self.gate.compute_logits(hidden_states) - moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None +from tensorrt_llm.quantization.mode import QuantAlgo - def _routed_output(): - # Latent down/up projections via the min-latency fused GEMM op: - # at <=16 tokens (decode graphs) it runs a single pipelined - # bf16 kernel per projection instead of cuBLAS's split-K GEMV + - # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); - # for larger token counts the op falls back to cuBLAS internally. - # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear - # (A/B escape hatch). When the FP8 weight-read conversion has - # replaced the projection module, call it directly: its weight is - # an e4m3 buffer the bf16 dsv3 op must not read, and its forward - # is already a single fused GEMM (fp8_swap_ab_gemm). - routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) - y = self.routed_experts( - routed_in, - router_logits, - all_rank_num_tokens=all_rank_num_tokens, - ) - if self._reduce_routed_output: - return y - # Communication-backed paths return a complete routed result. - y = self.routed_expert_norm(y) - return self._routed_projection(y, self.routed_expert_up_proj) +from ..modeling import K3TargetSpec, ModelingV2KimiK3 +from ..modeling import _check_construction as _check_target_construction - # Shared experts depend only on the block input, so overlap their GEMMs - # with the routed dispatch/expert/combine chain. Multi-stream engages - # only under CUDA graphs; otherwise both branches run in order on the - # default stream. The shared GatedMLP includes its output all-reduce on - # the auxiliary stream. The join below must precede the routed - # all-reduce: concurrent collectives on different streams can corrupt - # SYMM_MEM all-reduce state. - routed_out, shared_out = maybe_execute_in_parallel( - _routed_output, - lambda: self.shared_experts(identity), - self.moe_main_event, - self.moe_shared_event, - self.shared_expert_stream, - disable_on_compile=True, - ) - if self._reduce_routed_output: - routed_latent = moe_all_reduce(routed_out) - routed_latent = self.routed_expert_norm(routed_latent) - routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) - return routed_out + shared_out - - -def resolve_attention_quant_config( - config: ModelConfig | None, layer_idx: int, projection: str -) -> QuantConfig: - """Resolve a checkpoint projection, including mixed-precision exclusions.""" - if config is None: - return QuantConfig() - global_config = config.quant_config or QuantConfig() - names = [ - f"{prefix}layers.{layer_idx}.self_attn.{projection}" - for prefix in ("language_model.model.", "model.", "") - ] - if any(global_config.is_module_excluded_from_quantization(name) for name in names): - return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) - declarations = config.quant_config_dict or {} - matches = [declarations[name] for name in names if name in declarations] - if matches: - selected = matches[0] - if any( - (match.quant_algo, match.group_size) != (selected.quant_algo, selected.group_size) - for match in matches[1:] - ): - raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") - elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: - selected = QuantConfig() - else: - selected = global_config - if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): - raise ValueError( - f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " - f"quantization {selected.quant_algo}" - ) - if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): - raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") - result = copy.copy(selected) - result.kv_cache_quant_algo = global_config.kv_cache_quant_algo - return result - - -class KimiMLARuntime(nn.Module): - """Wraps K3 MLA and applies its external TP output reduction.""" - - def __init__( - self, - cfg: "PretrainedConfig", - layer_idx: int, - model_config: ModelConfig, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ) -> None: - super().__init__() - - max_positions = int( - os.environ.get( - _KIMI_K3_MLA_MAX_POSITIONS_ENV, - cfg.max_position_embeddings, - ) - ) - self.layer_idx = layer_idx - # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the - # final output reduction in this wrapper so the output gate remains - # between attention and the row-parallel o_proj. - mapping = model_config.mapping - reduce_output = mapping.tp_size > 1 - self._o_allreduce = ( - AllReduce( - mapping=mapping, - strategy=model_config.allreduce_strategy, - dtype=torch.bfloat16, - ) - if reduce_output - else None - ) - attention_config = copy.copy(model_config) - attention_config._frozen = False - attention_config.quant_config_dict = { - name: resolve_attention_quant_config(model_config, layer_idx, name) - for name in ( - "q_a_proj", - "kv_a_proj_with_mqa", - "q_b_proj", - "kv_b_proj", - "g_proj", - "o_proj", - ) - } - attention_config.quant_config = QuantConfig( - kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo - if model_config.quant_config is not None - else None - ) - attention_config._frozen = model_config._frozen - self.mixer = K3DecodeMLA( - hidden_size=cfg.hidden_size, - num_heads=cfg.num_attention_heads, - q_lora_rank=cfg.q_lora_rank, - kv_lora_rank=cfg.kv_lora_rank, - qk_nope_head_dim=cfg.qk_nope_head_dim, - qk_rope_head_dim=cfg.qk_rope_head_dim, - v_head_dim=cfg.v_head_dim, - rms_norm_eps=cfg.rms_norm_eps, - dtype=torch.bfloat16, - layer_idx=layer_idx, - use_output_gate=cfg.mla_use_output_gate, - max_position_embeddings=max_positions, - model_config=attention_config, - aux_stream_dict=aux_stream_dict, - ) - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" - return self.mixer.will_run_decode_branch(attn_metadata, step) - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - step: DecodeStep, - reduce_output: bool = True, - project_output: bool = True, - ) -> torch.Tensor: - """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the - mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" - # MLA.forward takes position_ids first; K3 is NoPE, so pass None. - out = self.mixer( - None, hidden_states, attn_metadata, step=step, project_output=project_output - ) - return self._reduce(out) if project_output and reduce_output else out - - def forward_generic( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - reduce_output: bool = True, - ) -> torch.Tensor: - out = self.mixer.forward_generic(None, hidden_states, attn_metadata) - return self._reduce(out) if reduce_output else out - - def _reduce(self, out: torch.Tensor) -> torch.Tensor: - # Head-sharded TP: sum the row-sharded o_proj partials across - # the head-shard group. - return self._o_allreduce(out) if self._o_allreduce is not None else out - - -class KimiLinearDecoderLayer(nn.Module): - def __init__( - self, - model_config: ModelConfig, - cfg, - layer_idx: int, - aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - ): - super().__init__() - self.layer_idx = layer_idx - self.hidden_size = cfg.hidden_size - dtype = torch.bfloat16 - - self.is_kda = _is_kda_layer(cfg, layer_idx) - is_mla = _is_mla_layer(cfg, layer_idx) - if self.is_kda == is_mla: - raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") - - if self.is_kda: - projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") - attention_config = copy.copy(model_config) - attention_config._frozen = False - attention_config.quant_config_dict = { - name: resolve_attention_quant_config(model_config, layer_idx, name) - for name in projection_names - } - attention_config._frozen = model_config._frozen - self.linear_attn = K3DecodeKDA( - cfg, - layer_idx, - mapping=model_config.mapping, - allreduce_strategy=model_config.allreduce_strategy, - aux_stream=aux_stream_dict[AuxStreamType.Attention], - model_config=attention_config, - ) - else: - self.self_attn = KimiMLARuntime( - cfg, - layer_idx, - model_config=model_config, - aux_stream_dict=aux_stream_dict, - ) - - self.is_moe = ( - cfg.num_experts is not None - and layer_idx >= cfg.first_k_dense_replace - and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 - ) - if self.is_moe: - self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) - else: - situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) - # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's - # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take - # (decode_gemv.SITES); otherwise it stays within one node. - spans_nodes = self._mnnvl_allreduce() is not None - if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: - self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) - mlp_model_config = copy.copy(model_config) - mlp_model_config.quant_config = QuantConfig() - # K3's dense layer is BF16, so a unit block size gives the same - # subgroup selection as DeepSeek-V3. - self.mlp = GatedMLP( - hidden_size=cfg.hidden_size, - intermediate_size=cfg.intermediate_size, - bias=False, - activation=SituAndMul( - beta=situ_beta, - linear_beta=situ_linear_beta, - use_fused_activation=True, - ), - dtype=dtype, - config=mlp_model_config, - overridden_tp_size=self.mlp_tp_size, - reduce_output=self.mlp_tp_size > 1, - layer_idx=layer_idx, - ) - self._situ = (situ_beta, situ_linear_beta) - # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics - # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied - # after downcast, use_gemma=False). - self.input_layernorm = RMSNorm( - hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.post_attention_layernorm = RMSNorm( - hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - - # Attention residual scheme (always on for K3). The res norms stay - # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by - # _apply_attn_res and the fused attn_res op, never called as - # modules. - self.attn_res_block_size = cfg.attn_res_block_size - assert self.attn_res_block_size is not None, ( - "Kimi K3 runtime expects attn_res_block_size to be set" - ) - self.self_attention_res_norm = KimiK3RMSNorm( - cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) - self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set - # by the target's post_load_weights. - self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None - self.sandwich_oproj = False - - def forward( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - attn_metadata: AttentionMetadata, - step: DecodeStep, - capture: Optional[Tuple[Any, int]] = None, - prenormed: bool = False, - pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, - defer_moe_tail: bool = False, - ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: - """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). - - ``block_residual`` is a preallocated snapshot bank in kernel-native - ``[K_max, M, H]`` layout. Returns the running prefix sum and the - number of valid bank rows. - - ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux - stream for the layer BEFORE this one: the aggregated stream for layer j - is by definition what its next consumer sees, so the mixture computed - below already is it. Reading it here beats recomputing it, and is only - possible because K3 asserts pp_size == 1 -- layer j+1 is always local. - PP support would need a recompute at the rank boundary. - - ``step`` is the step's classification (``decode_step``), handed to the - attention module. - - ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is - this layer's input norm, and the layer's input, the step's embedding, - already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). - - The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens - runs the attention's all-reduce and the residual update in one - collective once the target has built its ``decode_comm``: with o_proj - (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, - the step and the attention's decode branch, else on the attention's - unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). - - ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the - previous layer's MoE output, unreduced (``defer_moe_tail``); then - ``hidden_states`` is the prefix sum without it, and this layer's - pre-attention step reduces and adds it: a ``PendingTail`` in one - ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its - all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer - the ``sandwich_tail`` call also stores the tap into the capture slot - where the speculative metadata exposes it (``capture_view``); other - tapped layers take the tap after the step. - - ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` - instead, ``partial`` this layer's MoE output unreduced, for the next - consumer to reduce. - """ - prefix_sum = hidden_states - valid_block_residual = block_residual[:num_snapshots] - attn_res_max_tokens = _attn_res_max_tokens(step) - tail = ( - pending_moe_partial - if isinstance(pending_moe_partial, _decode_comm.PendingTail) - else None - ) - # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum - # straight into the bank row it snapshots. - snapshot_row = None - if tail is not None and self.layer_idx % self.attn_res_block_size == 0: - snapshot_row = block_residual[num_snapshots] - # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the - # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). - tap_view = None - if tail is not None and capture is not None: - tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) - - if prenormed: - assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 - assert pending_moe_partial is None and capture is None - elif tail is not None: - hidden_states, prefix_sum = self.decode_comm.sandwich_tail( - tail, - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - updated_out=snapshot_row, - tap=tap_view, - tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, - ) - elif pending_moe_partial is not None: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - attn_res_max_tokens, - ) - elif capture is not None: - # The mixture tap needs the PRE-norm value, which the fused - # attn-res + RMSNorm kernel does not expose. Keep the two steps - # split on captured layers only and fuse everywhere else. - if num_snapshots > 0: - hidden_states = _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers - # acceptance, silently. hidden_states is the pre-norm attn_res mixture; - # prefix_only wants the running prefix, already in hand as prefix_sum. - tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - hidden_states = self.input_layernorm(hidden_states) - elif num_snapshots > 0: - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - attn_res_max_tokens, - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - if capture is not None and pending_moe_partial is not None and tap_view is None: - # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output - # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. - tapped = ( - _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - if _AUX_ATTN_RES_STREAM_ENABLED - else prefix_sum - ) - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - - if self.layer_idx % self.attn_res_block_size == 0: - if not prenormed and snapshot_row is None: - block_residual[num_snapshots].copy_(prefix_sum) - num_snapshots += 1 - valid_block_residual = block_residual[:num_snapshots] - prefix_sum = None - attention = self.linear_attn if self.is_kda else self.self_attn - comm = self.decode_comm - if comm is not None and comm.takes_post_attention(hidden_states, step): - if ( - self.sandwich_oproj - and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS - and attention.will_run_decode_branch(attn_metadata, step) - ): - core = attention(hidden_states, attn_metadata, step=step, project_output=False) - hidden_states, prefix_sum = comm.sandwich_oproj( - core, - self._o_proj().weight, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) - hidden_states, prefix_sum = comm.allreduce_attn_res( - partial, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - if comm is not None and step.wide: - partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) - attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) - else: - attention_output = attention(hidden_states, attn_metadata, step=step) - if prefix_sum is None: - prefix_sum = attention_output - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - attn_res_max_tokens, - ) - else: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - attention_output, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - attn_res_max_tokens, - ) - all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) - if self.is_moe and defer_moe_tail: - partial = self.block_sparse_moe( - hidden_states, all_rank_num_tokens, partial_tail=True, step=step - ) - return prefix_sum, num_snapshots, partial - if self.is_moe: - hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) - else: - hidden_states = self._dense_mlp(hidden_states, step) - - prefix_sum = prefix_sum + hidden_states - return prefix_sum, num_snapshots - - def forward_generic( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - attn_metadata: AttentionMetadata, - capture: Optional[Tuple[Any, int]] = None, - ) -> Tuple[torch.Tensor, int]: - prefix_sum = hidden_states - valid_block_residual = block_residual[:num_snapshots] - if capture is not None: - if num_snapshots > 0: - hidden_states = _apply_attn_res( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - ) - tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum - capture[0].maybe_capture_hidden_states(capture[1], tapped, None) - hidden_states = self.input_layernorm(hidden_states) - elif num_snapshots > 0: - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.self_attention_res_proj, - self.self_attention_res_norm, - self.input_layernorm, - ) - else: - hidden_states = self.input_layernorm(hidden_states) - - if self.layer_idx % self.attn_res_block_size == 0: - block_residual[num_snapshots].copy_(prefix_sum) - num_snapshots += 1 - valid_block_residual = block_residual[:num_snapshots] - prefix_sum = None - attention = self.linear_attn if self.is_kda else self.self_attn - comm = self.decode_comm - if comm is not None and comm.takes_post_attention(hidden_states, None): - partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) - hidden_states, prefix_sum = comm.allreduce_attn_res( - partial, - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - elif prefix_sum is None: - prefix_sum = attention.forward_generic(hidden_states, attn_metadata) - hidden_states = _apply_attn_res_and_rmsnorm( - prefix_sum, - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - else: - prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( - prefix_sum, - attention.forward_generic(hidden_states, attn_metadata), - valid_block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - self.post_attention_layernorm, - ) - if self.is_moe: - all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) - hidden_states = self.block_sparse_moe.forward_generic( - hidden_states, all_rank_num_tokens - ) - else: - hidden_states = self.mlp(hidden_states) - return prefix_sum + hidden_states, num_snapshots - - def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: - """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode - kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" - gemvs = self.decode_gemvs - if gemvs is not None and step.small: - out = gemvs.dense_mlp( - hidden_states, - self.mlp.gate_up_proj.weight, - self.mlp.down_proj.weight, - *self._situ, - ) - if out is not None: - return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out - return self.mlp(hidden_states) - - def _mnnvl_allreduce(self): - """The MNNVL all-reduce of this layer's attention output, or None.""" - attention = self.linear_attn if self.is_kda else self.self_attn - return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) - - def _o_proj(self) -> nn.Module: - """This layer's attention output projection (row parallel).""" - return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj - - def _o_allreduce(self) -> AllReduce: - """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" - return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce - - def accepts_moe_partial(self, num_snapshots: int) -> bool: - """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target - built its collective state, and the snapshot bank is not empty.""" - return num_snapshots > 0 and self.decode_comm is not None - - def skip_forward( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. - - ``modeling_utils.skip_forward()`` only drops a module's weights when it - finds this attribute, so without it the layer-wise benchmarks would - allocate all 93 layers instead of the profiled slice. - """ - return hidden_states, block_residual - - -class KimiLinearModel(DecoderModel): - def __init__(self, model_config: ModelConfig): - super().__init__(model_config) - cfg = _get_text_config(model_config.pretrained_config) - self._text_cfg = cfg - dtype = torch.bfloat16 - - # Attention and MoE phases are sequential, so their branch-overlap - # roles share one stream; MoE-internal overlap roles remain separate. - aux_stream_list = [torch.cuda.Stream() for _ in range(4)] - self.aux_stream_dict = { - AuxStreamType.Attention: aux_stream_list[0], - AuxStreamType.MoeShared: aux_stream_list[0], - AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], - AuxStreamType.MoeBalancer: aux_stream_list[2], - AuxStreamType.MoeOutputMemset: aux_stream_list[3], - } - - self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) - self.layers = nn.ModuleList( - [ - KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) - for layer_idx in range(cfg.num_hidden_layers) - ] - ) - self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) - - # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) - # by _apply_attn_res and the fused attn_res op. - self.output_attn_res_norm = KimiK3RMSNorm( - cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype - ) - self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) - self.num_attn_res_snapshots = ( - cfg.num_hidden_layers + cfg.attn_res_block_size - 1 - ) // cfg.attn_res_block_size - # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the - # weights are final. None: every step embeds on the generic path. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - - # Which convention the drafter tap is on is not recoverable from the - # served output -- a mismatch only lowers acceptance -- so state it once - # at construction rather than leaving it to be inferred from an AL. - logger.info_once( - "Kimi K3 aux hidden capture: mode=" - f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " - f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", - key="kimi_k3_aux_capture_mode", - ) - - @property - def kda_draft_records(self) -> bool: - """Whether the hybrid cache manager keeps the last verify's draft records (``kda_state_tok``), from which - ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn`` replay the accepted drafts. The engine reads it once the weights - are loaded, to build the manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer - taking the K3 kernels. Otherwise the KDA verify replays the accepted drafts from the replay caches (the - built-in verify).""" - spec_config = getattr(self.model_config, "spec_config", None) - return bool( - spec_config is not None - and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) - and spec_config.tokens_per_gen_step in (2, 4, 6, 8) - and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) - ) - - def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, - from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from - load-time state, so every rank and every replay of the graph decides alike.""" - if not torch.cuda.is_current_stream_capturing(): - return False - kda = [layer.linear_attn for layer in self.layers if layer.is_kda] - mla = [layer.self_attn for layer in self.layers if not layer.is_kda] - return latent_push( - step, - capturing=True, - breakable=is_in_breakable_cuda_graph(), - kda_draft_records=self.kda_draft_records, - kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), - mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), - ) - - def _defer_moe_tail( - self, - i: int, - hidden_states: torch.Tensor, - num_snapshots: int, - spec_metadata, - capture_set, - step: DecodeStep, - ) -> bool: - """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE - takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after - the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps - its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated - tail.""" - layer = self.layers[i] - if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): - return False - if spec_metadata is not None and ( - capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) - ): - return False - return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) - - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - spec_metadata=None, - **kwargs, - ) -> torch.Tensor: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] - step = decode_step(attn_metadata, num_tokens) - if step is None: - return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) - if self._latent_push(attn_metadata, step): - step = replace(step, latent_push=True) - # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. - prenormed = None - if ( - inputs_embeds is None - and self.decode_gemvs is not None - and len(self.layers) > 0 - and self.num_attn_res_snapshots > 0 - ): - table = self.embed_tokens.weight - block_residual = table.new_empty( - self.num_attn_res_snapshots, num_tokens, table.shape[1] - ) - prenormed = self.decode_gemvs.embed_norm( - input_ids, table, self.layers[0].input_layernorm, block_residual - ) - if prenormed is not None: - hidden_states = prenormed - else: - hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - block_residual = hidden_states.new_empty( - self.num_attn_res_snapshots, - hidden_states.shape[0], - hidden_states.shape[1], - ) - num_snapshots = 0 - capture_set = getattr(spec_metadata, "_capture_layer_set", None) - # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) - # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. - pending_moe_partial = None - for i, layer in enumerate(self.layers): - defer_moe_tail = self._defer_moe_tail( - i, hidden_states, num_snapshots, spec_metadata, capture_set, step - ) - outputs = layer( - hidden_states, - block_residual, - num_snapshots, - attn_metadata, - step=step, - capture=self._capture(spec_metadata, capture_set, i), - prenormed=i == 0 and prenormed is not None, - pending_moe_partial=pending_moe_partial, - defer_moe_tail=defer_moe_tail, - ) - if defer_moe_tail: - hidden_states, num_snapshots, pending_moe_partial = outputs - else: - (hidden_states, num_snapshots), pending_moe_partial = outputs, None - - if isinstance(pending_moe_partial, _decode_comm.PendingTail): - normed, _ = self.layers[-1].decode_comm.sandwich_tail( - pending_moe_partial, - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - ) - return normed - if pending_moe_partial is not None: - _, normed = _apply_attn_res_add_and_rmsnorm( - hidden_states, - _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - _attn_res_max_tokens(step), - ) - return normed - return self._final_norm( - hidden_states, - block_residual, - num_snapshots, - spec_metadata, - capture_set, - _attn_res_max_tokens(step), - ) - - def _generic_forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor], - inputs_embeds: Optional[torch.FloatTensor], - spec_metadata, - ) -> torch.Tensor: - hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds - block_residual = hidden_states.new_empty( - self.num_attn_res_snapshots, - hidden_states.shape[0], - hidden_states.shape[1], - ) - num_snapshots = 0 - capture_set = getattr(spec_metadata, "_capture_layer_set", None) - for i, layer in enumerate(self.layers): - hidden_states, num_snapshots = layer.forward_generic( - hidden_states, - block_residual, - num_snapshots, - attn_metadata, - capture=self._capture(spec_metadata, capture_set, i), - ) - return self._final_norm( - hidden_states, block_residual, num_snapshots, spec_metadata, capture_set - ) - - def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: - # DFlash/DSpark hidden-state capture. The drafter is distilled on - # the aggregated stream value -- the pre-norm softmax mixture its - # next consumer sees -- not on the raw prefix sum a layer returns, - # which is SGLang's fallback for models without the - # attention-residual scheme. Capturing the prefix sum costs 4.5pt - # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). - # The tap fires inside layer i+1, which computes that tensor - # anyway; see its forward docstring. Ground truth: SGLang - # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 - # aggregate_stream_torch. - if ( - spec_metadata is not None - and i > 0 - and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) - ): - return (spec_metadata, self.layers[i - 1].layer_idx) - return None - - def _final_norm( - self, - hidden_states: torch.Tensor, - block_residual: torch.Tensor, - num_snapshots: int, - spec_metadata, - capture_set, - max_fused_tokens: Optional[int] = None, - ) -> torch.Tensor: - # The last layer has no successor, so this one recompute is - # unavoidable -- output-side score weights, matching SGLang's - # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set - # against 93 layers; kept so a set that does include the final layer - # gets the right tensor rather than the raw prefix sum. - if spec_metadata is not None and len(self.layers) > 0: - last = self.layers[-1] - if capture_set is None or last.layer_idx in capture_set: - tail = ( - _apply_attn_res( - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - ) - if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED - else hidden_states - ) - spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) - - return _apply_attn_res_and_rmsnorm( - hidden_states, - block_residual[:num_snapshots], - self.output_attn_res_proj, - self.output_attn_res_norm, - self.norm, - max_fused_tokens, - ) - - -# ---------------------------------------------------------------------------------------------------------------------- -# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. -# ---------------------------------------------------------------------------------------------------------------------- - - -class K3DecodeKDA(KimiKDALinearAttention): - """Quantized KDA with the target's output-reduction boundary.""" - - takes_k3_kernels = False - k3_buffers = None - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - return False - - def forward(self, hidden_states, attn_metadata, step, reduce_output=True, project_output=True): - if not project_output: - raise ValueError("Quantized KDA requires its own output projection") - return self.forward_generic(hidden_states, attn_metadata, reduce_output) - - def forward_generic(self, hidden_states, attn_metadata, reduce_output=True): - if reduce_output: - return super().forward(hidden_states, attn_metadata) - if self.register_to_config and is_in_breakable_cuda_graph(): - core = hidden_states.new_empty( - (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 - ) - maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) - else: - core = self._forward_impl(hidden_states, attn_metadata) - return self.o_proj(core.reshape(-1, self.proj_size)) - - -class K3DecodeMLA(KimiK3MLAAttention): - """Stock quantized MLA, including its KV-cache format handling.""" - - def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: - return False - - def forward(self, position_ids, hidden_states, attn_metadata, step, project_output=True): - if not project_output: - raise ValueError("Quantized MLA requires its own output projection") - return self.forward_generic(position_ids, hidden_states, attn_metadata) - - def forward_generic(self, position_ids, hidden_states, attn_metadata): - return super().forward(position_ids, hidden_states, attn_metadata) - - -# ---------------------------------------------------------------------------------------------------------------------- -# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. -# ---------------------------------------------------------------------------------------------------------------------- - -#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at -#: max_draft_len 7. -DRAFTER_ATTN_SPLITS = frozenset( - {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} -) -# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k -# RMSNorm's epsilon and the NeoX RoPE's base. -_DRAFTER_HEADS = (6, 1) -_DRAFTER_HEAD_DIM = 64 -_DRAFTER_PAGE = 64 -_DRAFTER_EPS = 1e-5 -_DRAFTER_ROPE_BASE = 10000.0 - - -def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: - """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input - feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" - if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: - return None - width = in_features // tp_size - return tp_rank * width, (tp_rank + 1) * width - - -class K3FcSlice(nn.Module): - """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is - ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's - ranks is ``fc``'s output.""" - - def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: - super().__init__() - self.weight = nn.Parameter(weight, requires_grad=False) - self.start = start - self.end = end - - @classmethod - def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: - """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns - (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" - columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) - if columns is None: - return None - start, end = columns - return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the - GEMM reads the column block in place, through the rows' stride.""" - return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) - - -class K3DSparkDrafter(GQADSparkForCausalLM): - """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. - - A block the entries take runs, per layer: - - * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; - * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to - its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the - worker writes the accepted rows' context k / v before a block reads them; - * the output projection on ``drafter_o``, then the module's all-reduce; - * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the - module's all-reduce. - - The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state - (below); a projection whose site does not take its rows runs its module. Every other block runs the stock - ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or - normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under - CUDA-graph capture. - - Where the target hands the drafter the TP group's collective state (``use_decode_comm``): - - * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and - ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; - * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the - post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the - block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the - projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection - after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is - followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers - whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. - - Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the - stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); - the context k / v and the Markov head stay upstream code. - """ - - def __init__( - self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" - ) -> None: - super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) - # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. - self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None - # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. - self._k3_take: Optional[bool] = None - # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes - # only these. - self._k3_attn_ran: set = set() - # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). - self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None - # The zero residual of the split context projection's all-reduce, up to a decode step's rows. - self._k3_zero_rows: Optional[torch.Tensor] = None - # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms - # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. - self._k3_norms_fuse = False - self._k3_sandwich_forms: frozenset = frozenset() - - def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: - """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the - target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are - loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. - - * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they - split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` - sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. - * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and - RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose - shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not - compile is not used.""" - self.decode_comm = comm - self._k3_split_fc() - self._k3_norms_fuse = self._k3_norms_take_comm() - forms = set() - if self._k3_norms_fuse: - layers = self.model.layers - for form, swiglu, weights in ( - ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), - ("down", True, [layer.mlp.down_proj.weight for layer in layers]), - ): - if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( - weights[0], swiglu=swiglu - ): - forms.add(form) - self._k3_sandwich_forms = frozenset(forms) - logger.info( - "Kimi K3 DSpark drafter: residual adds and RMSNorms " - + ( - f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " - "mnnvl_fusion_allreduce)" - if self._k3_norms_fuse - else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" - ) - ) - - def _k3_norms_take_comm(self) -> bool: - """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and - post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no - quantized output), and every output and down projection is row parallel with its all-reduce.""" - hidden = self.config.hidden_size - - def plain(norm: nn.Module) -> bool: - weight = getattr(norm, "weight", None) - return ( - isinstance(weight, torch.Tensor) - and weight.dtype == torch.bfloat16 - and tuple(weight.shape) == (hidden,) - and hasattr(norm, "variance_epsilon") - and not getattr(norm, "use_gemma", True) - and not getattr(norm, "is_nvfp4", True) - and not getattr(norm, "return_hp_output", True) - ) - - def reduces(linear: nn.Module) -> bool: - return ( - getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" - and getattr(linear, "reduce_output", False) - and getattr(linear, "all_reduce", None) is not None - ) - - return plain(self.model.norm) and all( - plain(layer.input_layernorm) - and plain(layer.post_attention_layernorm) - and reduces(layer.self_attn.o_proj) - and reduces(layer.mlp.down_proj) - for layer in self.model.layers - ) - - def _k3_tp_all_reduce(self) -> Optional[nn.Module]: - """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" - o_proj = self.model.layers[0].self_attn.o_proj - if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": - return None - return getattr(o_proj, "all_reduce", None) - - def _k3_split_fc(self) -> None: - """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero - residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is - a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP - all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" - fc = getattr(self, "fc", None) - if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): - return - mapping = self.model_config.mapping - weight = getattr(fc, "weight", None) - sliced = None - if ( - isinstance(weight, torch.Tensor) - and weight.dim() == 2 - and weight.dtype == torch.bfloat16 - and getattr(fc, "bias", None) is None - and self._k3_tp_all_reduce() is not None - ): - sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) - if sliced is None: - logger.info( - "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " - f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" - ) - return - self.fc = sliced - self._k3_zero_rows = weight.new_zeros( - MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] - ) - logger.info( - f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " - f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" - ) - - def load_weights(self, weights, weight_mapper=None, **kwargs): - """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" - result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) - self._k3_split_fc() - return result - - def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. - - With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group - with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode - step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then - ``hidden_norm``. Otherwise the stock projection.""" - fc = self.fc - if not isinstance(fc, K3FcSlice): - return super().project_target_hidden(hidden_states) - partial = fc(hidden_states.to(fc.weight.dtype)) - rows = partial.shape[0] - zeros = self._k3_zero_rows - if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): - normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) - return normed - if rows == 0: - return self.hidden_norm(partial) - return self.hidden_norm(self._k3_tp_all_reduce()(partial)) - - def dflash_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_k_cache: torch.Tensor, - ctx_v_cache: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor] = None, - ctx_page_table: Optional[torch.Tensor] = None, - ctx_rows_start: Optional[int] = None, - ) -> torch.Tensor: - """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class - docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows - ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows - as a view instead of gathering them.""" - keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) - if keys is None: - return super().dflash_forward( - noise_embedding, - query_positions, - num_ctx_per_req, - ctx_k_cache, - ctx_v_cache, - ctx_cache_batch_idx, - ctx_kv_cache, - ctx_page_table, - ) - out = self._k3_block_forward( - noise_embedding, - query_positions, - num_ctx_per_req, - ctx_cache_batch_idx, - ctx_kv_cache, - ctx_page_table, - ctx_rows_start, - ) - if not torch.cuda.is_current_stream_capturing(): - self._k3_attn_ran |= keys - return out - - def _k3_layers_take(self) -> bool: - """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head - layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, - non-causal unwindowed block attention, and no block convolution or post-attention gate.""" - if self._k3_take is None: - if self._fused_kv_weight is None: - self._build_fused_kv_buffers() - take = ( - self._use_fused_qk_norm_rope - and not self.has_block_conv - and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate - ) - for layer_idx, layer in enumerate(self.model.layers): - attn, mlp = layer.self_attn, layer.mlp - pos = getattr(attn, "pos_embd_params", None) - rope = getattr(pos, "rope", None) - norms = (attn.q_norm, attn.k_norm) - take = take and ( - (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS - and attn.head_dim == _DRAFTER_HEAD_DIM - and getattr(attn, "is_qk_norm", False) - and not getattr(attn, "use_gemma_rms_norm", True) - and rope is not None - and pos.is_neox - and getattr(pos, "mrope_section", None) is None - and rope.theta == _DRAFTER_ROPE_BASE - and all( - norm.variance_epsilon == _DRAFTER_EPS - and norm.weight.dtype == torch.bfloat16 - and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) - for norm in norms - ) - and attn.qkv_proj.bias is None - and attn.o_proj.bias is None - and mlp.activation is torch.nn.functional.silu - and mlp.swiglu_limit is None - and mlp.swiglu_alpha in (None, 1.0) - and mlp.swiglu_beta in (None, 0.0) - and mlp.gate_up_proj.bias is None - and mlp.down_proj.bias is None - and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) - ) - self._k3_take = bool(take) - logger.info( - f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " - + ( - "on k3_drafter_attn_qknorm and the drafter GEMV sites" - if self._k3_take - else "on the stock block forward (a layer the drafter entries do not take)" - ) - ) - return self._k3_take - - def _k3_block_keys( - self, - noise_embedding: torch.Tensor, - ctx_kv_cache: Optional[torch.Tensor], - ctx_page_table: Optional[torch.Tensor], - ) -> Optional[frozenset]: - """The attention's compile keys for this block when the drafter entries take it, else None.""" - if ( - self.dflash_attention_backend != "TRTLLM" - or ctx_kv_cache is None - or ctx_page_table is None - or noise_embedding.dtype != torch.bfloat16 - or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS - or is_in_breakable_cuda_graph() - or not self._k3_layers_take() - ): - return None - num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads - page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) - dense = ( - num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, - _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, - ) - strides = set() - for layer_idx in range(len(self.model.layers)): - cache = ctx_kv_cache[layer_idx] - if ( - cache.dtype != torch.bfloat16 - or cache.dim() != 5 - or tuple(cache.shape[1:]) != page - or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) - or cache.stride(0) % 8 - ): - return None - strides.add(cache.stride(0)) - keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) - if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: - return None - return keys - - def _k3_block_forward( - self, - noise_embedding: torch.Tensor, - query_positions: torch.Tensor, - num_ctx_per_req: torch.Tensor, - ctx_cache_batch_idx: torch.Tensor, - ctx_kv_cache: torch.Tensor, - ctx_page_table: torch.Tensor, - ctx_rows_start: Optional[int] = None, - ) -> torch.Tensor: - batch, block = noise_embedding.shape[:2] - rows = batch * block - ctx_len = num_ctx_per_req[:batch].to(torch.int32) - if ctx_rows_start is None: - page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) - else: - # The batch's page-table rows are one contiguous run: a view, no gather. - page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] - positions = query_positions.reshape(-1).contiguous() - hidden = noise_embedding.reshape(rows, -1) - layers = self.model.layers - comm = self._k3_fused_comm(hidden) - residual = None - if comm is not None: - # The fused all-reduces read the residual and return the updated one: the block's input serves as the - # first residual, uncopied. - residual = hidden - normed = layers[0].input_layernorm(hidden) - for layer_idx, layer in enumerate(layers): - attn = layer.self_attn - if comm is None: - if residual is None: - residual = hidden.clone() - normed = layer.input_layernorm(hidden) - else: - normed, residual = layer.input_layernorm(hidden, residual) - qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) - out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) - k3_drafter_attn_qknorm( - qkv, - attn.q_norm.weight, - attn.k_norm.weight, - positions, - attn.q_norm.variance_epsilon, - attn.pos_embd_params.rope.theta, - ctx_kv_cache[layer_idx], - page_table, - ctx_len, - attn.num_heads, - attn.num_key_value_heads, - out, - ) - if comm is not None: - # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after - # the last layer the final norm. - next_norm = ( - layers[layer_idx + 1].input_layernorm - if layer_idx + 1 < len(layers) - else self.model.norm - ) - normed, residual = self._k3_project_norm( - comm, out, attn.o_proj, residual, layer.post_attention_layernorm - ) - normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) - continue - hidden = self._k3_project("drafter_o", out, attn.o_proj) - hidden, residual = layer.post_attention_layernorm(hidden, residual) - hidden = self._k3_mlp(layer.mlp, hidden) - if comm is not None: - return normed - out, _ = self.model.norm(hidden, residual) - return out - - def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: - """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its - all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace - holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" - comm = self.decode_comm - if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): - return None - return comm - - def _k3_project_norm( - self, - comm: _decode_comm.K3DecodeComm, - x: torch.Tensor, - linear: nn.Module, - residual: torch.Tensor, - norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, - the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the - projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` - where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" - if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( - x, linear.weight, residual, norm - ): - return comm.sandwich_plain(x, linear.weight, residual, norm) - gemvs = self.decode_gemvs - partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) - if partial is None: - partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) - return comm.allreduce_norm(partial, residual, norm) - - def _k3_mlp_norm( - self, - comm: _decode_comm.K3DecodeComm, - mlp: nn.Module, - x: torch.Tensor, - residual: torch.Tensor, - norm: nn.Module, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down - projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul - form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down - projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; - then ``comm/mnnvl_fusion_allreduce``.""" - gemvs = self.decode_gemvs - gate_up = ( - None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) - ) - partial = None - if gate_up is not None: - if "down" in self._k3_sandwich_forms and comm.takes_plain( - gate_up, mlp.down_proj.weight, residual, norm, swiglu=True - ): - return comm.sandwich_plain( - gate_up, mlp.down_proj.weight, residual, norm, swiglu=True - ) - partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) - if partial is None: - partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) - return comm.allreduce_norm(partial, residual, norm) - - def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: - """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel - projection's all-reduce.""" - y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) - if y is None: - return linear(x) - return _row_parallel_reduce(linear, y) - - def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: - """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on - ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" - gemvs = self.decode_gemvs - gate_up = ( - None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) - ) - down = ( - None - if gate_up is None - else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) - ) - if down is None: - return mlp(x) - return _row_parallel_reduce(mlp.down_proj, down) - - -def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: - """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" - all_reduce = getattr(linear, "all_reduce", None) - if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: - return all_reduce(partial) - return partial - - -# ---------------------------------------------------------------------------------------------------------------------- -# The target: step classification, the construction checks and the class route returns. -# ---------------------------------------------------------------------------------------------------------------------- - - -def _text_model_config(model_config: ModelConfig) -> ModelConfig: - """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. - - The checkpoint names its language-model modules `language_model.`; the text model's are ``, with - `layers.*` under `model.`. - """ - config = model_config.pretrained_config - text = copy.copy(model_config) - text._frozen = False - text.pretrained_config = config.text_config - excluded = text.quant_config.exclude_modules - if excluded: - text.quant_config = copy.copy(text.quant_config) - renamed = [] - for name in excluded: - if name.startswith(_LANG_PREFIX): - name = name[len(_LANG_PREFIX) :] - if name.startswith("layers."): - name = "model." + name - renamed.append(name) - text.quant_config.exclude_modules = renamed - text.skip_create_weights_in_init = True - text._frozen = True - return text - - -@dataclass(frozen=True) -class DecodeStep: - """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` - generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" - - num_tokens: int - num_requests: Optional[int] = None - tokens_per_request: Optional[int] = None - #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` - #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. - latent_push: bool = False - - @property - def small(self) -> bool: - """Whether the step fits one token tile of the token-count kernels.""" - return self.num_tokens <= DECODE_MAX_TOKENS - - @property - def decode(self) -> bool: - """Whether the step is a pure decode step the request-aware kernels take.""" - return self.num_requests is not None - - @property - def wide(self) -> bool: - """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode - layout's MoE head and tail, on M-general ops.""" - return self.decode and not self.small - - -def latent_push( - step: DecodeStep, - *, - capturing: bool, - breakable: bool, - kda_draft_records: bool, - kda_decode_kernels: bool, - mla_decode_branch: bool, -) -> bool: - """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at - most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a - breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels - (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA - layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). - Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel - between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a - graph launch orders every pushing replay behind whatever ran before it.""" - if not (step.decode and step.small) or not capturing or breakable: - return False - if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): - return False - return mla_decode_branch - - -def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: - """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). - - ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers - and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input - here is fixed by that shape, so a captured step and its replays are classified alike. - """ - if num_tokens <= 0: - return None - requests = _decode_requests(attn_metadata, num_tokens) - if requests is not None: - return DecodeStep(num_tokens, requests, num_tokens // requests) - if num_tokens <= DECODE_MAX_TOKENS: - return DecodeStep(num_tokens) - return None - - -def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: - """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" - if attn_metadata.num_contexts != 0: - return None - requests = attn_metadata.num_generations - if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: - return None - tokens = num_tokens // requests - if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: - return None - seq_lens = getattr(attn_metadata, "seq_lens", None) - if seq_lens is not None and seq_lens.device.type == "cpu": - lens = seq_lens[:requests] - if lens.numel() != requests or bool((lens != tokens).any()): - return None - return requests +_TARGET = K3TargetSpec(sm=(10, 7), quant_algo=QuantAlgo.NVFP4, moe_tp=4, moe_ep=4) def _check_construction(model_config: ModelConfig) -> None: - """The settings this target is built for that are fixed before the first step.""" - capability = torch.cuda.get_device_capability() - assert capability == _SM, ( - f"this target requires sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" - ) - mapping = model_config.mapping - topology = ( - mapping.world_size, - mapping.tp_size, - mapping.pp_size, - mapping.moe_tp_size, - mapping.moe_ep_size, - mapping.enable_attention_dp, - ) - assert topology == (16, 16, 1, 4, 4, False), ( - "the tp16_moetp4ep4 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " - "moe_tensor_parallel_size 4, moe_expert_parallel_size 4 and enable_attention_dp false; the engine built " - f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" - ) - assert model_config.torch_dtype == torch.bfloat16, ( - f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" - ) - for layer_idx in range(1, 93): - quant = KimiK3MoERuntime._resolve_routed_quant_config(model_config, layer_idx) - if quant.quant_algo != QuantAlgo.NVFP4 or quant.group_size != 16: - raise ValueError( - f"NVFP4 target requires group-size-16 NVFP4 experts at layer {layer_idx}; " - f"got {quant.quant_algo}, group_size={quant.group_size}" - ) - strategy = model_config.allreduce_strategy - assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( - f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" - ) - - -class ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4(KimiLinearForCausalLM): - """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. - - It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through - weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model - keeps the built-in one's. - """ - - @classmethod - def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: - """The V2 hybrid manager holds the KDA states; the step contract requires it.""" - return "V2" - - def __init__(self, model_config: ModelConfig): - config = model_config.pretrained_config - assert getattr(config, "text_config", None) is not None, ( - "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " - "text_config" - ) - _check_construction(model_config) - text = _text_model_config(model_config) - spec_config = getattr(text, "spec_config", None) - assert ( - spec_config is None - or spec_config.spec_dec_mode.is_sa() - or spec_config.spec_dec_mode.is_dflash() - or spec_config.spec_dec_mode.is_dspark() - ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" - # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 - # weight-read conversion of the shared / latent MLPs. - self._fp8_weight_read_moe_mlp = False - self.mapping_with_cp = None - self._repurposed_tp_mapping = None - cfg = text.pretrained_config - SpecDecOneEngineForCausalLM.__init__( - self, - KimiLinearModel(text), - text, - hidden_size=cfg.hidden_size, - vocab_size=cfg.vocab_size, - ) - self._own_spec_branch(text) - self._step_checked = False - # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the - # processor the speculative worker's target logits and a parallel drafter's logits call too. - stock_logits_processor = self.logits_processor - self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) - draft_model = getattr(self, "draft_model", None) - if getattr(draft_model, "logits_processor", None) is stock_logits_processor: - draft_model.logits_processor = self.logits_processor - # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which - # must therefore be the text config, as the built-in wrapper leaves it. - model_config._frozen = False - model_config.pretrained_config = self.config - model_config._frozen = True - - def _own_spec_branch(self, model_config: ModelConfig) -> None: - """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a - standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is - `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" - spec_config = model_config.spec_config - if spec_config is None: - return - mode = spec_config.spec_dec_mode - standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) - if ( - standalone_dspark - and self.draft_config is not None - and is_mla(self.draft_config.pretrained_config) - ): - # The reference NVFP4 checkpoint uses an MLA DSpark drafter. Its - # compressed context cache does not meet the GQA worker contract. - return - if ( - standalone_dspark - and self.draft_config is not None - and not is_mla(self.draft_config.pretrained_config) - ): - drafter = K3DSparkDrafter( - self.draft_config, dflash_attention_backend=spec_config.attention_backend - ) - if mode.is_parallel_draft(): - drafter.logits_processor = self.logits_processor - self.epilogue[self.epilogue.index(self.draft_model)] = drafter - self.draft_model = drafter - if mode.is_dflash(): - worker_cls = _spec_worker.KimiK3DFlashWorker - elif standalone_dspark: - worker_cls = _spec_worker.KimiK3DSparkWorker - else: - return - worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) - worker.set_draft_model(self.draft_model) - self.epilogue[self.epilogue.index(self.spec_worker)] = worker - self.spec_worker = worker - - def load_weights(self, weights, *args, **kwargs): - _weights.load(self, weights) - - def cache_derived_state(self) -> None: - """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of - every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: - a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites - are among them, and the drafter gets the state too.""" - super().cache_derived_state() - if self.model.decode_gemvs is not None: - return - sites = ("dense_gate_up", "dense_down") - if isinstance(self.draft_model, K3DSparkDrafter): - sites += ("drafter_qkv", "drafter_o", "drafter_gate_up", "drafter_down") - gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head, sites=sites) - self.model.decode_gemvs = gemvs - self.logits_processor.gemvs = gemvs - for layer in self.model.layers: - if not layer.is_moe: - layer.decode_gemvs = gemvs - if isinstance(self.draft_model, K3DSparkDrafter): - self.draft_model.decode_gemvs = gemvs - - def post_load_weights(self) -> None: - """Prepare BF16 decode state and MNNVL collectives before graph capture.""" - super().post_load_weights() - if getattr(self, "_nvfp4_decode_state_ready", False): - return - layers = self.model.layers - comm = None - if all(layer._mnnvl_allreduce() is not None for layer in layers): - comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, None) - for layer in layers: - layer.decode_comm = comm - layer.sandwich_oproj = False - _decode_comm.use_decode_one_shot(self) - self._gate_spec_worker_kernels(comm) - self._gate_drafter_comm(comm) - self._nvfp4_decode_state_ready = True - logger.info( - "Kimi K3 NVFP4 SM107: stock quantized attention and NVFP4 experts; " - "BF16 decode GEMVs and DSpark specializations; " - f"attention-residual MNNVL fusion={comm is not None}" - ) - - def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: - """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, - ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input - that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's - ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the - vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks - each step's own conditions. Returns the setting (False without such a worker).""" - worker = getattr(self, "spec_worker", None) - if not hasattr(worker, "k3_decode"): - return False - gemvs = self.model.decode_gemvs - worker.k3_decode = ( - comm is not None and gemvs is not None and gemvs.head_workspace is not None - ) - logger.info( - f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " - + ( - "on" - if worker.k3_decode - else "off (no MNNVL decode state or no k3_head_gemv LM head)" - ) - ) - return worker.k3_decode + _check_target_construction(model_config, _TARGET) - def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: - """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built - it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context - projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms - in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on - the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate - requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` - workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the - drafter took the state (False without such a drafter).""" - drafter = getattr(self, "draft_model", None) - if comm is None or not isinstance(drafter, K3DSparkDrafter): - return False - drafter.use_decode_comm(comm) - return True - def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: - """First-forward checks of the engine surface and the per-engine settings.""" - objects = { - "attn_metadata": attn_metadata, - "kv_cache_manager": attn_metadata.kv_cache_manager, - } - missing = [ - f"{owner}.{name}" - for owner, names in REQUIRED_ENGINE_FIELDS.items() - for name in names - if not hasattr(objects[owner], name) - ] - assert not missing, f"engine fields this target reads are missing: {missing}" - manager = attn_metadata.kv_cache_manager - assert isinstance(manager, MambaHybridCacheManagerV2), ( - "this target needs the V2 hybrid KV / state cache manager " - "(kv_cache_config.use_kv_cache_manager_v2); the engine built " - f"{type(manager).__name__}" - ) - self._step_checked = True +class ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4(ModelingV2KimiK3): + """Concrete deployment of the shared Kimi K3 text model.""" - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - return_context_logits: bool = False, - spec_metadata=None, - resource_manager=None, - **kwargs, - ) -> torch.Tensor: - if kwargs.pop("multimodal_params", None): - raise ValueError( - "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" - ) - if not self._step_checked: - self._check_step_contract(attn_metadata) - return super().forward( - attn_metadata, - input_ids, - position_ids, - inputs_embeds, - return_context_logits, - spec_metadata, - resource_manager, - **kwargs, - ) + target_spec = _TARGET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py deleted file mode 100644 index abc1a1388213..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/spec_worker.py +++ /dev/null @@ -1,1343 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -import functools -import inspect -from typing import Optional - -import torch -from torch import nn - -from tensorrt_llm._torch.attention.backends import AttentionMetadata -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager -from tensorrt_llm._torch.speculative.dflash import ( - DFlashSpecMetadata, - DFlashWorker, - dflash_allocated_ctx_limit, - dflash_noise_block_embedding, -) -from tensorrt_llm._torch.speculative.dspark import DSparkWorker -from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase -from tensorrt_llm.logger import logger - - -@functools.lru_cache(maxsize=None) -def _takes_ctx_rows_start_cls(cls) -> bool: - return "ctx_rows_start" in inspect.signature(cls.dflash_forward).parameters - - -def _takes_ctx_rows_start(draft_model) -> bool: - """Whether the draft model's ``dflash_forward`` accepts ``ctx_rows_start``.""" - return _takes_ctx_rows_start_cls(type(draft_model)) - - -def dflash_ctx_rows_kwargs(draft_model, ctx_rows_start: Optional[int]) -> dict: - """The ``ctx_rows_start`` keyword (from ``prepare_1st_drafter_inputs``) for a draft model whose ``dflash_forward`` - takes it; no keyword for another draft model or without a start.""" - if ctx_rows_start is None or not _takes_ctx_rows_start(draft_model): - return {} - return {"ctx_rows_start": ctx_rows_start} - - -def capture_view( - spec_metadata: Optional[SpecMetadata], layer_id: int, num_tokens: int -) -> Optional[torch.Tensor]: - """``layer_id``'s slot of the DFlash capture buffer for ``num_tokens`` rows (a strided view a kernel can write the - tap into), or None when ``spec_metadata`` is not DFlash's or the layer is not captured.""" - if ( - not isinstance(spec_metadata, DFlashSpecMetadata) - or spec_metadata.captured_hidden_states is None - ): - return None - i = spec_metadata._layer_to_idx.get(layer_id) - if i is None: - return None - return spec_metadata.captured_hidden_states[ - :num_tokens, i * spec_metadata.hidden_size : (i + 1) * spec_metadata.hidden_size - ] - - -class KimiK3DFlashWorker(DFlashWorker): - # Set by a Kimi K3 target running its decode kernels: decode steps then run the acceptance and the drafter's - # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context - # K/V with trtllm::k3_ctx_kv. - k3_decode = False - - def _mask_token_id(self, draft_model) -> int: - """The drafter's mask token id, resolved once.""" - if self._resolved_mask_token_id is None: - if ( - hasattr(self.spec_config, "mask_token_id") - and self.spec_config.mask_token_id is not None - ): - self._resolved_mask_token_id = self.spec_config.mask_token_id - elif hasattr(draft_model, "mask_token_id"): - self._resolved_mask_token_id = draft_model.mask_token_id - elif hasattr(draft_model.model, "mask_token_id"): - self._resolved_mask_token_id = draft_model.model.mask_token_id - else: - raise ValueError( - "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " - "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." - ) - return self._resolved_mask_token_id - - @staticmethod - def _trained_mask_embedding(draft_model, mask_token_id: int) -> Optional[torch.Tensor]: - """The drafter's own trained mask row, if it kept one for the mask id in use.""" - if mask_token_id != getattr(draft_model, "mask_token_id", None): - return None - return getattr(draft_model, "mask_token_embedding", None) - - def _k3_accept_applies( - self, logits, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int - ) -> bool: - """Whether this step's acceptance and the drafter's inputs run as ``trtllm::k3_spec_accept``. - - That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no - rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target - logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block - table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. - """ - if not self._k3_accept_step_applies( - attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ): - return False - K = spec_metadata.runtime_draft_len - if ( - logits.dim() != 2 - or logits.dtype != torch.float32 - or logits.shape[0] != num_gens * (K + 1) - ): - return False - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - - embed = draft_model.draft_model_full.model.embed_tokens - return accept_op.supports( - logits.shape[1], num_gens, self._compute_block_size, K, embed.weight.shape[1] - ) - - def _k3_accept_step_applies( - self, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int - ) -> bool: - """Everything ``_k3_accept_applies`` checks but the target logits.""" - if not self.k3_decode: - return False - from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MambaHybridCacheManagerV2, - ) - - if num_contexts != 0 or not 0 < num_gens <= 8 or self.guided_decoder is not None: - return False - if ( - not spec_metadata.is_all_greedy_sample - or self._can_use_rejection_sampling(spec_metadata) - or getattr(spec_metadata, "enable_penalty", False) - or type(self)._reshape_draft_tokens_for_accept - is not SpecWorkerBase._reshape_draft_tokens_for_accept - or type(self)._reshape_logits_for_accept - is not SpecWorkerBase._reshape_logits_for_accept - ): - return False - K = spec_metadata.runtime_draft_len - draft_tokens = spec_metadata.draft_tokens - if ( - K <= 0 - or draft_tokens is None - or draft_tokens.dtype != torch.int32 - or draft_tokens.numel() != num_gens * K - ): - return False - mgr = attn_metadata.kv_cache_manager - mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) - if ( - not isinstance(mgr, MambaHybridCacheManagerV2) - or type(mgr).update_mamba_states is not MambaHybridCacheManagerV2.update_mamba_states - or not mgr.use_kda_replay_update - or getattr(mgr, "prev_num_accepted_tokens", None) is None - or getattr(mgr, "_dummy_request_mask", None) is None - or mamba_metadata is None - or mamba_metadata.state_indices.dtype != torch.int32 - ): - return False - if ( - self._ctx_block_tables is None - or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None - or getattr(attn_metadata, "kv_lens_cuda", None) is None - ): - return False - embed = draft_model.draft_model_full.model.embed_tokens - return getattr(embed, "tp_size", 1) == 1 and embed.weight.dtype == torch.bfloat16 - - def target_logits( - self, hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model - ) -> torch.Tensor: - """This rank's bf16 vocabulary shard of the target logits, [rows, vocab / TP] (``_k3_head_shard``: the head - GEMM without its all-gather and fp32 cast), for a step whose acceptance runs as the sharded - ``trtllm::k3_spec_accept`` (see ``_k3_logits_shard``); the gathered fp32 logits otherwise. At such a step the - acceptance is their only reader: the one-model spec sampler takes the worker's tokens and rejects requests for - logits or log probabilities, so the shard is also what the step returns as ``"logits"``.""" - self._k3_step_shard = None - shard = self._k3_logits_shard( - hidden_states, lm_head, attn_metadata, spec_metadata, draft_model - ) - if shard is None: - return super().target_logits( - hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model - ) - logits = self._k3_head_shard(logits_processor, lm_head, hidden_states) - self._k3_step_shard = (logits, shard) - return logits - - @staticmethod - def _k3_head_shard(logits_processor, lm_head, rows: torch.Tensor) -> torch.Tensor: - """This rank's bf16 vocabulary shard of ``lm_head(rows)``, without the head's all-gather: from the logits - processor's own head kernel where it has one that takes the rows (``lm_head_shard``; the Kimi K3 target's - ``gemm/k3_head_gemv``), so the shard holds the values of the processor's gathered logits; else from the head's - own GEMM.""" - head_shard = getattr(logits_processor, "lm_head_shard", None) - logits = None if head_shard is None else head_shard(rows, lm_head) - if logits is None: - logits = lm_head.apply_linear(rows, lm_head.bias) - return logits - - def _k3_logits_shard(self, hidden_states, lm_head, attn_metadata, spec_metadata, draft_model): - """``(workspace, first column)`` when this step's target logits stay vocabulary-sharded, else None. - - That needs a step ``_k3_accept_step_applies`` takes, plain TP whose all-reduces own an MNNVL workspace for - this mapping, and an unquantized, bias-free, unpadded, evenly split column-parallel bf16 lm_head with shards - the kernel splits. The exchange is collective, so every rank must decide alike: every condition is the - batch's or the configuration's. The workspace is allocated (collectively) and the kernel compiled outside - CUDA-graph capture only, so a graph whose warmup allocated no workspace is captured on the gathered path. - """ - if not self.k3_decode: - return None - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce - - mapping = self.mapping - if getattr(self, "_k3_shard_head", None) is not lm_head: - reason = self._k3_logits_shard_reason(lm_head) - if reason is None and MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None: - # Not kept: re-checked on later steps, since the model's MNNVL workspace may not exist yet. - logger.info_once( - "DFlash: the target logits are all-gathered (the TP all-reduces are not MNNVL)", - key="dflash_k3_logits_no_mnnvl", - ) - return None - if mapping is not None and mapping.tp_size > 1: - if reason is None: - logger.info( - f"DFlash: decode-step target logits stay vocabulary-sharded ({lm_head.weight.shape[0]} " - f"columns per rank); trtllm::k3_spec_accept exchanges the row maxima, and a request's " - f"logits post-processors get them gathered" - ) - else: - logger.info_once( - f"DFlash: the target logits are all-gathered ({reason})", key=reason - ) - self._k3_shard_head, self._k3_shard_ok = lm_head, reason is None - if not self._k3_shard_ok: - return None - num_contexts = attn_metadata.num_contexts - num_gens = attn_metadata.num_seqs - num_contexts - K = spec_metadata.runtime_draft_len - if hidden_states.shape[0] != num_gens * (K + 1) or not self._k3_accept_step_applies( - attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ): - return None - embed = draft_model.draft_model_full.model.embed_tokens - columns = lm_head.weight.shape[0] - if not accept_op.supports( - columns, num_gens, self._compute_block_size, K, embed.weight.shape[1], mapping.tp_size - ): - return None - if torch.cuda.is_current_stream_capturing(): - ws = accept_op.existing_workspace(mapping) - if ws is None: - return None - else: - ws = accept_op.workspace(mapping) - return ws, mapping.tp_rank * columns - - def _k3_logits_shard_reason(self, lm_head): - """Why the target logits cannot stay vocabulary-sharded under this configuration and head (None if they can, - given an MNNVL workspace for the mapping, which ``_k3_logits_shard`` checks).""" - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - from tensorrt_llm._torch.modules.linear import TensorParallelMode - - mapping = self.mapping - if mapping is None or mapping.tp_size <= 1: - return "no TP" - if mapping.enable_attention_dp or mapping.has_cp() or mapping.has_pp(): - return "attention DP, CP or PP" - if ( - getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN - or not getattr(lm_head, "gather_output", False) - or getattr(lm_head, "padding_size", 0) != 0 - or getattr(lm_head, "gather_output_sizes", None) is not None - ): - return "the lm_head is not an evenly split, unpadded column-parallel head" - if ( - lm_head.bias is not None - or lm_head.has_any_quant - or lm_head.weight.dtype != torch.bfloat16 - ): - return "the lm_head has a bias, is quantized or is not bf16" - if not accept_op.supports_columns(lm_head.weight.shape[0], mapping.tp_size): - columns = lm_head.weight.shape[0] - return f"trtllm::k3_spec_accept cannot split {columns} columns over {mapping.tp_size} ranks" - return None - - def _k3_static_arange(self, lo: int, hi: int) -> torch.Tensor: - """``arange(lo, hi)`` (int64), built once outside CUDA-graph capture.""" - cache = self.__dict__.setdefault("_k3_aranges", {}) - t = cache.get((lo, hi)) - if t is None: - t = torch.arange(lo, hi, dtype=torch.long, device="cuda") - if not torch.cuda.is_current_stream_capturing(): - cache[(lo, hi)] = t - return t - - def _k3_mask_row(self, draft_model) -> torch.Tensor: - """The noise block's mask embedding, bf16 [hidden], cached: the drafter's own trained mask row if it kept one - for the mask id in use, else the embedding's row (``dflash_noise_block_embedding`` reads the same row).""" - row = getattr(self, "_k3_mask_embedding", None) - if row is None: - embed = draft_model.draft_model_full.model.embed_tokens - mask_token_id = self._mask_token_id(draft_model) - trained = self._trained_mask_embedding(draft_model, mask_token_id) - if trained is not None: - row = trained.to(embed.weight.dtype).reshape(-1).contiguous() - else: - row = embed.weight[mask_token_id].contiguous() - if torch.cuda.is_current_stream_capturing(): - return row - self._k3_mask_embedding = row - return row - - def _k3_ctx_pool_key(self): - """The identity of the pool the ctx cache is bound to now (KV cache estimation rebinds it to a new one).""" - return ( - tuple(t.data_ptr() for t in self._ctx_kv_buf) if self._ctx_kv_buf is not None else None - ) - - def _k3_ctx_kv_applies(self, draft_model, projected: torch.Tensor) -> bool: - """Whether ``trtllm::k3_ctx_kv`` writes this step's context K/V: the manager-bound paged pool with K and V, - bf16, the fused K/V weight without bias or context input norm, k_norm, NeoX RoPE from flashinfer's fp32 - cache, a context the op can split (up to 64 tokens: B <= 8 requests of K + 1 <= 8, ``ctx_op.pick_split``), one - allocation for every layer's pool. Checked once per token count and - bound pool; the pool view is kept for the bound pool only, so a rebind never leaves the kernel writing the - replaced pool.""" - pool_key = self._k3_ctx_pool_key() - if getattr(self, "_k3_ctx_pool_bound", None) != pool_key: - self._k3_ctx_pool_bound = pool_key - self._k3_ctx_pool = None - self._k3_ctx_kv_ok = {} - cache = self._k3_ctx_kv_ok - n = projected.shape[0] - ok = cache.get(n) - if ok is not None: - return ok - from tensorrt_llm._torch.cute_dsl_kernels.k3_ctx_kv import op as ctx_op - from tensorrt_llm._torch.models import modeling_dflash - - reason = None - if not (self._ctx_paged and self._ctx_block_tables is not None): - reason = "the context pool is not the manager-bound paged pool" - elif projected.dtype != torch.bfloat16 or self._ctx_kv_buf[0].dtype != torch.bfloat16: - reason = "not bf16" - else: - if draft_model._fused_kv_weight is None: - draft_model._build_fused_kv_buffers() - view = ctx_op.pool_view(list(self._ctx_kv_buf)) - if ( - draft_model._fused_kv_bias is not None - or getattr(draft_model, "_input_ln_eps", None) is not None - ): - reason = "K/V bias or context input norm" - elif draft_model._k_norm_stacked is None or not draft_model._is_neox: - reason = "no k_norm or not NeoX RoPE" - elif ( - modeling_dflash._flashinfer_rope is None - or draft_model._get_cos_sin_cache().shape[1] != 64 - ): - reason = "RoPE is not flashinfer's full-rotary fp32 cache" - elif view is None or self._ctx_kv_buf[0].size(1) != 2: - reason = "the layers' pools are not K/V views of one allocation" - elif not ctx_op.pick_split( - draft_model._fused_kv_weight.shape[0], - projected.shape[1], - draft_model._num_kv_heads, - self._draft_tokens_per_req, - n, - projected.device, - ): - reason = f"{n} context tokens (at most 64) or the K/V weight shape" - else: - self._k3_ctx_pool = view - ok = reason is None - if not ok: - logger.info(f"DFlash: context K/V on the Python path ({reason})") - if not torch.cuda.is_current_stream_capturing(): - cache[n] = ok - return ok - - def _k3_ctx_kv( - self, draft_model, projected, ctx_positions, num_accepted, slots, rows - ) -> torch.Tensor: - """``trtllm::k3_ctx_kv``: every drafter layer's context K/V of this step into the paged pool and - ``_ctx_len += num_accepted`` (clamped); returns the context length each request may advertise.""" - if self._k3_ctx_pool is None or self._k3_ctx_pool_bound != self._k3_ctx_pool_key(): - raise RuntimeError("k3_ctx_kv: the context pool was rebound after the path was chosen") - flat, layer_off, page_stride, kv_stride, head_stride = self._k3_ctx_pool - return torch.ops.trtllm.k3_ctx_kv( - projected.contiguous(), - draft_model._fused_kv_weight, - draft_model._k_norm_stacked, - draft_model._get_cos_sin_cache(), - ctx_positions, - num_accepted, - self._ctx_len, - slots, - rows, - self._ctx_block_tables, - self._ctx_block_counts, - flat, - layer_off, - page_stride, - kv_stride, - head_stride, - draft_model._k_norm_eps, - self._max_ctx, - self._ctx_page_size, - self._compute_block_size, - draft_model._num_kv_heads, - ) - - def _on_acceptance( - self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata - ) -> None: - """Called with this step's acceptance when ``k3_spec_accept`` produced it (a hook for drafter families).""" - - def _k3_accept(self, logits, attn_metadata, spec_metadata, draft_model, shard=None): - """``trtllm::k3_spec_accept`` for a step ``_k3_accept_applies`` accepted: the acceptance, the block table, - the KDA replay record, kv_lens + 1 and the drafter's inputs in one launch. ``shard``: the exchange - workspace and first column when ``logits`` are this rank's vocabulary shard (``_k3_logits_shard``). - Returns (accepted tokens, num accepted, {rewind, bonus, qpos, cpos, noise}).""" - from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op - - num_gens = attn_metadata.num_seqs - K = spec_metadata.runtime_draft_len - mgr = attn_metadata.kv_cache_manager - is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() - force = 0.0 if is_warmup else float(self.force_num_accepted_tokens) - mode, _, _ = accept_op.force_mode(force, K) - if mode == accept_op._kernel_module().FORCE_FRAC: - self._ensure_force_accept_rng_state(logits.device) - pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter - else: - dummy_rng = getattr(self, "_k3_dummy_rng", None) - if dummy_rng is None: - dummy_rng = self._k3_dummy_rng = ( - torch.zeros(4, dtype=torch.float32, device=logits.device), - torch.zeros(2, dtype=torch.int64, device=logits.device), - ) - pool, counter = dummy_rng - embed = draft_model.draft_model_full.model.embed_tokens - accepted, num_acc, rewind, bonus, qpos, cpos, noise = torch.ops.trtllm.k3_spec_accept( - logits, - spec_metadata.draft_tokens.reshape(num_gens, K), - attn_metadata.draft_kv_cache_block_offsets, - self._ctx_pool_idx, - self._ctx_block_divisor, - self._ctx_block_counts, - self._ctx_block_tables, - mgr.prev_num_accepted_tokens, - attn_metadata.mamba_metadata.state_indices, - mgr._dummy_request_mask, - attn_metadata.kv_lens_cuda, - self._batch_to_slot, - self._ctx_len, - self._max_ctx, - embed.weight, - self._k3_mask_row(draft_model), - pool, - counter, - force, - self._compute_block_size, - *self._k3_shard_args(shard), - ) - self._on_acceptance(accepted, num_acc, attn_metadata, spec_metadata) - return ( - accepted, - num_acc, - dict(rewind=rewind, bonus=bonus, qpos=qpos, cpos=cpos, noise=noise), - ) - - @staticmethod - def _k3_shard_args(shard) -> tuple: - """``trtllm::k3_spec_accept``'s exchange arguments for a vocabulary-sharded step (none otherwise).""" - if shard is None: - return () - ws, first_column = shard - return ( - ws["uc"], - ws["mc"], - ws["flags"], - ws["rank"], - ws["slots"], - ws["push_copies"], - first_column, - ) - - def _forward_impl( - self, - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - resource_manager=None, - ): - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - # Set when target_logits kept this step's logits vocabulary-sharded (only k3_spec_accept reads them). - step_shard, self._k3_step_shard = getattr(self, "_k3_step_shard", None), None - shard = None - if step_shard is not None: - if step_shard[0] is not logits: - raise RuntimeError( - "DFlash: the target logits are not the vocabulary shard target_logits returned" - ) - shard = step_shard[1] - - raw_logits = logits - K = spec_metadata.runtime_draft_len - - if K == 0: - return self.skip_drafting( - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - ) - - # Lazy init buffers and attach worker reference for prepare() - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - self._lazy_init_ctx_buffers( - draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager - ) - spec_metadata._dflash_worker = self - # A decode step of the supported mode runs the acceptance and the drafter's inputs as one kernel - # (trtllm::k3_spec_accept), which also decodes the block table below. - if shard is None: - fused_accept = self._k3_accept_applies( - logits, attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ) - elif self._k3_accept_step_applies( - attn_metadata, spec_metadata, draft_model, num_contexts, num_gens - ): - fused_accept = True - else: - raise RuntimeError( - "DFlash: the target logits are vocabulary-sharded, but the fused acceptance no longer applies" - ) - # Before any store: prefill and decode both address pages through it. - # Returning False here means an empty batch -- the missing-offsets case - # raises inside, with the metadata type in the message. - if not fused_accept: - self._refresh_ctx_block_tables(attn_metadata, batch_size) - - # Save context lengths so both warmup and a failed forward can roll - # back the in-place _ctx_len updates made during drafting. - is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() - if not torch.cuda.is_current_stream_capturing(): - # Never allocate the snapshot while capturing a CUDA graph: the - # clone would live in the graph memory pool, and its replay-time - # writes could alias blocks reused by later captures. Rollback is - # only meaningful for eager/warmup forwards anyway; a failure - # during capture aborts the graph itself, and captured ops do not - # mutate _ctx_len until replay. - self._saved_ctx_len = self._ctx_len.clone() - self._saved_ctx_len_host = list(self._ctx_len_host) - self._saved_req_ctx_pos = dict(self._req_ctx_pos) - self._ctx_len_restore_pending = True - - self._execute_guided_decoder_if_present(logits) - - prep = None - if fused_accept: - # Saved before the kernel updates kv_lens (warmup also snapshots kv_lens_cuda here). - self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) - accepted_tokens, num_accepted_tokens, prep = self._k3_accept( - logits, attn_metadata, spec_metadata, draft_model, shard - ) - else: - accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( - logits, attn_metadata, spec_metadata - ) - - # Opt-in acceptance recording (env-gated; eager-mode measurement - # runs only). Skipped for CUDA-graph batches (capture/replay/warmup - # use synthetic requests and forbid the host sync). - if ( - self._accept_stats is not None - and num_gens > 0 - and not spec_metadata.is_cuda_graph - and not torch.cuda.is_current_stream_capturing() - and spec_metadata.request_ids is not None - ): - self._accept_stats.on_accept( - spec_metadata.request_ids[num_contexts:batch_size], - num_accepted_tokens[num_contexts:batch_size].tolist(), - ) - - if fused_accept: - # The kernel recorded the KDA replay acceptance and advanced kv_lens_cuda. - self._kv_rewind_amount = prep["rewind"] - self._kv_rewind_nc = num_contexts - self._kv_rewind_bs = batch_size - self._kv_rewind_pending = True - attn_metadata.update_for_spec_dec() - else: - # Update GDN/Mamba recurrent states to the accepted token's state. - if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): - attn_metadata.kv_cache_manager.update_mamba_states( - attn_metadata=attn_metadata, - num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices, - ) - - self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) - self._prepare_kv_for_draft_forward( - attn_metadata, num_accepted_tokens, num_contexts, batch_size - ) - - # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. - # The draft model uses standard 1D RoPE, so only scalar positions are needed. - if position_ids.ndim == 3: - position_ids = position_ids[0, 0] - else: - position_ids = position_ids.squeeze(0) - - # Get total tokens processed by target model (for hidden state extraction) - total_target_tokens = input_ids.shape[0] - - # Capture prefill (context) hidden states for future gen steps. - # This gives the draft model the full prompt context, not just gen tokens. - if num_contexts > 0: - self._store_prefill_context( - draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens - ) - # Rebuild batch_to_slot after prefill assigns new slots - if self._ctx_buf_inited and spec_metadata.request_ids: - self._rebuild_batch_to_slot(spec_metadata.request_ids) - - inputs = self.prepare_1st_drafter_inputs( - input_ids=input_ids, - position_ids=position_ids, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - total_target_tokens=total_target_tokens, - prep=prep, - ) - - if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states_out = draft_model.dflash_forward( - noise_embedding=inputs["noise_embedding"], - query_positions=inputs["query_positions"], - num_ctx_per_req=inputs["num_ctx_per_req"], - ctx_k_cache=inputs["ctx_k_cache"], - ctx_v_cache=inputs["ctx_v_cache"], - ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], - ctx_kv_cache=inputs["ctx_kv_cache"], - ctx_page_table=inputs["ctx_page_table"], - **dflash_ctx_rows_kwargs(draft_model, inputs["ctx_rows_start"]), - ) - - # Gather K logits per gen request from the block outputs. - # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. - # Which block slots carry them is a drafter-family convention, - # resolved through _draft_slot_ids. - block_size = self._compute_block_size - gen_hidden_states = self._draft_block_hidden_states( - draft_model, hidden_states_out, num_gens, block_size, K - ) - gen_logits = self._draft_block_logits( - draft_model, gen_hidden_states, attn_metadata, spec_metadata - ) - - vocab_size = gen_logits.shape[-1] - gen_logits = gen_logits.reshape(num_gens, K, vocab_size) - - gen_logits = self._refine_block_logits( - draft_model, gen_logits, inputs, spec_metadata - ) - - # DFlash 2: replace the independent per-position picks with one - # coherent path through the block (absent for plain DFlash). - if getattr(draft_model, "has_candidate_selector", False): - gen_logits = self._apply_dflash2_selector( - draft_model, - gen_logits, - gen_hidden_states.reshape(num_gens, K, -1), - inputs["first_prev_tokens"], - spec_metadata, - ) - vocab_size = gen_logits.shape[-1] - - gen_draft_tokens = self.sample_draft_tokens( - gen_logits, - spec_metadata, - batch_size, - num_contexts=num_contexts, - ) - - # Opt-in confidence-calibration recording (env-gated). The - # confidence values come from an optional provider on the - # recorder (None here -> no calibration rows collected; the - # DSpark confidence-scheduled verification MR supplies the - # real provider, see accept_stats.py). Guarded off for - # CUDA-graph batches and d2t vocab-mapped drafters. - if ( - self._accept_stats is not None - and self._accept_stats.confidence_provider is not None - and not spec_metadata.is_cuda_graph - and not torch.cuda.is_current_stream_capturing() - and spec_metadata.request_ids is not None - and self._d2t is None - ): - self._accept_stats.record_draft_confidence( - spec_metadata.request_ids[num_contexts:batch_size], - draft_model, - gen_hidden_states.reshape(num_gens, K, -1), - inputs["first_prev_tokens"], - gen_draft_tokens, - ) - - else: - gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") - - # Context requests are not drafted by the block worker (zero placeholder - # token); fill their draft-prob slot rows with a one-hot placeholder so - # they are a legal distribution when they become gen requests next iter. - gen_vocab = vocab_size if num_gens > 0 else None - self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) - - if num_contexts > 0 and num_gens > 0: - ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) - elif num_contexts > 0: - next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") - else: - next_draft_tokens = gen_draft_tokens - - self._restore_attn_metadata_from_spec_dec(attn_metadata) - self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) - - self._rollback_guided_decoder_after_verify(num_accepted_tokens) - - next_new_tokens = self._prepare_next_new_tokens( - accepted_tokens, - next_draft_tokens, - spec_metadata.batch_indices_cuda, - batch_size, - num_accepted_tokens, - ) - - # Restore context lengths after warmup; real runs keep the updates. - if is_warmup: - self._ctx_len.copy_(self._saved_ctx_len) - self._restore_ctx_len_host() - self._ctx_len_restore_pending = False - - outputs = { - "logits": raw_logits, - "new_tokens": accepted_tokens, - "new_tokens_lens": num_accepted_tokens, - "next_draft_tokens": next_draft_tokens, - "next_new_tokens": next_new_tokens, - } - if shard is not None: - # "logits" is this rank's vocabulary shard: the engine gathers it for any logits post-processor. - outputs["logits_vocab_shard"] = True - return outputs - - def _draft_block_hidden_states( - self, - draft_model, - hidden_states_out: torch.Tensor, - num_gens: int, - block_size: int, - num_draft_tokens: int, - ) -> torch.Tensor: - """The block-output rows that produce the K draft logits per gen request (``_draft_slot_ids``). - - The slot ids depend only on the shapes, so they are built once per shape, outside CUDA-graph - capture, and cached. When they are one contiguous run of rows (a single request, or the - shift_label convention with a block of K slots, whose requests' slots follow each other) the - rows are returned as a view, without a gather; otherwise they are gathered. - """ - rows = hidden_states_out.shape[0] - key = (num_gens, block_size, num_draft_tokens, rows) - cache = self.__dict__.setdefault("_draft_block_rows", {}) - entry = cache.get(key) - if entry is None: - ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) - # Shields only the last request: at block_size == K with - # shift_label off, slots run 1..K, so every request reads the - # next one's slot 0 and the last overruns. Degrades, never raises. - ids = ids.clamp(max=rows - 1) - if torch.cuda.is_current_stream_capturing(): - return hidden_states_out[ids] - host = ids.tolist() - if host and host == list(range(host[0], host[0] + len(host))): - entry = host[0] - else: - entry = ids - cache[key] = entry - if isinstance(entry, int): - return hidden_states_out[entry : entry + num_gens * num_draft_tokens] - return hidden_states_out[entry] - - def _draft_block_logits( - self, - draft_model, - gen_hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - spec_metadata, - ) -> torch.Tensor: - """fp32 logits of the block positions over the full vocabulary (the TP lm_head all-gathers its - shards). A drafter family whose refinement and draft sampler reduce across ranks themselves may - return this rank's vocabulary shard instead.""" - return draft_model.logits_processor( - gen_hidden_states, draft_model.lm_head, attn_metadata, True - ) - - def _ctx_rows_start(self, num_contexts: int, num_gens: int) -> Optional[int]: - """The first row of the gen requests' page-table rows where those rows are one contiguous run, else None. - - The manager's block table is keyed by batch position, so the gen requests' rows are [num_contexts, - num_contexts + num_gens): a drafter can view them instead of gathering ``ctx_cache_batch_idx``'s rows. The - private arena's table is keyed by slot, which need not be contiguous. - """ - if num_gens > 0 and self._ctx_block_tables is not None: - return num_contexts - return None - - def prepare_1st_drafter_inputs( - self, - input_ids: torch.LongTensor, - position_ids: torch.LongTensor, - hidden_states: torch.Tensor, - accepted_tokens: torch.Tensor, - num_accepted_tokens: torch.Tensor, - attn_metadata: AttentionMetadata, - spec_metadata: DFlashSpecMetadata, - draft_model: nn.Module, - total_target_tokens: int = 0, - prep: Optional[dict] = None, - ): - """Prepare inputs for DFlash's draft forward. - - ``prep``: ``trtllm::k3_spec_accept``'s bonus tokens, positions and noise embedding for a decode step it - took (see ``_k3_accept``); computed here otherwise. - - For gen requests, builds: - - noise_embedding: token embeddings for [accepted + mask] tokens - - query_positions: position IDs for the query tokens - - num_ctx_per_req: per-request context length in the pool - - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed - views of the persistent per-layer K/V pool. - - ctx_rows_start: where ctx_cache_batch_idx is the contiguous rows - [ctx_rows_start, ctx_rows_start + num_gens) of ctx_page_table, their - start (``_ctx_rows_start``), else None. - """ - num_contexts = attn_metadata.num_contexts - batch_size = attn_metadata.num_seqs - num_gens = batch_size - num_contexts - - mask_token_id = self._mask_token_id(draft_model) - - # Get the embed_tokens layer from the draft model - embed_tokens = draft_model.draft_model_full.model.embed_tokens - hidden_dim = ( - spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] - ) - - if num_gens > 0: - gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] - gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] - - total_tokens_per_req = self._draft_tokens_per_req # K+1 - K = spec_metadata.runtime_draft_len - - # Get captured multi-layer hidden states from spec_metadata - captured_hs = spec_metadata.get_hidden_states(total_target_tokens) - has_target_features = ( - captured_hs is not None - and hasattr(draft_model, "fc") - and hasattr(draft_model, "hidden_norm") - ) - - # Use cached block_size (resolved once on first call) - block_size = self._compute_block_size - query_tokens_per_req = block_size - - # Get slots for gen requests from pre-computed mapping - slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] - K_plus_1 = K + 1 - - if prep is not None: - gen_rows_out = self._k3_static_arange(num_contexts, num_contexts + num_gens) - offsets_kp1 = self._k3_static_arange(0, K_plus_1) - bonus = prep["bonus"] - query_position_ids = prep["qpos"] - ctx_position_ids = prep["cpos"] - noise_embed_2d = prep["noise"] - else: - gen_rows_out = torch.arange( - num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" - ) - - bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) - bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() - - ctx_len_gen = self._ctx_len[slots] - j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") - offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - - # _ctx_len is clamped to _max_ctx only AFTER this step's accepted - # tokens are folded in (see the update below), so the running length - # used here has to be clamped on its own -- otherwise a request that - # already sits at the ceiling indexes num_accepted positions past - # any position the sequence can legitimately reach. - ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) - query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) - ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) - - noise_embed_2d = dflash_noise_block_embedding( - embed_tokens, - bonus, - mask_token_id, - query_tokens_per_req, - self._trained_mask_embedding(draft_model, mask_token_id), - ) - - # Accumulate new accepted features into context buffers - fused_ctx_kv = False - if has_target_features: - gen_start = attn_metadata.num_ctx_tokens - # Target now processes exactly K+1 tokens per gen req, so the - # captured slice is already the full set we need to project. - gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] - gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) - projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) - fused_ctx_kv = prep is not None and self._k3_ctx_kv_applies( - draft_model, projected_to_store - ) - if fused_ctx_kv: - num_ctx_fused = self._k3_ctx_kv( - draft_model, - projected_to_store, - ctx_position_ids, - gen_num_accepted, - slots, - gen_rows_out, - ) - # k3_ctx_kv adds the accepted tokens in the kernel; empty the dummy slot after - # it, as _advance_ctx_len does after the add on the torch path. - self._ctx_len[self._dummy_slot].zero_() - if has_target_features and not fused_ctx_kv: - gen_num_accepted_long = gen_num_accepted.long() - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) - write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - if self._ctx_block_tables is not None: - # Bound by what the manager handed THIS request: the masked - # entries below are still written (fixed-size, for graph - # safety) and a placeholder page resolves to another's. - gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size - col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - else: - col_idx = col_idx.clamp(max=self._max_ctx - 1) - - # Fixed-size writes for CUDA graph compatibility: - # Write ALL entries but zero out invalid ones. Invalid - # writes land at clamped column indices beyond valid - # _ctx_len range, so they're harmless. - slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) - col_flat = col_idx.reshape(-1) - proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] - pos_flat = ctx_position_ids.long().reshape(-1) - mask_1d = write_mask.reshape(-1) - - # Fast path: store the pre-projected/pre-RoPE'd K/V. - # dflash_forward reads these directly via cache_batch_idx. - cache_dtype = ( - self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype - ) - k_new, v_new = draft_model.precompute_context_kv( - proj_flat.to(cache_dtype), pos_flat - ) - mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) - k_new.mul_(mask_bc) - if v_new is not None: - v_new.mul_(mask_bc) - slot_long = slot_flat.long() - col_long = col_flat.long() - if self._ctx_paged: - if self._ctx_block_tables is not None: - # Batch positions of the gen requests, matching the - # per-request block table's row order. - rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) - else: - rows_long = slot_long - self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - else: # VANILLA DFlash backend (FlashAttention) - self._ctx_k_buf[slot_long, :, col_long] = k_new - if v_new is not None: - self._ctx_v_buf[slot_long, :, col_long] = v_new - - self._advance_ctx_len(slots, gen_num_accepted_long) - - if fused_ctx_kv: - num_ctx_per_req_t = num_ctx_fused - else: - num_ctx_per_req_t = self._ctx_len[slots] - if not fused_ctx_kv and self._ctx_block_tables is not None: - # The write above clamps columns to the request's allocation, so - # a context that outruns it has its tail written over the last - # valid slot. Truncate what is advertised to match, or the read - # attends that clobbered value -- or another request's page. - allocated = dflash_allocated_ctx_limit( - self._ctx_block_counts, self._ctx_page_size, block_size - ) - num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) - noise_embedding = noise_embed_2d - query_positions = query_position_ids.long() - - # Update seq_lens for gen requests to K+1 (the number of tokens - # the target forward actually processed). - attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( - total_tokens_per_req - ) - attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req - else: - noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) - query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") - num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") - slots = torch.empty(0, dtype=torch.long, device="cuda") - gen_rows_out = slots - bonus = torch.empty(0, dtype=torch.long, device="cuda") - - return { - "noise_embedding": noise_embedding, - "query_positions": query_positions, - "num_ctx_per_req": num_ctx_per_req_t, - "ctx_k_cache": self._ctx_k_buf, - "ctx_v_cache": self._ctx_v_buf, - # Slots index the private arena; the manager's block table is keyed - # by batch position, so the drafter reads its own rows there. - "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, - # Anchor token per gen request (block slot 0): last accepted - # token. The dspark Markov chain conditions its first step on it. - "first_prev_tokens": bonus, - "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": ( - self._ctx_block_tables - if self._ctx_block_tables is not None - else self._ctx_page_table - ), - "ctx_rows_start": self._ctx_rows_start(num_contexts, num_gens), - } - - -class KimiK3DSparkWorker(KimiK3DFlashWorker, DSparkWorker): - # Whether this step's block logits are the Kimi K3 decode path's vocabulary shard (see ``_draft_block_logits``). - _k3_sharded_block_logits = False - - def _apply_dspark_markov_bias( - self, - draft_model, - gen_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - spec_metadata, - ) -> torch.Tensor: - """Apply the dspark vanilla-Markov intra-block bias to block logits. - - Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): - step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head - logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} - is the greedy token from step i-1's biased logits. Greedy per-position - argmax of the returned logits therefore reproduces the reference - sampled chain; the rejection-sampling path samples from the same - biased distributions (proposal conditioned on the greedy chain). - - Block logits the Kimi K3 decode path kept vocab-sharded (see ``_draft_block_logits``) run the - whole chain as ``trtllm::k3_markov``, which also returns the greedy draft tokens and - next_new_tokens. Other TP-sharded logits slice markov_w2's rows to this rank's contiguous - shard and chain through the TP-aware global argmax. - """ - # The d2t guard lives in set_draft_model: it is model-static, so raising - # it here would surface a load-time config error per decode step. - # Unlike the d2t guard this one cannot move to set_draft_model: it - # keys on the runtime logits width, and reproducing that at init would - # duplicate the draft head's sharding rules. A standalone drafter - # borrows the target lm_head, whose gather_output defaults to True, so - # the logits normally arrive full-vocab and this branch is skipped. - self._k3_markov = None - k3_sharded, self._k3_sharded_block_logits = self._k3_sharded_block_logits, False - full_vocab = draft_model.markov_w2.shape[0] - shard = gen_logits.shape[-1] - vocab_slice = None - if shard != full_vocab: - mapping = self.mapping - if ( - mapping is None - or getattr(mapping, "enable_attention_dp", False) - or shard * mapping.tp_size != full_vocab - ): - raise NotImplementedError( - f"DSpark Markov head: draft logits width {shard} does not " - f"match the drafter vocab {full_vocab} and is not a plain " - "TP column shard of it." - ) - vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) - if k3_sharded: - return self._k3_markov_chain( - draft_model, gen_logits, first_prev_tokens, vocab_slice - ) - - def argmax_fn(step_logits): - # Full-vocab token ids (TP-aware when sharded); tokens stay in - # draft-vocab space, which is what markov_w1 indexes. - return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() - - return draft_model.apply_markov_chain_logits( - gen_logits, - first_prev_tokens, - argmax_fn=argmax_fn, - vocab_slice=vocab_slice, - ) - - def _keep_draft_logits_sharded(self, draft_model, spec_metadata, num_gens: int) -> bool: - """Keep the draft logits vocab-sharded: ``trtllm::k3_markov`` reduces across the ranks itself. - - The chain's global argmaxes are the greedy draft tokens, so the lm_head's all-gather of the - full logits and the draft sampler's gather are both skipped. Needs a Kimi K3 target running its - decode kernels (``k3_decode``), plain TP, greedy drafting (rejection sampling reads - full-vocabulary probabilities), an unquantized bias-free column-parallel bf16 lm_head whose - shards tile the Markov vocabulary, bf16 Markov weights of the kernel's rank, MNNVL (the model's - all-reduces own an MNNVL workspace for this mapping), and a block, shard and batch of - ``num_gens`` requests the kernel can split. - """ - if not self.k3_decode: - return False - from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op - from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce - from tensorrt_llm._torch.modules.linear import TensorParallelMode - - mapping = self.mapping - lm_head = getattr(draft_model, "lm_head", None) - if ( - lm_head is None - or not getattr(draft_model, "has_markov_head", False) - or mapping is None - or mapping.tp_size <= 1 - or mapping.enable_attention_dp - or spec_metadata.wants_advanced_draft_sampling - or getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN - or not getattr(lm_head, "gather_output", False) - or getattr(lm_head, "bias", None) is not None - or lm_head.weight.dtype != torch.bfloat16 - or draft_model.markov_w1.dtype != torch.bfloat16 - or draft_model.markov_w2.dtype != torch.bfloat16 - or lm_head.weight.shape[0] * mapping.tp_size != draft_model.markov_w2.shape[0] - or MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None - ): - return False - shard = lm_head.weight.shape[0] - block = spec_metadata.runtime_draft_len - rank = k3_markov_op._kernel_module().MARKOV_RANK - if ( - not 0 < block <= k3_markov_op.WORKSPACE_MAX_BLOCK - or draft_model.markov_w1.dim() != 2 - or draft_model.markov_w1.shape[1] != rank - or tuple(draft_model.markov_w2.shape[1:]) != (rank,) - ): - return False - if k3_markov_op.pick_grid(shard, block, num_gens) == 0: - logger.warning_once( - f"DSpark Markov head: trtllm::k3_markov cannot split a {shard}-row vocab shard for " - f"{num_gens} requests; the draft logits are all-gathered and the chain runs unfused.", - key=f"dspark_k3_markov_shard_{num_gens}", - ) - return False - return True - - def _draft_block_logits( - self, - draft_model, - gen_hidden_states: torch.Tensor, - attn_metadata, - spec_metadata, - ) -> torch.Tensor: - """This rank's bf16 shard of the block logits when ``k3_markov`` takes them (it converts them to fp32 - exactly, as ``.float()`` would), from the drafter's logits processor's head kernel where it takes the rows, - else the head's own GEMM (``_k3_head_shard``); otherwise the base class' fp32 logits.""" - num_gens = gen_hidden_states.shape[0] // max(spec_metadata.runtime_draft_len, 1) - self._k3_sharded_block_logits = self._keep_draft_logits_sharded( - draft_model, spec_metadata, num_gens - ) - if self._k3_sharded_block_logits: - return self._k3_head_shard( - draft_model.logits_processor, draft_model.lm_head, gen_hidden_states - ) - return super()._draft_block_logits( - draft_model, gen_hidden_states, attn_metadata, spec_metadata - ) - - def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): - """The base acceptance; under ``k3_decode`` its outputs are kept for ``k3_markov``'s next_new_tokens.""" - accepted_tokens, num_accepted_tokens = super().sample_and_accept_draft_tokens( - logits, attn_metadata, spec_metadata - ) - if self.k3_decode: - self._on_acceptance(accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata) - return accepted_tokens, num_accepted_tokens - - def _on_acceptance( - self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata - ) -> None: - """Keeps this step's acceptance for ``k3_markov``'s next_new_tokens (and the metadata whose KV lengths it - rewinds).""" - self._k3_acceptance = ( - accepted_tokens, num_accepted_tokens, attn_metadata.num_contexts, spec_metadata, attn_metadata - ) # fmt: skip - - def _k3_markov_chain( - self, - draft_model, - gen_logits: torch.Tensor, - first_prev_tokens: torch.Tensor, - vocab_slice: slice, - ) -> torch.Tensor: - """The Markov chain as ``trtllm::k3_markov`` on this rank's shard of the block logits. - - Returns the corrected logits (fp32); keeps the kernel's greedy tokens for - ``sample_draft_tokens`` and its next_new_tokens (built from this step's acceptance) for - ``_prepare_next_new_tokens``. - """ - from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op - - acceptance = getattr(self, "_k3_acceptance", None) - if acceptance is None: - raise RuntimeError("DSpark Markov head: k3_markov needs this step's acceptance first") - accepted_tokens, num_accepted_tokens, num_contexts, spec_metadata, attn_metadata = ( - acceptance - ) - num_gens = gen_logits.shape[0] - # The draft forward's pending KV-length rewind (see _apply_kv_rewind_after_draft) runs in the kernel: every - # reader of kv_lens_cuda in the draft forward precedes it. In warmup the restore of kv_lens_cuda that - # follows overwrites it, as it would the Python rewind. - rewind = getattr(self, "_kv_rewind_amount", None) - fold = ( - getattr(self, "_kv_rewind_pending", False) - and rewind is not None - and getattr(attn_metadata, "kv_lens_cuda", None) is not None - and rewind.dtype == torch.int32 - and self._kv_rewind_bs - self._kv_rewind_nc == rewind.numel() <= num_gens - ) - corrected, tokens, next_new = k3_markov_op.markov_chain( - self.mapping, - gen_logits.contiguous(), - first_prev_tokens.long(), - draft_model.markov_w1, - draft_model.markov_w2[vocab_slice], - vocab_slice.start, - accepted_tokens, - num_accepted_tokens[num_contexts : num_contexts + num_gens], - spec_metadata.batch_indices_cuda[num_contexts : num_contexts + num_gens], - kv_lens=attn_metadata.kv_lens_cuda if fold else None, - rewind=rewind if fold else None, - rewind_first=self._kv_rewind_nc if fold else 0, - ) - if fold: - self._kv_rewind_amount = None - self._kv_rewind_pending = False - self._k3_markov = (corrected, tokens, next_new, accepted_tokens, num_accepted_tokens) - return corrected - - def sample_draft_tokens( - self, - logits, - spec_metadata, - batch_size, - *, - num_contexts=0, - draft_step=None, - mapping_lm_head_tp=None, - ): - """Greedy block drafts straight from ``k3_markov``. - - When ``logits`` are the corrected logits the kernel just returned, its per-position global - argmax (first maximum, lowest vocabulary index among equal values) is the token the - TP-gathered greedy sampler would pick. Anything else goes to the base sampler. - """ - chain = getattr(self, "_k3_markov", None) - if ( - chain is not None - and chain[0] is logits - and mapping_lm_head_tp is None - and not spec_metadata.wants_advanced_draft_sampling - ): - self._k3_markov_next = (chain[1], chain[2], chain[3], chain[4]) - return chain[1] - self._k3_markov_next = None - return super().sample_draft_tokens( - logits, - spec_metadata, - batch_size, - num_contexts=num_contexts, - draft_step=draft_step, - mapping_lm_head_tp=mapping_lm_head_tp, - ) - - def _prepare_next_new_tokens( - self, - accepted_tokens, - next_draft_tokens, - batch_indices_cuda, - batch_size, - num_accepted_tokens, - ): - """``k3_markov``'s next_new_tokens when the drafts are its tokens for the whole batch (no context - requests); otherwise the base assembly.""" - chain = getattr(self, "_k3_markov_next", None) - # The step ends here: release its acceptance (which holds the step's attention and spec metadata) and the - # kernel's outputs, so they do not keep a captured graph's pool or metadata alive after the step. - self._k3_markov_next = None - self._k3_acceptance = None - self._k3_markov = None - if ( - chain is not None - and chain[0] is next_draft_tokens - and chain[2] is accepted_tokens - and chain[3] is num_accepted_tokens - and next_draft_tokens.shape[0] == batch_size - ): - return chain[1] - return super()._prepare_next_new_tokens( - accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py index 950a29442166..24233d194709 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py @@ -1,42 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Weight loading: Kimi K3 (NVFP4) / sm_107 / tp16_moetp4ep4. +"""The shared Kimi K3 text checkpoint loader.""" -The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and -`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: - -* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the - routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under - moe_tp 4 x moe_ep 4), and checks its own key coverage. -* The vision tower and the projector are a predicted non-load: listed here, never read. -* Any other key fails the load, naming it, rather than being dropped. -""" - -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM -from tensorrt_llm._torch.models.modeling_utils import filter_weights - -_LANG_PREFIX = "language_model." - -# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. -PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") - - -def load(model, weights) -> None: - """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" - unknown = sorted( - k - for k in weights.keys() - if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) - ) - assert not unknown, ( - f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " - f"first {unknown[:5]}" - ) - lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) - assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" - checkpoint_dir = getattr(weights, "checkpoint_dir", None) - if checkpoint_dir is not None: - lm_weights.checkpoint_dir = checkpoint_dir - lm_weights.checkpoint_prefix = _LANG_PREFIX - KimiLinearForCausalLM.load_weights(model, lm_weights) +from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/modeling.py new file mode 100644 index 000000000000..bd742b6c5859 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/modeling.py @@ -0,0 +1,3694 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared Kimi K3 text model for MXFP4 and NVFP4 deployment targets. + +MXFP4 uses the K3 fused attention and routed-expert decode kernels. NVFP4 +preserves the checkpoint's per-layer quantization and uses the existing NVFP4 +MoE and quantized attention kernels. Both formats share the decoder, attention +residuals, BF16 GEMVs, checkpoint streaming and speculative decoding. +Concrete targets specify the SM and expert parallelism independently. +""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( + k3_drafter_attn_qknorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( + k3_mla_attn_vb_out, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( + K3MlaAttnWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import k3_mla_qkv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_attn import k3_kda_attn +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import K3KdaBuffers +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( + k3_kda_decode_attn, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify +from tensorrt_llm._torch.attention.backends import AttentionMetadata + +# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. +from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.decode_view import k3_mla_decode_view +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.models.modeling_dspark import ( + GQADSparkForCausalLM, + draft_is_embedded_in_target, +) +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.moe.fused_moe import ( + ConfigurableMoE, + SiTuActivation, + TRTLLMGenFusedMoE, + create_moe, +) +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.config_utils import is_mla +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from . import decode_comm as _decode_comm +from . import decode_gemv as _decode_gemv +from . import decode_moe as _decode_moe +from . import spec_worker as _spec_worker +from . import weights as _weights + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, +# and the certification is per arch. + +#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in +#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router +#: and fused-A GEMMs), then the decode kernels'. +REQUIRED_TRTLLM_OPS = ( + "attn_res_fwd", + "attn_res_rmsnorm_fwd", + "attn_res_add_rmsnorm_fwd", + "attn_res_add_rmsnorm_persistent_fwd", + "kda_prefill", + "kda_decode", + "kda_mtp_decode", + "dsv3_router_gemm_op", + "dsv3_fused_a_gemm_op", + "k3_kda_decode_attn", + "k3_kda_attn", + "k3_kda_verify", + "k3_mla_qkv", + "k3_mla_attn_vb_out", + # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). + "k3_decode_gemv", + "k3_ctm_gemv_wide", + "k3_ctm_gemv_long", + "k3_situ_mul", + "k3_head_gemv", + "k3_embed_norm", + "allgather", + # The decode path's collectives (decode_comm.py) and MoE (decode_moe.py). + "mnnvl_allreduce_attn_res", + "k3_sandwich_oproj", + "k3_sandwich_tail", + "k3_moe_front", + "k3_moe", + "k3_latent_reduce", + "k3_route_quant", + "mnnvl_allgather_split", + # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces + # with the residual add and RMSNorm (the split context projection's with hidden_norm). + "k3_ctm_gemv", + "k3_ctm_gemv_swiglu", + "k3_drafter_attn_qknorm", + "k3_sandwich_plain", + "mnnvl_fusion_allreduce", +) + +#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. +#: A renamed field upstream fails here, loudly, instead of reading as a default. +REQUIRED_ENGINE_FIELDS = { + "attn_metadata": ( + "num_contexts", + "num_generations", + "num_seqs", + "num_tokens", + "seq_lens", + "tokens_per_block", + "kv_cache_manager", + "mamba_metadata", + ), + "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), +} + +#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every +#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry +#: leaves the list when a catalog entry replaces it. +UNCERTIFIED_GENERIC_CALLS = ( + # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. + "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", + "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", + "tensorrt_llm._torch.models.modeling_utils.DecoderModel", + # The text model's stock modules. + "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", + "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", + "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", + "tensorrt_llm._torch.moe.fused_moe.create_moe", + "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", + "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", + "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", + "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", + "tensorrt_llm._torch.modules.situ.SituAndMul", + "tensorrt_llm._torch.modules.rms_norm.RMSNorm", + "tensorrt_llm._torch.distributed.AllReduce", + "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", + # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries + # do not take a block, its context projection without the TP group's collective state, its context k / v, its + # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. + "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", + "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", + "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", + "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", +) + +# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware +# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). +DECODE_MAX_TOKENS = 8 +MAX_REQUESTS = 8 +MAX_TOKENS_PER_REQUEST = 8 + +# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own +# get_model_defaults sets it for the same reason). +_TOKENS_PER_BLOCK = 64 + +_LANG_PREFIX = "language_model." + + +# ---------------------------------------------------------------------------------------------------------------------- +# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and +# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, +# `K3DecodeMLA` below). +# ---------------------------------------------------------------------------------------------------------------------- + +# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections +# instead of the min-latency fused GEMM op (read once at import). +_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( + os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" +) + + +# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table +# holds cos=1/sin=0), but the chunked-context path indexes the table by +# absolute position, so it must cover max_position_embeddings (~512MB per +# backend for the 1M-position checkpoint); a smaller table is read out of +# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context +# deployments. +_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" + + +class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): + """DeepSeek-V3 routing, computed outside the TRTLLM-Gen MoE kernel. + + The kernel's own top-16 of 896 experts is the slow part of a generic MoE call of a few tokens. Outside it, the + scheduler routes with the backend's fused route + MXFP8 quantize (``trtllm::k3_route_quant``) up to 64 tokens and + with ``noaux_tc_op`` above, and the kernel takes the expert ids and weights. Both compute the kernel's arithmetic + and break ties the same way; a bf16 weight can differ by one ulp where it lies within fp32 rounding of a bf16 + rounding boundary. From about a thousand tokens on, ``noaux_tc_op`` takes longer than the kernel's routing. + """ + + @property + def requires_separated_routing(self) -> bool: + return True + + +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + @property + def routing_method(self) -> KimiK3MoeRoutingMethod: + """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return KimiK3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + +def _get_text_config(pretrained_config: "PretrainedConfig"): + """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" + if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( + not hasattr(pretrained_config, "linear_attn_config") + and hasattr(pretrained_config, "text_config") + ): + return pretrained_config.text_config + return pretrained_config + + +def _is_kda_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] + + +def _is_mla_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" + + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" + + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" + + +KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" + + +_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") + + +# Turning the port on should not require knowing which of the three topologies +# is the right one: "1" selects the measured policy (``split``), so enabling the +# feature and choosing the policy are one action through one variable. +_ATTN_RES_TOPOLOGY_ON = "1" + + +def _read_attn_res_topology() -> str: + """Default ``per_token``: the persistent kernel is opt-in. + + ``1`` is the only accepted on-value and resolves to ``split``. The named + topologies stay available for measurement: ``persistent`` uses the + persistent kernel at every shape it implements, ``per_token`` at none. + """ + raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") + if raw == _ATTN_RES_TOPOLOGY_ON: + return "split" + if raw not in _ATTN_RES_TOPOLOGIES: + # Loudly, for the same reason as the token ceiling below: a mistyped A/B + # arm that silently fell back to the default would measure one side + # twice and report no difference. + raise ValueError( + f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " + f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " + f"(which means 'split'), got {raw!r}" + ) + return raw + + +def _read_fused_attn_res_max_tokens() -> int: + """Resolved after the topology, because its default follows it. + + With the persistent port off -- the default -- the ceiling is 1, the + pre-existing gate: the fused epilogue is taken at the single-token decode + shape and nowhere else. Enabling the port raises it to 32, the top of the + measured range, so that "off" keeps meaning "unchanged". + """ + default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" + raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) + try: + value = int(raw) + except ValueError: + # Failing loudly matters more than usual here: a mistyped A/B arm that + # silently fell back to the default would measure the candidate twice + # and report no difference. + raise ValueError( + f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value < 1: + raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") + return value + + +_ATTN_RES_TOPOLOGY = _read_attn_res_topology() + + +_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() + + +def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: + """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. + + No token ceiling: the persistent grid is sized by the SM count, not by the + token count, so prefill is the case it exists for. + """ + del M # deliberately unused; see above + return H == 7168 and 2 <= N <= 9 + + +def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: + """Pick between the two fused kernels for this call site. + + ``persistent`` takes the persistent kernel at every shape it implements; + ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in + for the prefill/decode boundary. Shapes the persistent kernel does not + implement fall through to the caller's existing gate and land on the + unfused path. + """ + if not _persistent_attn_res_applicable(M, H, N): + return False + if _ATTN_RES_TOPOLOGY == "persistent": + return True + return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> Optional[torch.Tensor]: + """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. + + Returns ``None`` when the call falls outside the fused kernel's + contract (dtype/shape/arch) so the caller can use the exact fp32 reference + instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` + layout. Candidate order matches the reference: snapshots first, the + running prefix sum last. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + try: + attn_res_op = torch.ops.trtllm.attn_res_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output, _rsigma, _probs, _logits = attn_res_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + ) + return output.reshape(M, H) + + +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that + the shape gate let the call through, and a rejected call looks exactly like + a disabled one in the logs. Emitted at debug level: it is once per distinct + shape, not once per process, so it is a diagnostic rather than a summary. + """ + logger.debug_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, + default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, + # which is the measured window; larger token counts have not been measured + # and fall back to the unfused add + attn_res_fwd + RMSNorm path. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + _, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + None, + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) + return output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in + ``_apply_attn_res_rmsnorm_fused``). + + The production residual add produces a BF16 tensor that remains live across + the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual output, + while avoiding a separate add launch and a re-read of the intermediate by + attention-residual selection. + """ + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # Same measured window as _apply_attn_res_rmsnorm_fused above. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + updated_prefix_sum, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + addend.reshape(M, 1, H).contiguous(), + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = attn_res_add_rmsnorm_op( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + +def _apply_attn_res( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> torch.Tensor: + """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). + + prefix_sum: ``[num_tokens, hidden_size]`` + block_residual: ``[num_snapshots, num_tokens, hidden_size]`` + + Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's + contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. + Only the fallback boundary restores the HF ``[M, K, H]`` layout. + """ + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + block_residual_hf = block_residual.transpose(0, 1) + v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.eps) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the + fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_rmsnorm_fused( + prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in + ``_apply_attn_res_and_rmsnorm``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + + +# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where +# they are faster than the add -> attn_res -> RMSNorm chain. +_WIDE_ATTN_RES_MAX_TOKENS = 32 + + +def _attn_res_max_tokens(step: DecodeStep) -> int: + """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step + ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" + return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS + + +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + + +# The routed experts' module path prefixes. ``exclude_modules`` matches with +# wildcards and walks ancestor prefixes, so an empty prefix would widen what +# matches instead of just missing. +_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") + + +@dataclass(frozen=True) +class K3TargetSpec: + """The fixed hardware, expert format and parallelism of a concrete target.""" + + sm: tuple[int, int] + quant_algo: QuantAlgo + moe_tp: int = 4 + moe_ep: int = 4 + + +def _uses_nvfp4(config: ModelConfig) -> bool: + return config.quant_config is not None and config.quant_config.quant_algo in ( + QuantAlgo.MIXED_PRECISION, + QuantAlgo.NVFP4, + ) + + +# Routed-expert quantization used when the checkpoint declares nothing per +# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors +# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that +# checkpoint is what this default has always served. +_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 + + +def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_packed_mxfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), + w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), + w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), + w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), + w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), + w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), + ) + + +def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_streaming_nvfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + **{ + f"{w}_{kind}": get_tensor(f"{base}.{expert_idx}.{w}.{kind}") + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.NVFP4: _K3ExpertCkptSpec( + kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"), + loader=_load_nvfp4_expert, + loaded_slots_attr="_streamed_expert_slots", + needs_layer_finalize=True, + ), + QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( + kinds=("weight_packed", "weight_scale"), + loader=_load_packed_mxfp4_expert, + loaded_slots_attr="_packed_mxfp4_loaded_slots", + needs_layer_finalize=False, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + self.num_experts = cfg.num_experts + self.top_k = cfg.num_experts_per_token + self.moe_hidden_size = cfg.routed_expert_hidden_size + # ValueError (not assert): these guard unsupported checkpoint + # configurations and must stay active under ``python -O``. + if self.moe_hidden_size is None: + raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") + if not getattr(cfg, "latent_moe_use_norm", False): + raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") + + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) + self._situ_betas = (situ_beta, situ_linear_beta) + dtype = torch.bfloat16 + + # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 + # accumulate/output (checkpoint stores the gate weight in bf16; saves a + # per-layer input cast + fp32 splitK pair on the bs1 decode path). + # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True + self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) + + routed_moe_model_config = self._routed_moe_model_config(model_config) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) + routed_moe_kwargs = dict( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_intermediate_size, + dtype=dtype, + # Kimi owns the latent reduction so it can order that collective + # after the shared expert's auxiliary-stream reduction. + reduce_results=False, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + # Let CommunicationFactory select the best available strategy. + communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. As measured once: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver + # instead, which reports the rejection trail. + # + # CUTLASS is absent on purpose: it is the fallback target, so + # "degraded to CUTLASS" is not a thing that can happen to it. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + ) + self._check_trtllm_situ_quant( + routed_moe_model_config.moe_backend, routed_quant_config.quant_algo + ) + + self.routed_experts = create_moe(**routed_moe_kwargs) + if not isinstance(self.routed_experts, ConfigurableMoE): + raise RuntimeError( + "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." + ) + if self.routed_experts.layer_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "dynamic EPLB or replicated expert slots." + ) + local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) + if local_expert_ids != list( + range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) + ): + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming currently requires a " + "contiguous static expert partition." + ) + self.local_expert_ids = tuple(local_expert_ids) + self.experts_per_rank = len(local_expert_ids) + self.expert_lo = local_expert_ids[0] + self.expert_hi = self.expert_lo + self.experts_per_rank + + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = model_config.mapping.tp_size > 1 + self._reduce_routed_output = ( + use_shared_tp + and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + ) + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self.shared_experts = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=shared_model_config, + reduce_output=use_shared_tp, + layer_idx=layer_idx, + is_shared_expert=True, + ) + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.moe_main_event = torch.cuda.Event() + self.moe_shared_event = torch.cuda.Event() + self.routed_expert_down_proj = nn.Linear( + cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype + ) + # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer + # fallback is the same fp32-variance eager math as KimiK3RMSNorm). + self.routed_expert_norm = RMSNorm( + hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. + self.decode_moe: Optional[_decode_moe.K3DecodeMoeLayer] = None + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, projection.weight.t(), None, None + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """The expert split, validated against the concrete deployment target.""" + return mapping.moe_tp_size, mapping.moe_ep_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Resolve per-layer NVFP4 metadata or the compressed-tensors MXFP4 default. + + An exclusion outranks the default and is checked against actual module names. + """ + if _uses_nvfp4(model_config): + global_quant = model_config.quant_config + if any( + global_quant.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in ("language_model.model.", "model.", "") + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + raise ValueError( + f"NVFP4 experts at layer {layer_idx} are excluded from quantization" + ) + declarations = model_config.quant_config_dict or {} + matches = [ + declarations[name] + for prefix in ("language_model.model.", "model.", "") + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + if (name := f"{prefix}layers.{layer_idx}.{suffix}") in declarations + ] + if matches: + if any( + quant.quant_algo != matches[0].quant_algo + or quant.group_size != matches[0].group_size + for quant in matches[1:] + ): + raise ValueError( + f"Conflicting expert quantization aliases at layer {layer_idx}" + ) + return matches[0] + if global_quant.quant_algo == QuantAlgo.NVFP4: + return global_quant + raise ValueError(f"Missing NVFP4 expert declaration at layer {layer_idx}") + quant_config = model_config.quant_config + if quant_config is not None and any( + quant_config.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + logger.debug( + "Kimi K3 layer %d routed experts: excluded from quantization, " + "keeping them unquantized", + layer_idx, + ) + return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) + + logger.debug( + "Kimi K3 layer %d routed experts: no per-layer quant config in the " + "checkpoint, defaulting to %s", + layer_idx, + _K3_DEFAULT_ROUTED_QUANT_ALGO, + ) + return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) + + @staticmethod + def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: + """Reject formats without a fused SiTU cubin in the resolved TRTLLM backend. + + Read the supported formats from the backend so model construction and + kernel selection agree, including when AUTO resolves to TRTLLM. + """ + situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() + if moe_backend != "TRTLLM" or quant_algo in situ_supported: + return + supported = ", ".join(sorted(algo.name for algo in situ_supported)) + raise ValueError( + f"Kimi K3 routed experts are quantized as {quant_algo}, which the " + "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " + f"exist only for {supported}. Set moe_config.backend to CUTLASS " + "or MEGAMOE_CUTEDSL." + ) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config, with the split of ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. + supported_backends = { + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + f"got {model_config.moe_backend!r}." + ) + if model_config.moe_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "EPLB or replicated expert slots." + ) + mapping = model_config.mapping + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + logger.info_once( + f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " + f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", + key="kimi_k3_moe_tp_ep_split", + ) + + mapping_dict = mapping.to_dict() + mapping_dict["moe_cluster_size"] = 1 + mapping_dict["moe_tp_size"] = moe_tp + mapping_dict["moe_ep_size"] = moe_ep + routed_mapping = Mapping.from_dict(mapping_dict) + + routed_model_config = copy.copy(model_config) + routed_model_config._frozen = False + routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) + routed_model_config.mapping = routed_mapping + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend in { + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + }: + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) + if configured_moe_max_num_tokens < default_moe_max_num_tokens: + logger.info_once( + "Kimi K3 MegaMoE raises moe_max_num_tokens from " + f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " + "because the global DP SymmBuffer requires capacity for " + "max_num_tokens * dp_size.", + key=( + "kimi_k3_megamoe_capacity_override_" + f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" + ), + ) + routed_model_config.moe_max_num_tokens = max( + configured_moe_max_num_tokens, + default_moe_max_num_tokens, + ) + routed_model_config._frozen = True + return routed_model_config + + def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: + """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel + share (``partial_tail``): the decode path takes the step with that tail.""" + return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + partial_tail: bool = False, + *, + step: DecodeStep, + ) -> Union[torch.Tensor, _decode_comm.PendingTail]: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification + (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where + ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at + most 8 tokens, a tensor on a wide decode step.""" + decode = self.decode_moe + if decode is not None and decode.takes(hidden_states, step, partial_tail): + return decode.forward( + self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push + ) + if partial_tail: + raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") + return self.forward_generic(hidden_states, all_rank_num_tokens) + + def forward_generic( + self, hidden_states: torch.Tensor, all_rank_num_tokens=None + ) -> torch.Tensor: + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None + + def _routed_output(): + # Latent down/up projections via the min-latency fused GEMM op: + # at <=16 tokens (decode graphs) it runs a single pipelined + # bf16 kernel per projection instead of cuBLAS's split-K GEMV + + # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); + # for larger token counts the op falls back to cuBLAS internally. + # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear + # (A/B escape hatch). When the FP8 weight-read conversion has + # replaced the projection module, call it directly: its weight is + # an e4m3 buffer the bf16 dsv3 op must not read, and its forward + # is already a single fused GEMM (fp8_swap_ab_gemm). + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + if self._reduce_routed_output: + return y + # Communication-backed paths return a complete routed result. + y = self.routed_expert_norm(y) + return self._routed_projection(y, self.routed_expert_up_proj) + + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. The shared GatedMLP includes its output all-reduce on + # the auxiliary stream. The join below must precede the routed + # all-reduce: concurrent collectives on different streams can corrupt + # SYMM_MEM all-reduce state. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + disable_on_compile=True, + ) + if self._reduce_routed_output: + routed_latent = moe_all_reduce(routed_out) + routed_latent = self.routed_expert_norm(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + return routed_out + shared_out + + +def resolve_attention_quant_config( + config: ModelConfig | None, layer_idx: int, projection: str +) -> QuantConfig: + """Resolve a checkpoint projection, including mixed-precision exclusions.""" + if config is None: + return QuantConfig() + global_config = config.quant_config or QuantConfig() + names = [ + f"{prefix}layers.{layer_idx}.self_attn.{projection}" + for prefix in ("language_model.model.", "model.", "") + ] + if any(global_config.is_module_excluded_from_quantization(name) for name in names): + return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) + declarations = config.quant_config_dict or {} + matches = [declarations[name] for name in names if name in declarations] + if matches: + selected = matches[0] + if any( + (match.quant_algo, match.group_size) != (selected.quant_algo, selected.group_size) + for match in matches[1:] + ): + raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") + elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: + selected = QuantConfig() + else: + selected = global_config + if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " + f"quantization {selected.quant_algo}" + ) + if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): + raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") + result = copy.copy(selected) + result.kv_cache_quant_algo = global_config.kv_cache_quant_algo + return result + + +class KimiMLARuntime(nn.Module): + """Wraps K3 MLA and applies its external TP output reduction.""" + + def __init__( + self, + cfg: "PretrainedConfig", + layer_idx: int, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the + # final output reduction in this wrapper so the output gate remains + # between attention and the row-parallel o_proj. + mapping = model_config.mapping + reduce_output = mapping.tp_size > 1 + self._o_allreduce = ( + AllReduce( + mapping=mapping, + strategy=model_config.allreduce_strategy, + dtype=torch.bfloat16, + ) + if reduce_output + else None + ) + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in ( + "q_a_proj", + "kv_a_proj_with_mqa", + "q_b_proj", + "kv_b_proj", + "g_proj", + "o_proj", + ) + } + attention_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + attention_config._frozen = model_config._frozen + mla_cls = K3StockMLA if _uses_nvfp4(model_config) else K3DecodeMLA + self.mixer = mla_cls( + hidden_size=cfg.hidden_size, + num_heads=cfg.num_attention_heads, + q_lora_rank=cfg.q_lora_rank, + kv_lora_rank=cfg.kv_lora_rank, + qk_nope_head_dim=cfg.qk_nope_head_dim, + qk_rope_head_dim=cfg.qk_rope_head_dim, + v_head_dim=cfg.v_head_dim, + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + use_output_gate=cfg.mla_use_output_gate, + max_position_embeddings=max_positions, + model_config=attention_config, + aux_stream_dict=aux_stream_dict, + ) + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" + return self.mixer.will_run_decode_branch(attn_metadata, step) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the + mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer( + None, hidden_states, attn_metadata, step=step, project_output=project_output + ) + return self._reduce(out) if project_output and reduce_output else out + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + out = self.mixer.forward_generic(None, hidden_states, attn_metadata) + return self._reduce(out) if reduce_output else out + + def _reduce(self, out: torch.Tensor) -> torch.Tensor: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + return self._o_allreduce(out) if self._o_allreduce is not None else out + + +class KimiLinearDecoderLayer(nn.Module): + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + dtype = torch.bfloat16 + + self.is_kda = _is_kda_layer(cfg, layer_idx) + is_mla = _is_mla_layer(cfg, layer_idx) + if self.is_kda == is_mla: + raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") + + if self.is_kda: + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in projection_names + } + attention_config._frozen = model_config._frozen + kda_cls = K3StockKDA if _uses_nvfp4(model_config) else K3DecodeKDA + self.linear_attn = kda_cls( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + ) + + self.is_moe = ( + cfg.num_experts is not None + and layer_idx >= cfg.first_k_dense_replace + and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 + ) + if self.is_moe: + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's + # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take + # (decode_gemv.SITES); otherwise it stays within one node. + spans_nodes = self._mnnvl_allreduce() is not None + if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. + self.mlp = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, + ) + self._situ = (situ_beta, situ_linear_beta) + # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics + # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied + # after downcast, use_gemma=False). + self.input_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3). The res norms stay + # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by + # _apply_attn_res and the fused attn_res op, never called as + # modules. + self.attn_res_block_size = cfg.attn_res_block_size + assert self.attn_res_block_size is not None, ( + "Kimi K3 runtime expects attn_res_block_size to be set" + ) + self.self_attention_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set + # by the target's post_load_weights. + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + self.sandwich_oproj = False + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + step: DecodeStep, + capture: Optional[Tuple[Any, int]] = None, + prenormed: bool = False, + pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, + defer_moe_tail: bool = False, + ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + ``block_residual`` is a preallocated snapshot bank in kernel-native + ``[K_max, M, H]`` layout. Returns the running prefix sum and the + number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. + + ``step`` is the step's classification (``decode_step``), handed to the + attention module. + + ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is + this layer's input norm, and the layer's input, the step's embedding, + already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). + + The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens + runs the attention's all-reduce and the residual update in one + collective once the target has built its ``decode_comm``: with o_proj + (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, + the step and the attention's decode branch, else on the attention's + unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). + + ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the + previous layer's MoE output, unreduced (``defer_moe_tail``); then + ``hidden_states`` is the prefix sum without it, and this layer's + pre-attention step reduces and adds it: a ``PendingTail`` in one + ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its + all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer + the ``sandwich_tail`` call also stores the tap into the capture slot + where the speculative metadata exposes it (``capture_view``); other + tapped layers take the tap after the step. + + ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` + instead, ``partial`` this layer's MoE output unreduced, for the next + consumer to reduce. + """ + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + attn_res_max_tokens = _attn_res_max_tokens(step) + tail = ( + pending_moe_partial + if isinstance(pending_moe_partial, _decode_comm.PendingTail) + else None + ) + # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum + # straight into the bank row it snapshots. + snapshot_row = None + if tail is not None and self.layer_idx % self.attn_res_block_size == 0: + snapshot_row = block_residual[num_snapshots] + # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the + # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). + tap_view = None + if tail is not None and capture is not None: + tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) + + if prenormed: + assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 + assert pending_moe_partial is None and capture is None + elif tail is not None: + hidden_states, prefix_sum = self.decode_comm.sandwich_tail( + tail, + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + updated_out=snapshot_row, + tap=tap_view, + tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, + ) + elif pending_moe_partial is not None: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + elif capture is not None: + # The mixture tap needs the PRE-norm value, which the fused + # attn-res + RMSNorm kernel does not expose. Keep the two steps + # split on captured layers only and fuse everywhere else. + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if capture is not None and pending_moe_partial is not None and tap_view is None: + # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output + # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. + tapped = ( + _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + if _AUX_ATTN_RES_STREAM_ENABLED + else prefix_sum + ) + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + + if self.layer_idx % self.attn_res_block_size == 0: + if not prenormed and snapshot_row is None: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, step): + if ( + self.sandwich_oproj + and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS + and attention.will_run_decode_branch(attn_metadata, step) + ): + core = attention(hidden_states, attn_metadata, step=step, project_output=False) + hidden_states, prefix_sum = comm.sandwich_oproj( + core, + self._o_proj().weight, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + if comm is not None and step.wide: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) + else: + attention_output = attention(hidden_states, attn_metadata, step=step) + if prefix_sum is None: + prefix_sum = attention_output + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention_output, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + if self.is_moe and defer_moe_tail: + partial = self.block_sparse_moe( + hidden_states, all_rank_num_tokens, partial_tail=True, step=step + ) + return prefix_sum, num_snapshots, partial + if self.is_moe: + hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) + else: + hidden_states = self._dense_mlp(hidden_states, step) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def forward_generic( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + ) -> Tuple[torch.Tensor, int]: + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + if capture is not None: + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, None): + partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + elif prefix_sum is None: + prefix_sum = attention.forward_generic(hidden_states, attn_metadata) + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention.forward_generic(hidden_states, attn_metadata), + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + if self.is_moe: + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + hidden_states = self.block_sparse_moe.forward_generic( + hidden_states, all_rank_num_tokens + ) + else: + hidden_states = self.mlp(hidden_states) + return prefix_sum + hidden_states, num_snapshots + + def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: + """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode + kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" + gemvs = self.decode_gemvs + if gemvs is not None and step.small: + out = gemvs.dense_mlp( + hidden_states, + self.mlp.gate_up_proj.weight, + self.mlp.down_proj.weight, + *self._situ, + ) + if out is not None: + return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out + return self.mlp(hidden_states) + + def _mnnvl_allreduce(self): + """The MNNVL all-reduce of this layer's attention output, or None.""" + attention = self.linear_attn if self.is_kda else self.self_attn + return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) + + def _o_proj(self) -> nn.Module: + """This layer's attention output projection (row parallel).""" + return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj + + def _o_allreduce(self) -> AllReduce: + """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" + return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce + + def accepts_moe_partial(self, num_snapshots: int) -> bool: + """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target + built its collective state, and the snapshot bank is not empty.""" + return num_snapshots > 0 and self.decode_comm is not None + + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + + +class KimiLinearModel(DecoderModel): + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = _get_text_config(model_config.pretrained_config) + self._text_cfg = cfg + dtype = torch.bfloat16 + + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream() for _ in range(4)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + } + + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + + # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) + # by _apply_attn_res and the fused attn_res op. + self.output_attn_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.num_attn_res_snapshots = ( + cfg.num_hidden_layers + cfg.attn_res_block_size - 1 + ) // cfg.attn_res_block_size + # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the + # weights are final. None: every step embeds on the generic path. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + + @property + def kda_draft_records(self) -> bool: + """Whether the hybrid cache manager keeps the last verify's draft records (``kda_state_tok``), from which + ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn`` replay the accepted drafts. The engine reads it once the weights + are loaded, to build the manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer + taking the K3 kernels. Otherwise the KDA verify replays the accepted drafts from the replay caches (the + built-in verify).""" + spec_config = getattr(self.model_config, "spec_config", None) + return bool( + spec_config is not None + and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) + and spec_config.tokens_per_gen_step in (2, 4, 6, 8) + and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) + ) + + def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, + from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from + load-time state, so every rank and every replay of the graph decides alike.""" + if not torch.cuda.is_current_stream_capturing(): + return False + kda = [layer.linear_attn for layer in self.layers if layer.is_kda] + mla = [layer.self_attn for layer in self.layers if not layer.is_kda] + return latent_push( + step, + capturing=True, + breakable=is_in_breakable_cuda_graph(), + kda_draft_records=self.kda_draft_records, + kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), + mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), + ) + + def _defer_moe_tail( + self, + i: int, + hidden_states: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + step: DecodeStep, + ) -> bool: + """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE + takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after + the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps + its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated + tail.""" + layer = self.layers[i] + if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): + return False + if spec_metadata is not None and ( + capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) + ): + return False + return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] + step = decode_step(attn_metadata, num_tokens) + if step is None: + return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) + if self._latent_push(attn_metadata, step): + step = replace(step, latent_push=True) + # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. + prenormed = None + if ( + inputs_embeds is None + and self.decode_gemvs is not None + and len(self.layers) > 0 + and self.num_attn_res_snapshots > 0 + ): + table = self.embed_tokens.weight + block_residual = table.new_empty( + self.num_attn_res_snapshots, num_tokens, table.shape[1] + ) + prenormed = self.decode_gemvs.embed_norm( + input_ids, table, self.layers[0].input_layernorm, block_residual + ) + if prenormed is not None: + hidden_states = prenormed + else: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) + # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. + pending_moe_partial = None + for i, layer in enumerate(self.layers): + defer_moe_tail = self._defer_moe_tail( + i, hidden_states, num_snapshots, spec_metadata, capture_set, step + ) + outputs = layer( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + step=step, + capture=self._capture(spec_metadata, capture_set, i), + prenormed=i == 0 and prenormed is not None, + pending_moe_partial=pending_moe_partial, + defer_moe_tail=defer_moe_tail, + ) + if defer_moe_tail: + hidden_states, num_snapshots, pending_moe_partial = outputs + else: + (hidden_states, num_snapshots), pending_moe_partial = outputs, None + + if isinstance(pending_moe_partial, _decode_comm.PendingTail): + normed, _ = self.layers[-1].decode_comm.sandwich_tail( + pending_moe_partial, + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + return normed + if pending_moe_partial is not None: + _, normed = _apply_attn_res_add_and_rmsnorm( + hidden_states, + _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + _attn_res_max_tokens(step), + ) + return normed + return self._final_norm( + hidden_states, + block_residual, + num_snapshots, + spec_metadata, + capture_set, + _attn_res_max_tokens(step), + ) + + def _generic_forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor], + inputs_embeds: Optional[torch.FloatTensor], + spec_metadata, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + for i, layer in enumerate(self.layers): + hidden_states, num_snapshots = layer.forward_generic( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + capture=self._capture(spec_metadata, capture_set, i), + ) + return self._final_norm( + hidden_states, block_residual, num_snapshots, spec_metadata, capture_set + ) + + def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + return (spec_metadata, self.layers[i - 1].layer_idx) + return None + + def _final_norm( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + max_fused_tokens: Optional[int] = None, + ) -> torch.Tensor: + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) + + return _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + max_fused_tokens, + ) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. +# ---------------------------------------------------------------------------------------------------------------------- + + +class K3StockKDA(KimiKDALinearAttention): + """Quantized KDA with the target's output-reduction boundary.""" + + takes_k3_kernels = False + k3_buffers = None + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + return False + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + if not project_output: + raise ValueError("Quantized KDA requires its own output projection") + return self.forward_generic(hidden_states, attn_metadata, reduce_output) + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + if reduce_output: + return super().forward(hidden_states, attn_metadata) + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + else: + core = self._forward_impl(hidden_states, attn_metadata) + return self.o_proj(core.reshape(-1, self.proj_size)) + + +class K3StockMLA(KimiK3MLAAttention): + """Stock quantized MLA, including its KV-cache format handling.""" + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + return False + + def forward( + self, + position_ids: torch.Tensor | None, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + project_output: bool = True, + ) -> torch.Tensor: + if not project_output: + raise ValueError("Quantized MLA requires its own output projection") + return self.forward_generic(position_ids, hidden_states, attn_metadata) + + def forward_generic( + self, + position_ids: torch.Tensor | None, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + return super().forward(position_ids, hidden_states, attn_metadata) + + +class K3DecodeKDA(KimiKDALinearAttention): + """Kimi K3's KDA attention: the built-in module, with the decode kernels on the steps they take. + + * A decode step of one token per request runs the fused input projection and the plain decode in one + ``ssm/k3_kda_decode_attn`` launch. + * With the cache manager's per-token states (``KimiLinearModel.kda_draft_records``), every verify of the layer, on + any step, runs the kernels that keep them: ``ssm/k3_kda_attn`` for one request of 8 tokens (the projection fused + in), else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it takes + them). The built-in verify replays drafts from caches these kernels do not fill, so the two never run on one + manager. + * On every step ``decode_step`` classifies, ``o_proj`` runs on the ``o_proj`` decode GEMV site where it takes the + rows, then the module's all-reduce. + + Every other step runs the built-in module. The kernels read one ``[q | k | v | g | f_a | b]`` weight, built at the + checkpoint load from the module's own, and the device's ``K3KdaBuffers`` and decode GEMVs' state, which the + target sets in ``post_load_weights``. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # The fused [q | k | v | g | f_a | b | pad] projection weight; the six projections' weights and the built-in + # fused [q | k | v | g] and [f_a | b | pad] ones are views of it. + self.k3_proj_weight: Optional[torch.Tensor] = None + # The fused projection's Lamport buffers: one set per device, shared by every KDA layer. + self.k3_buffers: Optional[K3KdaBuffers] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @property + def takes_k3_kernels(self) -> bool: + """Whether the decode kernels can run this layer: its fused projection weight is built.""" + return self.k3_proj_weight is not None + + def finalize_decode_weights(self) -> None: + """The built-in fused weights, then one ``[q | k | v | g | f_a | b | pad]`` weight of both.""" + super().finalize_decode_weights() + assert ( + self.use_full_rank_gate + and self.gate_lower_bound is not None + and self._qkvg_proj_weight is not None + and self._bfa_proj_weight is not None + and self._qkvg_proj_weight.dtype == self._bfa_proj_weight.dtype == torch.bfloat16 + ), ( + f"Kimi K3 KDA layer {self.layer_idx}: the decode kernels read the bf16 fused projections the built-in " + "module builds on CUDA at head dim 128, with a full-rank output gate and a gate lower bound" + ) + rows = self._qkvg_proj_weight.shape[0] + with torch.no_grad(): + fused = self._merge_projection_weights( + (self.q_proj, self.k_proj, self.v_proj, self.g_proj, self.f_a_proj, self.b_proj), + pad_rows_to=8, + ) + self.k3_proj_weight = fused + self._qkvg_proj_weight, self._bfa_proj_weight = fused[:rows], fused[rows:] + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode branch: outside a breakable CUDA graph.""" + return not is_in_breakable_cuda_graph() + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic`` where ``will_run_decode_branch`` does not hold. On the decode branch: the plain decode + on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the built-in dispatch; then + ``o_proj`` on its decode GEMV site and the TP all-reduce. + + ``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce), and, on the decode branch only, + ``project_output=False`` the post-o_norm core ``[N, H * 128]`` (no ``o_proj``).""" + if not self.will_run_decode_branch(attn_metadata, step): + return self.forward_generic(hidden_states, attn_metadata, reduce_output) + if ( + step.decode + and step.tokens_per_request == 1 + and self.takes_k3_kernels + and self.k3_buffers is not None + ): + core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) + else: + core = self._forward_impl(hidden_states, attn_metadata) + if not project_output: + return core.reshape(-1, self.proj_size) + return self._k3_project_output(core, reduce_output) + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + if reduce_output: + return super().forward(hidden_states, attn_metadata) + return self.o_proj( + self._builtin_core(hidden_states, attn_metadata).reshape(-1, self.proj_size) + ) + + def _builtin_core( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """The built-in forward's core ``[N, H, 128]``: inside a breakable CUDA graph from the eager + ``maybe_bcg_kda_core_inplace``, else from ``_forward_impl``.""" + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + return core + return self._forward_impl(hidden_states, attn_metadata) + + def _k3_project_output(self, core: torch.Tensor, reduce_output: bool = True) -> torch.Tensor: + """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then, with + ``reduce_output``, the TP all-reduce.""" + core2d = core.reshape(-1, self.proj_size) + out = None + if self.decode_gemvs is not None: + out = self.decode_gemvs.project("o_proj", core2d, self.o_proj.weight) + if out is None: + if reduce_output: + return self._project_output(core) + out = self.o_proj(core2d) + if reduce_output and self._o_allreduce is not None: + out = self._o_allreduce(out) + return out + + def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: + """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; + each slot's conv window and state advance in place.""" + mamba_metadata = attn_metadata.mamba_metadata + slots = getattr(mamba_metadata, "generation_state_indices", None) + if slots is None: + slots = mamba_metadata.state_indices[: x.shape[0]] + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + w_q, w_k, w_v = self._get_mtp_conv_weights() + core = k3_kda_decode_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.conv, + layer_cache.temporal, + slots, + self.k3_buffers, + float(self.gate_lower_bound), + self.head_k_dim**-0.5, + float(self.o_norm.eps), + ) + # Speculative decoding's replay caches keep their committed conv window in step with the pool's. + self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) + return core + + def forward_verify( + self, + x2d, + num_steps, + layer_cache, + conv_pool, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The built-in verify, unless the cache manager keeps the per-token states; then the decode kernels'.""" + if not (layer_cache.has_kda_replay_caches and layer_cache.kda_state_tok is not None): + return super().forward_verify( + x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices, output=output + ) + assert self.takes_k3_kernels and self.k3_buffers is not None, ( + f"Kimi K3 KDA layer {self.layer_idx}: the cache manager keeps the per-token verify states, which only " + "the decode kernels write, and this layer has no fused projection weight or buffers" + ) + core = self._k3_verify(x2d, num_steps, layer_cache, ssm_pool, slot_indices) + return self._store_core(core, output) + + def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch.Tensor: + """The core output ``[N T, H, 128]`` of N requests of T verify tokens: ``ssm/k3_kda_attn`` for one request of + 8 tokens, else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it + takes them). Each slot's state after the golden token, its drafts' states and its conv window are written in + place.""" + num_requests = x.shape[0] // num_steps + slots = slot_indices[:num_requests] + w_q, w_k, w_v = self._get_mtp_conv_weights() + constants = (float(self.gate_lower_bound), self.head_k_dim**-0.5, float(self.o_norm.eps)) + if num_requests == 1 and num_steps == 8: + out = k3_kda_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + self.k3_buffers, + num_steps - 1, + *constants, + ) + else: + rows = None + if self.decode_gemvs is not None: + rows = self.decode_gemvs.project("kda_proj", x, self.k3_proj_weight) + if rows is None: + rows = torch.nn.functional.linear(x, self.k3_proj_weight) + out = k3_kda_verify( + rows, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + num_steps - 1, + *constants, + ) + return out.view(-1, self.num_heads, self.head_dim) + + +class K3DecodeMLA(KimiK3MLAAttention): + """Kimi K3's MLA attention: the built-in module, with a decode step's attention on the decode kernels. + + A decode step runs ``x [W_a; W_g]^T`` with the gate columns through a sigmoid (the ``mla_ag`` decode GEMV site + where it takes the rows, else one GEMM), then ``attention/k3_mla_qkv`` (the q_a / kv_a RMSNorms, q_b and the k_b + absorption into the fused query, the step's latent rows stored into the paged cache) and + ``attention/k3_mla_attn_vb_out`` (the attention over the paged cache, v_b and the output gate in one launch), then + ``o_proj`` (the ``o_proj`` site where it takes the rows, else the module). Every other step, and a decode step whose + cache the kernels do not read (``k3_mla_decode_view`` says why), runs the built-in module. + + ``[W_a; W_g]`` is built at load from the module's weights, which become views of it. The attention workspace is + the device's ``K3MlaAttnWorkspace``; it and the decode GEMVs' state are set by the target in + ``post_load_weights``. + """ + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + # [W_a; W_g]: the fused q_a / kv_a projection's rows, then the output gate's. + self.k3_ag_weight: Optional[torch.Tensor] = None + # The decode attention's workspace: one per device, shared by every MLA layer. + self.k3_workspace: Optional[K3MlaAttnWorkspace] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + def post_load_weights(self) -> None: + """The built-in post-load, then ``[W_a; W_g]`` (once: CUDA graphs captured since read it).""" + super().post_load_weights() + if self.k3_ag_weight is not None: + return + gaps = self._k3_layout_gaps() + assert not gaps, ( + f"Kimi K3 MLA layer {self.layer_idx}: the decode kernels do not take {'; '.join(gaps)}" + ) + qkv_a, gate = self.kv_a_proj_with_mqa, self.g_proj + rows = qkv_a.weight.shape[0] + with torch.no_grad(): + fused = torch.cat([qkv_a.weight, gate.weight]) + qkv_a.weight = nn.Parameter(fused[:rows], requires_grad=False) + gate.weight = nn.Parameter(fused[rows:], requires_grad=False) + self.k3_ag_weight = fused + + def _k3_layout_gaps(self) -> list: + """What of this layer the decode kernels do not take (empty when they take all of it).""" + if not (self.use_output_gate and self.fuse_qkv_a_proj and not self.is_lite): + return ["a layer without the output gate or the fused q_a / kv_a projection"] + linears = (self.kv_a_proj_with_mqa, self.g_proj, self.q_b_proj, self.o_proj) + checks = ( + (not self.mapping.has_cp_helix(), "helix context parallelism"), + (not self.apply_rotary_emb and not self.llama_4_scaling, "RoPE or llama-4 scaling"), + (self.sparse_attn_hooks is None, "sparse attention"), + ( + self.kv_cache_dtype != "fp8_ds_mla" + and not getattr(self.mqa, "has_fp8_kv_cache", False) + and not getattr(self.mqa, "has_fp4_kv_cache", False), + "a quantized KV cache", + ), + ( + all(m.weight.dtype == torch.bfloat16 and m.bias is None for m in linears) + and self.k_b_proj_trans.dtype == self.v_b_proj.dtype == torch.bfloat16, + "projections other than bf16 and unbiased", + ), + ( + not getattr(self.q_a_layernorm, "is_nvfp4", False) + and not getattr(self.kv_a_layernorm, "use_gemma", False), + "an NVFP4 q_a norm or a Gemma kv_a norm", + ), + ( + self.num_heads_tp % 6 == 0 + and self.kv_lora_rank == 512 + and self.qk_rope_head_dim == 64 + and self.q_lora_rank == 1536, + f"{self.num_heads_tp} heads, latent {self.kv_lora_rank}, rope {self.qk_rope_head_dim}, " + f"q_lora {self.q_lora_rank}", + ), + ) + return [why for ok, why in checks if not ok] + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode kernels (see ``_k3_step_view``). Only there does it return + the gated attention output before ``o_proj``; the KV writes forbid running the attention twice, so a caller + that needs it asks first.""" + return self._k3_step_view(attn_metadata, step) is not None + + def _k3_step_view(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> Optional[dict]: + """The paged-cache view the decode kernels read on ``step``: a decode step, this layer's ``[W_a; W_g]`` and + workspace built, outside a breakable CUDA graph, and a cache ``k3_mla_decode_view`` takes. Else None.""" + if not ( + step.decode + and self.k3_ag_weight is not None + and self.k3_workspace is not None + and not is_in_breakable_cuda_graph() + ): + return None + return self._k3_decode_view(attn_metadata, step.num_tokens) + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic``, except on a decode step whose cache the decode kernels read + (``will_run_decode_branch``). There only, ``project_output=False`` returns the gated attention output + ``[M, H * 128]``, the input of ``o_proj``.""" + view = self._k3_step_view(attn_metadata, step) + if view is None: + if not project_output: + raise ValueError("project_output=False needs a step will_run_decode_branch takes") + return self.forward_generic(position_ids, hidden_states, attn_metadata) + x = hidden_states[: step.num_tokens].contiguous() + rows = self.kv_a_proj_with_mqa.weight.shape[0] + gemvs = self.decode_gemvs + ag = None if gemvs is None else gemvs.project("mla_ag", x, self.k3_ag_weight) + if ag is None: + ag = torch.nn.functional.linear(x, self.k3_ag_weight) + ag[:, rows:].sigmoid_() + fused_q = k3_mla_qkv( + ag, + self.q_a_layernorm.weight, + float(self.q_a_layernorm.variance_epsilon), + self.q_b_proj.weight, + self.k_b_proj_trans, + self.kv_a_layernorm.weight, + float(self.kv_a_layernorm.variance_epsilon), + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + ) + attn_output = self.create_output(x, 0) + k3_mla_attn_vb_out( + fused_q, + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + view["softmax_scale"], + self.v_b_proj, + attn_output, + self.k3_workspace, + gate=ag, + gate_col0=rows, + ) + if not project_output: + return attn_output + out = None if gemvs is None else gemvs.project("o_proj", attn_output, self.o_proj.weight) + if out is None: + out = self._project_output([attn_output], position_ids, attn_metadata, None) + return out + + def forward_generic( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + return super().forward(position_ids, hidden_states, attn_metadata) + + def _k3_decode_view(self, attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[dict]: + """The paged latent cache as the decode kernels read it this step, or None when they do not read it (the + reason is logged once).""" + view = k3_mla_decode_view(self.mqa, attn_metadata, num_tokens) + if isinstance(view, str): + logger.info_once( + f"Kimi K3 MLA: the built-in path for a decode step the decode kernels do not read ({view})", + key=f"k3_mla_decode_view_{view}", + ) + return None + return view + + +# ---------------------------------------------------------------------------------------------------------------------- +# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. +# ---------------------------------------------------------------------------------------------------------------------- + +#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at +#: max_draft_len 7. +DRAFTER_ATTN_SPLITS = frozenset( + {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} +) +# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k +# RMSNorm's epsilon and the NeoX RoPE's base. +_DRAFTER_HEADS = (6, 1) +_DRAFTER_HEAD_DIM = 64 +_DRAFTER_PAGE = 64 +_DRAFTER_EPS = 1e-5 +_DRAFTER_ROPE_BASE = 10000.0 + + +def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: + """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input + feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" + if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: + return None + width = in_features // tp_size + return tp_rank * width, (tp_rank + 1) * width + + +class K3FcSlice(nn.Module): + """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is + ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's + ranks is ``fc``'s output.""" + + def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: + super().__init__() + self.weight = nn.Parameter(weight, requires_grad=False) + self.start = start + self.end = end + + @classmethod + def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: + """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns + (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" + columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) + if columns is None: + return None + start, end = columns + return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the + GEMM reads the column block in place, through the rows' stride.""" + return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) + + +class K3DSparkDrafter(GQADSparkForCausalLM): + """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. + + A block the entries take runs, per layer: + + * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; + * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to + its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the + worker writes the accepted rows' context k / v before a block reads them; + * the output projection on ``drafter_o``, then the module's all-reduce; + * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the + module's all-reduce. + + The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state + (below); a projection whose site does not take its rows runs its module. Every other block runs the stock + ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or + normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under + CUDA-graph capture. + + Where the target hands the drafter the TP group's collective state (``use_decode_comm``): + + * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and + ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; + * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the + post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the + block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the + projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection + after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is + followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers + whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. + + Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the + stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); + the context k / v and the Markov head stay upstream code. + """ + + def __init__( + self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" + ) -> None: + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. + self._k3_take: Optional[bool] = None + # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes + # only these. + self._k3_attn_ran: set = set() + # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + # The zero residual of the split context projection's all-reduce, up to a decode step's rows. + self._k3_zero_rows: Optional[torch.Tensor] = None + # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms + # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. + self._k3_norms_fuse = False + self._k3_sandwich_forms: frozenset = frozenset() + + def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: + """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the + target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are + loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. + + * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they + split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` + sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. + * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and + RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose + shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not + compile is not used.""" + self.decode_comm = comm + self._k3_split_fc() + self._k3_norms_fuse = self._k3_norms_take_comm() + forms = set() + if self._k3_norms_fuse: + layers = self.model.layers + for form, swiglu, weights in ( + ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), + ("down", True, [layer.mlp.down_proj.weight for layer in layers]), + ): + if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( + weights[0], swiglu=swiglu + ): + forms.add(form) + self._k3_sandwich_forms = frozenset(forms) + logger.info( + "Kimi K3 DSpark drafter: residual adds and RMSNorms " + + ( + f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " + "mnnvl_fusion_allreduce)" + if self._k3_norms_fuse + else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" + ) + ) + + def _k3_norms_take_comm(self) -> bool: + """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and + post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no + quantized output), and every output and down projection is row parallel with its all-reduce.""" + hidden = self.config.hidden_size + + def plain(norm: nn.Module) -> bool: + weight = getattr(norm, "weight", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (hidden,) + and hasattr(norm, "variance_epsilon") + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "return_hp_output", True) + ) + + def reduces(linear: nn.Module) -> bool: + return ( + getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" + and getattr(linear, "reduce_output", False) + and getattr(linear, "all_reduce", None) is not None + ) + + return plain(self.model.norm) and all( + plain(layer.input_layernorm) + and plain(layer.post_attention_layernorm) + and reduces(layer.self_attn.o_proj) + and reduces(layer.mlp.down_proj) + for layer in self.model.layers + ) + + def _k3_tp_all_reduce(self) -> Optional[nn.Module]: + """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" + o_proj = self.model.layers[0].self_attn.o_proj + if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": + return None + return getattr(o_proj, "all_reduce", None) + + def _k3_split_fc(self) -> None: + """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero + residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is + a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP + all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" + fc = getattr(self, "fc", None) + if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): + return + mapping = self.model_config.mapping + weight = getattr(fc, "weight", None) + sliced = None + if ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and getattr(fc, "bias", None) is None + and self._k3_tp_all_reduce() is not None + ): + sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) + if sliced is None: + logger.info( + "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " + f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" + ) + return + self.fc = sliced + self._k3_zero_rows = weight.new_zeros( + MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] + ) + logger.info( + f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " + f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" + ) + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" + result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + self._k3_split_fc() + return result + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. + + With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group + with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode + step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then + ``hidden_norm``. Otherwise the stock projection.""" + fc = self.fc + if not isinstance(fc, K3FcSlice): + return super().project_target_hidden(hidden_states) + partial = fc(hidden_states.to(fc.weight.dtype)) + rows = partial.shape[0] + zeros = self._k3_zero_rows + if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): + normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) + return normed + if rows == 0: + return self.hidden_norm(partial) + return self.hidden_norm(self._k3_tp_all_reduce()(partial)) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class + docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows + ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows + as a view instead of gathering them.""" + keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) + if keys is None: + return super().dflash_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ) + out = self._k3_block_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ctx_rows_start, + ) + if not torch.cuda.is_current_stream_capturing(): + self._k3_attn_ran |= keys + return out + + def _k3_layers_take(self) -> bool: + """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head + layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, + non-causal unwindowed block attention, and no block convolution or post-attention gate.""" + if self._k3_take is None: + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + take = ( + self._use_fused_qk_norm_rope + and not self.has_block_conv + and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate + ) + for layer_idx, layer in enumerate(self.model.layers): + attn, mlp = layer.self_attn, layer.mlp + pos = getattr(attn, "pos_embd_params", None) + rope = getattr(pos, "rope", None) + norms = (attn.q_norm, attn.k_norm) + take = take and ( + (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS + and attn.head_dim == _DRAFTER_HEAD_DIM + and getattr(attn, "is_qk_norm", False) + and not getattr(attn, "use_gemma_rms_norm", True) + and rope is not None + and pos.is_neox + and getattr(pos, "mrope_section", None) is None + and rope.theta == _DRAFTER_ROPE_BASE + and all( + norm.variance_epsilon == _DRAFTER_EPS + and norm.weight.dtype == torch.bfloat16 + and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) + for norm in norms + ) + and attn.qkv_proj.bias is None + and attn.o_proj.bias is None + and mlp.activation is torch.nn.functional.silu + and mlp.swiglu_limit is None + and mlp.swiglu_alpha in (None, 1.0) + and mlp.swiglu_beta in (None, 0.0) + and mlp.gate_up_proj.bias is None + and mlp.down_proj.bias is None + and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) + ) + self._k3_take = bool(take) + logger.info( + f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " + + ( + "on k3_drafter_attn_qknorm and the drafter GEMV sites" + if self._k3_take + else "on the stock block forward (a layer the drafter entries do not take)" + ) + ) + return self._k3_take + + def _k3_block_keys( + self, + noise_embedding: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor], + ctx_page_table: Optional[torch.Tensor], + ) -> Optional[frozenset]: + """The attention's compile keys for this block when the drafter entries take it, else None.""" + if ( + self.dflash_attention_backend != "TRTLLM" + or ctx_kv_cache is None + or ctx_page_table is None + or noise_embedding.dtype != torch.bfloat16 + or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS + or is_in_breakable_cuda_graph() + or not self._k3_layers_take() + ): + return None + num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads + page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) + dense = ( + num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + ) + strides = set() + for layer_idx in range(len(self.model.layers)): + cache = ctx_kv_cache[layer_idx] + if ( + cache.dtype != torch.bfloat16 + or cache.dim() != 5 + or tuple(cache.shape[1:]) != page + or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) + or cache.stride(0) % 8 + ): + return None + strides.add(cache.stride(0)) + keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) + if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: + return None + return keys + + def _k3_block_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: torch.Tensor, + ctx_page_table: torch.Tensor, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + batch, block = noise_embedding.shape[:2] + rows = batch * block + ctx_len = num_ctx_per_req[:batch].to(torch.int32) + if ctx_rows_start is None: + page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) + else: + # The batch's page-table rows are one contiguous run: a view, no gather. + page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] + positions = query_positions.reshape(-1).contiguous() + hidden = noise_embedding.reshape(rows, -1) + layers = self.model.layers + comm = self._k3_fused_comm(hidden) + residual = None + if comm is not None: + # The fused all-reduces read the residual and return the updated one: the block's input serves as the + # first residual, uncopied. + residual = hidden + normed = layers[0].input_layernorm(hidden) + for layer_idx, layer in enumerate(layers): + attn = layer.self_attn + if comm is None: + if residual is None: + residual = hidden.clone() + normed = layer.input_layernorm(hidden) + else: + normed, residual = layer.input_layernorm(hidden, residual) + qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) + out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) + k3_drafter_attn_qknorm( + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + positions, + attn.q_norm.variance_epsilon, + attn.pos_embd_params.rope.theta, + ctx_kv_cache[layer_idx], + page_table, + ctx_len, + attn.num_heads, + attn.num_key_value_heads, + out, + ) + if comm is not None: + # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after + # the last layer the final norm. + next_norm = ( + layers[layer_idx + 1].input_layernorm + if layer_idx + 1 < len(layers) + else self.model.norm + ) + normed, residual = self._k3_project_norm( + comm, out, attn.o_proj, residual, layer.post_attention_layernorm + ) + normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) + continue + hidden = self._k3_project("drafter_o", out, attn.o_proj) + hidden, residual = layer.post_attention_layernorm(hidden, residual) + hidden = self._k3_mlp(layer.mlp, hidden) + if comm is not None: + return normed + out, _ = self.model.norm(hidden, residual) + return out + + def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: + """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its + all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace + holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" + comm = self.decode_comm + if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): + return None + return comm + + def _k3_project_norm( + self, + comm: _decode_comm.K3DecodeComm, + x: torch.Tensor, + linear: nn.Module, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, + the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the + projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` + where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" + if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( + x, linear.weight, residual, norm + ): + return comm.sandwich_plain(x, linear.weight, residual, norm) + gemvs = self.decode_gemvs + partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) + if partial is None: + partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_mlp_norm( + self, + comm: _decode_comm.K3DecodeComm, + mlp: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down + projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul + form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down + projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; + then ``comm/mnnvl_fusion_allreduce``.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + partial = None + if gate_up is not None: + if "down" in self._k3_sandwich_forms and comm.takes_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ): + return comm.sandwich_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ) + partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + if partial is None: + partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: + """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel + projection's all-reduce.""" + y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) + if y is None: + return linear(x) + return _row_parallel_reduce(linear, y) + + def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on + ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + down = ( + None + if gate_up is None + else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + ) + if down is None: + return mlp(x) + return _row_parallel_reduce(mlp.down_proj, down) + + +def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: + """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" + all_reduce = getattr(linear, "all_reduce", None) + if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: + return all_reduce(partial) + return partial + + +# ---------------------------------------------------------------------------------------------------------------------- +# The target: step classification, the construction checks and the class route returns. +# ---------------------------------------------------------------------------------------------------------------------- + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. + + The checkpoint names its language-model modules `language_model.`; the text model's are ``, with + `layers.*` under `model.`. + """ + config = model_config.pretrained_config + text = copy.copy(model_config) + text._frozen = False + text.pretrained_config = config.text_config + excluded = text.quant_config.exclude_modules + if excluded: + text.quant_config = copy.copy(text.quant_config) + renamed = [] + for name in excluded: + if name.startswith(_LANG_PREFIX): + name = name[len(_LANG_PREFIX) :] + if name.startswith("layers."): + name = "model." + name + renamed.append(name) + text.quant_config.exclude_modules = renamed + text.skip_create_weights_in_init = True + text._frozen = True + return text + + +@dataclass(frozen=True) +class DecodeStep: + """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` + generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" + + num_tokens: int + num_requests: Optional[int] = None + tokens_per_request: Optional[int] = None + #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` + #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. + latent_push: bool = False + + @property + def small(self) -> bool: + """Whether the step fits one token tile of the token-count kernels.""" + return self.num_tokens <= DECODE_MAX_TOKENS + + @property + def decode(self) -> bool: + """Whether the step is a pure decode step the request-aware kernels take.""" + return self.num_requests is not None + + @property + def wide(self) -> bool: + """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode + layout's MoE head and tail, on M-general ops.""" + return self.decode and not self.small + + +def latent_push( + step: DecodeStep, + *, + capturing: bool, + breakable: bool, + kda_draft_records: bool, + kda_decode_kernels: bool, + mla_decode_branch: bool, +) -> bool: + """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at + most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a + breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels + (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA + layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). + Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel + between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a + graph launch orders every pushing replay behind whatever ran before it.""" + if not (step.decode and step.small) or not capturing or breakable: + return False + if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): + return False + return mla_decode_branch + + +def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: + """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). + + ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers + and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input + here is fixed by that shape, so a captured step and its replays are classified alike. + """ + if num_tokens <= 0: + return None + requests = _decode_requests(attn_metadata, num_tokens) + if requests is not None: + return DecodeStep(num_tokens, requests, num_tokens // requests) + if num_tokens <= DECODE_MAX_TOKENS: + return DecodeStep(num_tokens) + return None + + +def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: + """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" + if attn_metadata.num_contexts != 0: + return None + requests = attn_metadata.num_generations + if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: + return None + tokens = num_tokens // requests + if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: + return None + seq_lens = getattr(attn_metadata, "seq_lens", None) + if seq_lens is not None and seq_lens.device.type == "cpu": + lens = seq_lens[:requests] + if lens.numel() != requests or bool((lens != tokens).any()): + return None + return requests + + +def _check_construction(model_config: ModelConfig, target: K3TargetSpec) -> None: + """Validate the concrete target before constructing modules or workspaces.""" + capability = torch.cuda.get_device_capability() + assert capability == target.sm, ( + f"this target requires sm_{target.sm[0]}{target.sm[1]}, running on " + f"sm_{capability[0]}{capability[1]}" + ) + mapping = model_config.mapping + topology = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.enable_attention_dp, + ) + parallel = f"tp16_moetp{target.moe_tp}ep{target.moe_ep}" + assert topology == (16, 16, 1, target.moe_tp, target.moe_ep, False), ( + f"the {parallel} target needs world_size 16, tensor_parallel_size 16, " + f"pipeline_parallel_size 1, moe_tensor_parallel_size {target.moe_tp}, " + f"moe_expert_parallel_size {target.moe_ep} and enable_attention_dp false; " + f"the engine built (world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" + ) + if target.moe_tp == 16: + assert mapping.moe_tp_ep_user_specified, ( + f"the {parallel} target needs moe_tensor_parallel_size 16 and " + "moe_expert_parallel_size 1 set explicitly" + ) + assert getattr(model_config, "spec_config", None) is None, ( + f"the {parallel} target decodes without speculation; for DSpark, DFlash " + "or SA set moe_tensor_parallel_size 4 and moe_expert_parallel_size 4" + ) + assert model_config.torch_dtype == torch.bfloat16, ( + f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" + ) + if target.quant_algo == QuantAlgo.NVFP4: + for layer_idx in range(1, 93): + quant = KimiK3MoERuntime._resolve_routed_quant_config(model_config, layer_idx) + if quant.quant_algo != target.quant_algo or quant.group_size != 16: + raise ValueError( + f"NVFP4 target requires group-size-16 NVFP4 experts at layer {layer_idx}; " + f"got {quant.quant_algo}, group_size={quant.group_size}" + ) + else: + quant = model_config.quant_config + assert quant.kv_cache_quant_algo is None, ( + "this target's MLA kernels read a bf16 KV pool; " + f"kv_cache_config.dtype resolved to {quant.kv_cache_quant_algo}" + ) + assert quant.quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( + "this target loads the MXFP4 checkpoint with no per-layer declarations; " + f"the engine read {quant.quant_algo} with " + f"{len(model_config.quant_config_dict or {})} per-layer declarations" + ) + strategy = model_config.allreduce_strategy + assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( + f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" + ) + + +class ModelingV2KimiK3(KimiLinearForCausalLM): + """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. + + It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through + weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model + keeps the built-in one's. + """ + + target_spec: K3TargetSpec + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The V2 hybrid manager holds the KDA states; the step contract requires it.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + config = model_config.pretrained_config + assert getattr(config, "text_config", None) is not None, ( + "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " + "text_config" + ) + _check_construction(model_config, self.target_spec) + text = _text_model_config(model_config) + spec_config = getattr(text, "spec_config", None) + assert ( + spec_config is None + or spec_config.spec_dec_mode.is_sa() + or spec_config.spec_dec_mode.is_dflash() + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" + # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 + # weight-read conversion of the shared / latent MLPs. + self._fp8_weight_read_moe_mlp = False + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + cfg = text.pretrained_config + SpecDecOneEngineForCausalLM.__init__( + self, + KimiLinearModel(text), + text, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self._own_spec_branch(text) + self._step_checked = False + self._decode_state_ready = False + # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the + # processor the speculative worker's target logits and a parallel drafter's logits call too. + stock_logits_processor = self.logits_processor + self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) + draft_model = getattr(self, "draft_model", None) + if getattr(draft_model, "logits_processor", None) is stock_logits_processor: + draft_model.logits_processor = self.logits_processor + # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which + # must therefore be the text config, as the built-in wrapper leaves it. + model_config._frozen = False + model_config.pretrained_config = self.config + model_config._frozen = True + + def _own_spec_branch(self, model_config: ModelConfig) -> None: + """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a + standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is + `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" + spec_config = model_config.spec_config + if spec_config is None: + return + mode = spec_config.spec_dec_mode + standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) + if ( + standalone_dspark + and self.draft_config is not None + and is_mla(self.draft_config.pretrained_config) + ): + # The MLA drafter's compressed context cache requires its stock worker. + return + if ( + standalone_dspark + and self.draft_config is not None + and not is_mla(self.draft_config.pretrained_config) + ): + drafter = K3DSparkDrafter( + self.draft_config, dflash_attention_backend=spec_config.attention_backend + ) + if mode.is_parallel_draft(): + drafter.logits_processor = self.logits_processor + self.epilogue[self.epilogue.index(self.draft_model)] = drafter + self.draft_model = drafter + if mode.is_dflash(): + worker_cls = _spec_worker.KimiK3DFlashWorker + elif standalone_dspark: + worker_cls = _spec_worker.KimiK3DSparkWorker + else: + return + worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) + worker.set_draft_model(self.draft_model) + self.epilogue[self.epilogue.index(self.spec_worker)] = worker + self.spec_worker = worker + + def load_weights(self, weights, *args, **kwargs): + _weights.load(self, weights) + + def cache_derived_state(self) -> None: + """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of + every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: + a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites + are among them, and the drafter gets the state too.""" + super().cache_derived_state() + if self.model.decode_gemvs is not None: + return + sites = tuple(_decode_gemv.SITES) + if self.target_spec.quant_algo == QuantAlgo.NVFP4: + sites = ("dense_gate_up", "dense_down") + if isinstance(self.draft_model, K3DSparkDrafter): + sites += ("drafter_qkv", "drafter_o", "drafter_gate_up", "drafter_down") + gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head, sites=sites) + self.model.decode_gemvs = gemvs + self.logits_processor.gemvs = gemvs + for layer in self.model.layers: + if not layer.is_moe: + layer.decode_gemvs = gemvs + if isinstance(self.draft_model, K3DSparkDrafter): + self.draft_model.decode_gemvs = gemvs + + def post_load_weights(self) -> None: + """The state the decode kernels share, built once per device before any CUDA-graph capture and handed to + every layer of its kind: the KDA projection's Lamport buffers and the MLA decode attention's workspace; the + decode GEMVs' state (built by ``cache_derived_state``) handed to every attention module; and, where every + attention all-reduce runs over MNNVL, the TP group's collective state (``K3DecodeComm``, collective: every + rank builds it here) handed to every layer, with whether its sandwich takes the layer's o_proj, the decode + path's one-shot ceiling on every stock MNNVL all-reduce (``use_decode_one_shot``), and the MoE decode path + (``_build_decode_moe``); then the speculative worker's decode kernels (``_gate_spec_worker_kernels``) and the + DSpark drafter's collectives (``_gate_drafter_comm``).""" + super().post_load_weights() + if self._decode_state_ready: + return + if self.target_spec.quant_algo == QuantAlgo.NVFP4: + layers = self.model.layers + comm = None + if all(layer._mnnvl_allreduce() is not None for layer in layers): + comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, None) + for layer in layers: + layer.decode_comm = comm + layer.sandwich_oproj = False + _decode_comm.use_decode_one_shot(self) + self._gate_spec_worker_kernels(comm) + self._gate_drafter_comm(comm) + self._decode_state_ready = True + logger.info( + f"Kimi K3 NVFP4 SM{self.target_spec.sm[0]}{self.target_spec.sm[1]}: " + "stock quantized attention and NVFP4 experts; " + "BF16 decode GEMVs and DSpark specializations; " + f"attention-residual MNNVL fusion={comm is not None}" + ) + return + kda = [layer.linear_attn for layer in self.model.layers if layer.is_kda] + mla = [layer.self_attn.mixer for layer in self.model.layers if not layer.is_kda] + if kda[0].k3_buffers is not None: + return # built by an earlier call; CUDA graphs captured since hold it + device = self.model.embed_tokens.weight.device + buffers = K3KdaBuffers.create(device) + workspace = K3MlaAttnWorkspace.create(device, mla[0].num_heads_tp // 6) + for module in kda: + module.k3_buffers = buffers + for module in mla: + module.k3_workspace = workspace + for module in kda + mla: + module.decode_gemvs = self.model.decode_gemvs + layers = self.model.layers + comm = None + if all(layer._mnnvl_allreduce() is not None for layer in layers): + for layer in layers: + layer.sandwich_oproj = _decode_comm.K3DecodeComm.takes_oproj( + layer._o_proj(), self.model.num_attn_res_snapshots + ) + oproj = next((layer._o_proj().weight for layer in layers if layer.sandwich_oproj), None) + comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, oproj) + for layer in layers: + layer.decode_comm = comm + _decode_comm.use_decode_one_shot(self) + moe_layers = 0 if comm is None else self._build_decode_moe(comm) + logger.info( + "Kimi K3 decode kernels: KDA on k3_kda_decode_attn, k3_kda_attn and k3_kda_verify " + f"({sum(m.takes_k3_kernels for m in kda)} / {len(kda)} layers take them), MLA on k3_mla_qkv and " + f"k3_mla_attn_vb_out ({len(mla)} layers); the post-attention all-reduce of at most " + f"{_decode_comm.AR_ATTN_RES_MAX_TOKENS} tokens " + + ( + "with the residual update: k3_sandwich_oproj with o_proj at most " + f"{_decode_comm.SANDWICH_MAX_TOKENS} decode tokens " + f"({sum(layer.sandwich_oproj for layer in layers)} / {len(layers)} layers), else " + "mnnvl_allreduce_attn_res" + if comm is not None + else "unfused (an attention all-reduce does not run over MNNVL)" + ) + + f"; MoE on k3_moe_front, k3_moe and the row-parallel tail ({moe_layers} layers)" + ) + self._gate_spec_worker_kernels(comm) + self._gate_drafter_comm(comm) + self._decode_state_ready = True + + def _build_decode_moe(self, comm: _decode_comm.K3DecodeComm) -> int: + """The MoE decode path (``decode_moe.py``) on every MoE layer it takes: the shared state (collective: the + front's head workspace), each layer's decode weights with its latent norm folded into the latent up + projection, the decode GEMVs' state; then one call of every kernel of the path (collective: the front and + the sandwich tail). Every rank builds it here. Returns the number of MoE layers it takes.""" + mapping = self.model_config.mapping + decode_moe = _decode_moe + if mapping.moe_tp_size == 16: + from . import decode_moe_tp16 as decode_moe + moes = [layer.block_sparse_moe for layer in self.model.layers if layer.is_moe] + gaps = { + id(moe): decode_moe.layout_gaps(moe, mapping.tp_size, self.model.num_attn_res_snapshots) + for moe in moes + } + takes = [moe for moe in moes if not gaps[id(moe)]] + for reason in sorted({gap for moe in moes for gap in gaps[id(moe)]}): + logger.info_once( + f"Kimi K3 MoE decode path off on some layers: {reason}", + key=f"k3_decode_moe_off_{reason}", + ) + if not takes: + return 0 + backend = takes[0].routed_experts.backend + state = decode_moe.K3DecodeMoe.create( + mapping, + backend.w3_w1_weight.device, + backend.w3_w1_weight.shape[1] // 2, + backend.expert_size_per_partition, + comm.mnnvl, + ) + for moe in takes: + decode_moe.fold_latent_norm(moe) + moe.decode_moe = decode_moe.K3DecodeMoeLayer.create( + moe, state, mapping.tp_rank, mapping.tp_size + ) + moe.decode_gemvs = self.model.decode_gemvs + first = takes[0].decode_moe + first.warm_up(takes[0]) + logger.info( + "Kimi K3 MoE decode path: the latent all-reduce at <= 8 tokens on " + + ( + "the routed experts' push form and k3_latent_reduce on pushing steps (a pure decode step captured " + "into a CUDA graph on the decode kernels), else the routed experts' all-reduce" + if state.exchange is not None + else f"the routed experts' all-reduce (no latent exchange at TP {mapping.tp_size})" + ) + ) + comm.compile_tail(takes[0].moe_hidden_size, first.shared_cols, first.tail_weight) + return len(takes) + + def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, + ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input + that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's + ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the + vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks + each step's own conditions. Returns the setting (False without such a worker).""" + worker = getattr(self, "spec_worker", None) + if not hasattr(worker, "k3_decode"): + return False + gemvs = self.model.decode_gemvs + worker.k3_decode = ( + comm is not None and gemvs is not None and gemvs.head_workspace is not None + ) + logger.info( + f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " + + ( + "on" + if worker.k3_decode + else "off (no MNNVL decode state or no k3_head_gemv LM head)" + ) + ) + return worker.k3_decode + + def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built + it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context + projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms + in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on + the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate + requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` + workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the + drafter took the state (False without such a drafter).""" + drafter = getattr(self, "draft_model", None) + if comm is None or not isinstance(drafter, K3DSparkDrafter): + return False + drafter.use_decode_comm(comm) + return True + + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """First-forward checks of the engine surface and the per-engine settings.""" + objects = { + "attn_metadata": attn_metadata, + "kv_cache_manager": attn_metadata.kv_cache_manager, + } + missing = [ + f"{owner}.{name}" + for owner, names in REQUIRED_ENGINE_FIELDS.items() + for name in names + if not hasattr(objects[owner], name) + ] + assert not missing, f"engine fields this target reads are missing: {missing}" + manager = attn_metadata.kv_cache_manager + assert isinstance(manager, MambaHybridCacheManagerV2), ( + "this target needs the V2 hybrid KV / state cache manager " + "(kv_cache_config.use_kv_cache_manager_v2); the engine built " + f"{type(manager).__name__}" + ) + if self.target_spec.quant_algo != QuantAlgo.NVFP4: + assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( + f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " + f"{attn_metadata.tokens_per_block}" + ) + assert not manager.enable_block_reuse, ( + "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" + ) + kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) + state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype + assert state_dtype == torch.float32, ( + "this target's KDA kernels keep fp32 recurrent states " + "(kv_cache_config.mamba_ssm_cache_dtype float32 or " + f"auto); the engine built a {state_dtype} state pool" + ) + self._step_checked = True + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + if kwargs.pop("multimodal_params", None): + raise ValueError( + "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" + ) + if not self._step_checked: + self._check_step_contract(attn_metadata) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py index 09abb6fd26f4..d6b30f6c68bc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/routing.py @@ -20,11 +20,6 @@ from ..._router_index import NoTarget -# The GPU architecture the MXFP4 targets are written for. sm is part of a -# target's identity, not a knob: a different SM is a different target. The K3 -# decode kernels (the tcgen05 GEMVs, the fused MoE) are built for GB200 only. -_SM = (10, 0) - # The MXFP4 checkpoint's config shape; see the note in the sibling gpt_oss # routing module for what shape-sniffing does and does not pin. # @@ -145,10 +140,24 @@ def route( ) quant_algo = _quant_algo(pretrained_config) parallel = _parallel(mapping) - mxfp4_on_sm100 = ( - compute_capability == _SM and shape == _SHAPE and quant_algo == QuantAlgo.W4A16_MXFP4 + mxfp4_on_blackwell = ( + compute_capability in ((10, 0), (10, 3)) + and shape == _SHAPE + and quant_algo == QuantAlgo.W4A16_MXFP4 ) + if ( + compute_capability == (10, 3) + and shape == _SHAPE + and parallel == "tp16_moetp4ep4" + and _has_nvfp4_experts(pretrained_config) + ): + from .kimi_k3_nvfp4__sm_103__tp16_moetp4ep4.modeling import ( + ModelingV2KimiK3Nvfp4Sm103Tp16Moetp4ep4, + ) + + return ModelingV2KimiK3Nvfp4Sm103Tp16Moetp4ep4 + if ( compute_capability == (10, 7) and shape == _SHAPE @@ -161,14 +170,28 @@ def route( return ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4 - if mxfp4_on_sm100 and parallel == "tp16_moetp4ep4": + if mxfp4_on_blackwell and compute_capability == (10, 3) and parallel == "tp16_moetp4ep4": + from .kimi_k3_mxfp4__sm_103__tp16_moetp4ep4.modeling import ( + ModelingV2KimiK3Mxfp4Sm103Tp16Moetp4ep4, + ) + + return ModelingV2KimiK3Mxfp4Sm103Tp16Moetp4ep4 + + if mxfp4_on_blackwell and compute_capability == (10, 3) and parallel == "tp16_moetp16ep1": + from .kimi_k3_mxfp4__sm_103__tp16_moetp16ep1.modeling import ( + ModelingV2KimiK3Mxfp4Sm103Tp16Moetp16ep1, + ) + + return ModelingV2KimiK3Mxfp4Sm103Tp16Moetp16ep1 + + if mxfp4_on_blackwell and parallel == "tp16_moetp4ep4": from .kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4, ) return ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4 - if mxfp4_on_sm100 and parallel == "tp16_moetp16ep1": + if mxfp4_on_blackwell and parallel == "tp16_moetp16ep1": from .kimi_k3_mxfp4__sm_100__tp16_moetp16ep1.modeling import ( ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1, ) @@ -182,7 +205,7 @@ def route( f"moe_tp={mapping.moe_tp_size} moe_ep={mapping.moe_ep_size} " f"split_set={getattr(mapping, 'moe_tp_ep_user_specified', False)} " f"attention_dp={mapping.enable_attention_dp}; the kimi_k3_vl targets are the MXFP4 " - "Kimi K3 checkpoint on sm100 (GB200), 16 ranks of tensor-parallel attention, with the " + "Kimi K3 checkpoint on sm100/sm103, 16 ranks of tensor-parallel attention, with the " "experts split moe_tp 4 x moe_ep 4 or, set explicitly, moe_tp 16 x moe_ep 1; " - "and group-size-16 NVFP4 Kimi K3 on sm107 with tp16 attention and moe_tp 4 x moe_ep 4" + "and group-size-16 NVFP4 Kimi K3 on sm103/sm107 with tp16 attention and moe_tp 4 x moe_ep 4" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/spec_worker.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/spec_worker.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/weights.py new file mode 100644 index 000000000000..915a6df08011 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/weights.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight loading for the Kimi K3 text targets. + +The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and +`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: + +* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the + routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under + moe_tp 4 x moe_ep 4), and checks its own key coverage. +* The vision tower and the projector are a predicted non-load: listed here, never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_utils import filter_weights + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") + + +def load(model, weights) -> None: + """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) + assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = _LANG_PREFIX + KimiLinearForCausalLM.load_weights(model, lm_weights) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_fused_moe/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_fused_moe/op.py index 813e9ef4b7d2..69a33c857ffa 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/k3_fused_moe/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_fused_moe/op.py @@ -87,8 +87,11 @@ def is_supported(w3_w1_weight: torch.Tensor, w3_w1_weight_scale: torch.Tensor, w """Whether these TRTLLM-Gen W4A8_MXFP4_MXFP8 buffers fit the fused kernels. Only metadata is read, so the buffers may still be on the meta device.""" - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - return False, "needs sm_100" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ( + (10, 0), + (10, 3), + ): + return False, "needs sm_100 or sm_103" try: import cutlass # noqa: F401 except ImportError: @@ -698,8 +701,11 @@ def m1_supported( logical intermediate; the buffers may hold it zero-padded to a multiple of 128, as TRT-LLM's loader lays out an unaligned shard (e.g. 192 -> 256 at TP16). Only metadata is read, so the buffers may still be on the meta device.""" - if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): - return False, "needs sm_100" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ( + (10, 0), + (10, 3), + ): + return False, "needs sm_100 or sm_103" try: import cutlass # noqa: F401 except ImportError: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index e0590d9d25f8..0b3e337dd31e 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -62,6 +62,8 @@ l0_cpu: # compute capability, no device. The rest of unittest/_torch/modeling_v2 needs a GPU. - unittest/_torch/modeling_v2/test_modeling_v2_claims.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py + - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py + - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py # l0_b300 selects unittest/_torch/modeling_v2 with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order diff --git a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py index d1d5dc3a72ec..a9d177df213c 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_decode_comm_op_matrix.py @@ -1017,8 +1017,7 @@ def _run_one_rank(args) -> int: import importlib T = importlib.import_module( - "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." - "kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling" + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling" ) # First, before anything exists as an inference tensor: the decode MoE as post_load_weights builds it. code = ls.run_checks(R, [check_decode_moe_from_model_parameters]) diff --git a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py b/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py index afda4d0112f7..2f7c6caef9ae 100644 --- a/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/comm/_kimi_k3_route_b_decode_moe_op_matrix.py @@ -377,10 +377,11 @@ def _run_one_rank(args) -> int: import importlib T = importlib.import_module( - "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." - "kimi_k3_mxfp4__sm_100__tp16_moetp16ep1.modeling" + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling" + ) + DM = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.decode_moe_tp16" ) - DM = T._decode_moe # First, before anything exists as an inference tensor: the decode MoE as post_load_weights builds it. code = ls.run_checks(R, [check_build]) if code: diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 59d0f48a8559..9d20ff46b988 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -128,13 +128,9 @@ def test_every_target_ships_both_products(): assert (target_dir / product).is_file(), f"{name}: missing {product}" -def test_targets_do_not_share_files(): - """Isolation is the property being demonstrated; measure it. - - Targets are allowed to import the catalog and nothing else of each - other's. A shared helper between two targets is the first step back to - the abstraction this package exists to avoid. - """ +def test_targets_only_import_their_declared_family_implementation(): + """K3 targets may share their family implementation, never another family's.""" + k3_shared_modules = {"modeling", "weights", "decode_comm", "decode_gemv", "spec_worker"} for arch in _ARCHS: for target_dir, name in _routed_targets(arch): for source_file in target_dir.glob("*.py"): @@ -143,6 +139,12 @@ def test_targets_do_not_share_files(): dots, tail = match.group(1), match.group(2) if len(dots) == 1: continue # sibling within the target + if ( + target_dir.parent.name == "kimi_k3_vl" + and len(dots) == 2 + and tail in k3_shared_modules + ): + continue assert tail.startswith("catalog"), ( f"{name}: {source_file.name} reaches outside its own " f"directory for {tail!r}; targets may import the " @@ -189,7 +191,10 @@ def test_uncertified_generic_calls_name_every_stock_import(): checked = 0 for arch in _ARCHS: for target_dir, name in _routed_targets(arch): - tree = ast.parse((target_dir / "modeling.py").read_text()) + source_file = target_dir / "modeling.py" + if target_dir.parent.name == "kimi_k3_vl": + source_file = target_dir.parent / "modeling.py" + tree = ast.parse(source_file.read_text()) declared = _declared_tuple(tree, "UNCERTIFIED_GENERIC_CALLS") if declared is None: continue diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py index d50944d4a23b..53d3fc4b5dd6 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_construction.py @@ -3,11 +3,13 @@ """The Kimi K3 targets' construction checks, host-side: each target requires its own parallel layout, and ``tp16_moetp16ep1`` also an explicit expert split and no speculative decoding.""" +import importlib import types import pytest import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as shared from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 modeling as route_a, ) @@ -76,14 +78,50 @@ def test_tp16_moetp16ep1_decodes_without_speculation(): route_b._check_construction(_config(16, 1, spec_config=object())) +@pytest.mark.parametrize("moe_tp,moe_ep", [(4, 4), (16, 1)]) +def test_sm103_mxfp4_requires_its_hardware(monkeypatch, moe_tp, moe_ep): + target = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + f"kimi_k3_mxfp4__sm_103__tp16_moetp{moe_tp}ep{moe_ep}.modeling" + ) + with pytest.raises(AssertionError, match="requires sm_103"): + target._check_construction(_config(moe_tp, moe_ep)) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (10, 3)) + target._check_construction(_config(moe_tp, moe_ep)) + + +@pytest.mark.parametrize("sm", [103, 107]) +def test_nvfp4_checks_every_expert_layer(monkeypatch, sm): + from tensorrt_llm.models.modeling_utils import QuantConfig + + target = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + f"kimi_k3_nvfp4__sm_{sm}__tp16_moetp4ep4.modeling" + ) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (sm // 10, sm % 10)) + config = _config(4, 4) + config.quant_config = QuantConfig( + quant_algo=QuantAlgo.MIXED_PRECISION, kv_cache_quant_algo=QuantAlgo.NVFP4 + ) + config.quant_config_dict = { + f"model.layers.{i}.mlp.experts": QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16) + for i in range(1, 93) + } + target._check_construction(config) + config.quant_config_dict["model.layers.92.mlp.experts"].group_size = 32 + with pytest.raises(ValueError, match="group-size-16 NVFP4 experts at layer 92"): + target._check_construction(config) + + def test_tp16_moetp16ep1_keeps_no_per_token_kda_states(monkeypatch): """The text model answers ``kda_draft_records`` (False, without speculation), which a captured step's latent-push decision reads: a pure decode step on the decode kernels pushes, as on ``tp16_moetp4ep4`` without per-token states.""" - model = route_b.KimiLinearModel.__new__(route_b.KimiLinearModel) + model = shared.KimiLinearModel.__new__(shared.KimiLinearModel) + object.__setattr__(model, "model_config", types.SimpleNamespace(spec_config=None)) assert model.kda_draft_records is False monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - monkeypatch.setattr(route_b, "is_in_breakable_cuda_graph", lambda: False) + monkeypatch.setattr(shared, "is_in_breakable_cuda_graph", lambda: False) kda = types.SimpleNamespace( is_kda=True, linear_attn=types.SimpleNamespace(takes_k3_kernels=True, k3_buffers=object()), @@ -93,5 +131,5 @@ def test_tp16_moetp16ep1_keeps_no_per_token_kda_states(monkeypatch): self_attn=types.SimpleNamespace(will_run_decode_branch=lambda metadata, step: True), ) object.__setattr__(model, "layers", [kda, mla]) - assert model._latent_push(None, route_b.DecodeStep(1, 1, 1)) - assert not model._latent_push(None, route_b.DecodeStep(8, 1, 8)) + assert model._latent_push(None, shared.DecodeStep(1, 1, 1)) + assert not model._latent_push(None, shared.DecodeStep(8, 1, 8)) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py index 195562ee38a0..3f3c7c2fbaa1 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_gemv.py @@ -34,7 +34,7 @@ k3_ctm_gemv_wide, ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 decode_gemv, ) from tensorrt_llm._torch.modules.embedding import LMHead diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py index 6ebe162c626d..3dc3deb5918f 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_decode_step.py @@ -9,18 +9,15 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 - modeling as route_a, +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 + modeling as target, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( # noqa: E501 DecodeStep, _attn_res_max_tokens, decode_step, latent_push, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 - modeling as route_b, -) pytestmark = pytest.mark.cpu_only @@ -156,7 +153,6 @@ def forward_generic( return hidden_states, num_snapshots -@pytest.mark.parametrize("target", [route_a, route_b], ids=["tp16_moetp4ep4", "tp16_moetp16ep1"]) @pytest.mark.parametrize( "seq_lens,num_contexts,rows,path", [ @@ -174,7 +170,7 @@ def forward_generic( ([8, 4], 0, 12, "generic"), ], ) -def test_each_step_takes_one_path_through_every_layer(target, seq_lens, num_contexts, rows, path): +def test_each_step_takes_one_path_through_every_layer(seq_lens, num_contexts, rows, path): model = object.__new__(target.KimiLinearModel) torch.nn.Module.__init__(model) calls = [] diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py index dbaf003daf05..2e3322ad7a49 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drafter.py @@ -38,11 +38,11 @@ import torch from transformers import Qwen3Config -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 decode_comm, decode_gemv, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 modeling as target, ) from tensorrt_llm._torch.model_config import ModelConfig diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py index 716249d68953..dd9f989873d0 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py @@ -1,135 +1,90 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The Kimi K3 `tp16_moetp16ep1` target is a copy of `tp16_moetp4ep4`'s files, changed only in marked blocks. +"""K3 targets share the model lifecycle while preserving each checkpoint's layout.""" -Targets share no files (test_modeling_v2_claims.py), so `tp16_moetp16ep1` (route B) ships its own copy of every -module of `tp16_moetp4ep4` (route A). Each place the copy differs is a block: - - # >>> route B: - - # <<< route B - -This reads both targets' files and imports neither. It fails on any difference outside the blocks, so a change to -route A's modules reaches route B's copy, or becomes one of its blocks, in the same commit. -""" - -from __future__ import annotations - -import difflib -from pathlib import Path -from typing import List, Tuple +import importlib +from types import SimpleNamespace import pytest - -import tensorrt_llm._torch._experimental.modeling_v2 as _modeling_v2 - -# The package, not this file: the targets live under tensorrt_llm/ while this test lives under tests/. -_FAMILY = Path(_modeling_v2.__file__).resolve().parent / "models" / "kimi_k3_vl" -_ROUTE_A = _FAMILY / "kimi_k3_mxfp4__sm_100__tp16_moetp4ep4" -_ROUTE_B = _FAMILY / "kimi_k3_mxfp4__sm_100__tp16_moetp16ep1" - -_BEGIN = "# >>> route B" -_END = "# <<< route B" - -# Each target's package marker holds only its own one-line description. -_NOT_COPIED = {"__init__.py"} - - -def _modules(target: Path) -> List[str]: - return sorted(p.name for p in target.glob("*.py") if p.name not in _NOT_COPIED) - - -def _blocks(lines: List[str]) -> List[Tuple[int, int]]: - """The route B blocks as (begin, end) indices of their marker lines; asserts the markers are well formed.""" - blocks, begin = [], None - for i, line in enumerate(lines): - text = line.strip() - if text.startswith(_BEGIN): - assert begin is None, ( - f"line {i + 1}: a block opens inside the one opened on line {begin + 1}" - ) - assert text[len(_BEGIN) :].startswith(":") and text[len(_BEGIN) + 1 :].strip(), ( - f"line {i + 1}: a block opens with '{_BEGIN}: '" - ) - begin = i - elif text.startswith(_END): - assert text == _END, f"line {i + 1}: a block closes with '{_END}' alone" - assert begin is not None, f"line {i + 1}: a block closes without one open" - blocks.append((begin, i)) - begin = None - assert begin is None, f"line {begin + 1}: a block is never closed" - return blocks - - -def _spans(lines: List[str], blocks: List[Tuple[int, int]]) -> List[Tuple[int, int]]: - """Each block with the blank lines around it, which the formatters add and drop next to a comment.""" - spans = [] - for begin, end in blocks: - while begin > 0 and not lines[begin - 1].strip(): - begin -= 1 - while end + 1 < len(lines) and not lines[end + 1].strip(): - end += 1 - spans.append((begin, end)) - return spans - - -def _inside(spans: List[Tuple[int, int]], j1: int, j2: int) -> bool: - """Whether route B's lines [j1, j2), or the point between lines j1 - 1 and j1 when j1 == j2, are in one block.""" - if j1 == j2: - return any(begin < j1 <= end for begin, end in spans) - return any(begin <= j1 and j2 <= end + 1 for begin, end in spans) - - -def _drift(a: List[str], b: List[str]) -> List[str]: - """Each difference between route A's lines ``a`` and route B's ``b`` outside route B's blocks, as text.""" - assert not _blocks(a) and not any(_BEGIN in line or _END in line for line in a), ( - "route A's files have no route B markers" +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize( + "sm,fmt,moe_tp,moe_ep", + [ + (100, "mxfp4", 4, 4), + (100, "mxfp4", 16, 1), + (103, "mxfp4", 4, 4), + (103, "mxfp4", 16, 1), + (103, "nvfp4", 4, 4), + (107, "nvfp4", 4, 4), + ], +) +def test_targets_share_the_model_lifecycle(sm: int, fmt: str, moe_tp: int, moe_ep: int) -> None: + module = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + f"kimi_k3_{fmt}__sm_{sm}__tp16_moetp{moe_tp}ep{moe_ep}.modeling" ) - blocks = _spans(b, _blocks(b)) - drift = [] - matcher = difflib.SequenceMatcher(None, a, b, autojunk=False) - for tag, i1, i2, j1, j2 in matcher.get_opcodes(): - if tag == "equal" or _inside(blocks, j1, j2): - continue - shown = [f" route A {i + 1}: {a[i]}" for i in range(i1, min(i2, i1 + 5))] - shown += [f" route B {j + 1}: {b[j]}" for j in range(j1, min(j2, j1 + 5))] - drift.append( - f"route A lines {i1 + 1}-{i2}, route B lines {j1 + 1}-{j2}:\n" + "\n".join(shown) + target = getattr(module, f"ModelingV2KimiK3{fmt.capitalize()}Sm{sm}Tp16Moetp{moe_tp}ep{moe_ep}") + assert target.forward is modeling.ModelingV2KimiK3.forward + assert target.load_weights is modeling.ModelingV2KimiK3.load_weights + assert target.cache_derived_state is modeling.ModelingV2KimiK3.cache_derived_state + assert target.post_load_weights is modeling.ModelingV2KimiK3.post_load_weights + assert target.target_spec.sm == (sm // 10, sm % 10) + assert (target.target_spec.moe_tp, target.target_spec.moe_ep) == (moe_tp, moe_ep) + + +@pytest.mark.parametrize("algo", [QuantAlgo.W4A8_MXFP4_MXFP8, QuantAlgo.NVFP4]) +def test_streaming_loader_preserves_checkpoint_tensors(algo: QuantAlgo) -> None: + """Each loader passes the exact packed tensors and scales to its quant method.""" + calls = [] + + def record(backend: object, **kwargs: object) -> None: + calls.append(kwargs) + + backend = SimpleNamespace( + quant_method=SimpleNamespace( + load_packed_mxfp4_expert=record, + load_streaming_nvfp4_expert=record, ) - return drift - - -def test_route_b_copies_every_module(): - assert _modules(_ROUTE_A), f"no modules in {_ROUTE_A}" - assert _modules(_ROUTE_B) == _modules(_ROUTE_A) - - -@pytest.mark.parametrize("name", _modules(_ROUTE_A)) -def test_route_b_matches_route_a_outside_its_blocks(name): - a = (_ROUTE_A / name).read_text(encoding="utf-8").splitlines() - b = (_ROUTE_B / name).read_text(encoding="utf-8").splitlines() - drift = _drift(a, b) - assert not drift, ( - f"{name}: route B's copy differs from route A's outside its blocks. Carry route A's " - "change into the copy, or mark route B's lines as a block that says why they differ:\n" - + "\n".join(drift) ) - - -def test_the_check_sees_drift_only_outside_the_blocks(): - a = ["x = 1", "y = 2", "z = 3"] - changed = ["x = 1", "# >>> route B: why", "y = 4", "# <<< route B", "z = 3"] - dropped = ["x = 1", "", "# >>> route B: why", "# <<< route B", "", "z = 3"] - unmarked = ["x = 1", "# >>> route B: why", "y = 2", "# <<< route B", "z = 4"] - added = ["x = 1", "y = 2", "w = 0", "z = 3"] - spaced = ["x = 1", "", "y = 2", "z = 3"] - assert not _drift(a, changed) - assert not _drift(a, dropped) - assert _drift(a, unmarked) - assert _drift(a, added) - assert _drift(a, spaced) - assert _drift(a[:2], a) - with pytest.raises(AssertionError, match="never closed"): - _drift(a, ["x = 1", "# >>> route B: why", "y = 2", "z = 3"]) - with pytest.raises(AssertionError, match="why it differs"): - _drift(a, ["x = 1", "# >>> route B", "y = 2", "# <<< route B", "z = 3"]) + spec = modeling._k3_expert_ckpt_spec(algo) + tensors = { + f"experts.7.{weight}.{kind}": object() + for weight in ("w1", "w2", "w3") + for kind in spec.kinds + } + spec.loader(backend, "experts", 7, 2, tensors.__getitem__) + assert len(calls) == 1 + loaded = calls[0] + assert loaded.pop("global_expert_id") == 7 + assert loaded.pop("local_slot_id") == 2 + for weight in ("w1", "w2", "w3"): + for kind in spec.kinds: + argument = "weight" if kind == "weight_packed" else kind + assert loaded.pop(f"{weight}_{argument}") is tensors[f"experts.7.{weight}.{kind}"] + assert not loaded + assert spec.needs_layer_finalize == (algo == QuantAlgo.NVFP4) + + +@pytest.mark.parametrize("sm", [(10, 0), (10, 3), (10, 7)]) +def test_mxfp4_fused_kernel_gates_hardware_and_scale_layout(monkeypatch, sm) -> None: + from tensorrt_llm._torch.cute_dsl_kernels.k3_fused_moe.op import is_supported + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: sm) + weights = torch.empty(224, 1536, 1792, dtype=torch.uint8, device="meta") + fc1_scales = torch.empty(224, 1536, 112, dtype=torch.uint8, device="meta") + down = torch.empty(224, 3584, 384, dtype=torch.uint8, device="meta") + fc2_scales = torch.empty(224, 3584, 24, dtype=torch.uint8, device="meta") + supported, _ = is_supported(weights, fc1_scales, down, fc2_scales, 224) + assert supported == (sm in ((10, 0), (10, 3))) + # Group-16 scales cannot be consumed by this group-32 MXFP4 kernel. + nvfp4_scales = torch.empty(224, 1536, 224, dtype=torch.uint8, device="meta") + supported, _ = is_supported(weights, nvfp4_scales, down, fc2_scales, 224) + assert not supported diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py index aef529d49750..31747c0fba80 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_moe_routing.py @@ -12,17 +12,12 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 - modeling as route_a, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 - modeling as route_b, +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 + modeling as shared, ) from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod, RoutingMethodType -TARGETS = pytest.mark.parametrize( - "target", [route_a, route_b], ids=["tp16_moetp4ep4", "tp16_moetp16ep1"] -) +TARGETS = pytest.mark.parametrize("target", [shared], ids=["shared-k3"]) NUM_EXPERTS, TOP_K, HIDDEN, LATENT = 896, 16, 7168, 3584 # One tp16_moetp16ep1 rank's expert width: 192 of 3072, zero-padded to 256 by the loader. diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py index 87ad2b3cb531..c3da9be219c9 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_route_b_moe.py @@ -11,8 +11,8 @@ import torch import tensorrt_llm._torch.cute_dsl_kernels.k3_fused_moe as _k3_fused_moe -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp16ep1 import ( # noqa: E501 - decode_moe as route_b_moe, +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( + decode_moe_tp16 as route_b_moe, ) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index ca11a8d2f831..05fb5eb7ff69 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -406,6 +406,18 @@ def test_kimi_k3_targets_match(mode, text_as_dict, mapping_kwargs, target): assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == target +@pytest.mark.parametrize("mode", ["auto", "require"]) +@pytest.mark.parametrize( + "mapping_kwargs,parallel", + [(_TP16_MOETP4EP4, "Moetp4ep4"), (_TP16_MOETP16EP1, "Moetp16ep1")], +) +def test_kimi_k3_mxfp4_routes_on_sm103(mode, mapping_kwargs, parallel): + config = _model_config(_k3_config(), **mapping_kwargs) + target = modeling_v2_resolve(config, _llm_args(mode)) + assert target.__name__ == f"ModelingV2KimiK3Mxfp4Sm103Tp16{parallel}" + assert target.target_spec.sm == _SM103 + + @pytest.mark.usefixtures("_on_sm100") @pytest.mark.parametrize( "quantization", @@ -452,11 +464,12 @@ def test_kimi_k3_near_misses_do_not_match(text_overrides, mapping_kwargs, missed modeling_v2_resolve(config, _llm_args("require")) -def test_kimi_k3_on_another_sm_does_not_match(): - """The autouse fixture routes as a GB300 (sm 10.3); the K3 targets are sm 10.0 only.""" +def test_kimi_k3_on_another_sm_does_not_match(monkeypatch): + """The MXFP4 fused decode path is scoped to SM100/SM103.""" + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) config = _model_config(_k3_config(), **_TP16_MOETP4EP4) assert modeling_v2_resolve(config, _llm_args("auto")) is None - with pytest.raises(ValueError, match=r"compute capability \(10, 3\)"): + with pytest.raises(ValueError, match=r"compute capability \(10, 7\)"): modeling_v2_resolve(config, _llm_args("require")) @@ -483,21 +496,23 @@ def _k3_nvfp4_mixed_config(): @pytest.mark.parametrize("mode", ["auto", "require"]) @pytest.mark.parametrize("quant_method", ["modelopt", "modelopt_mixed"]) -def test_kimi_k3_nvfp4_routes_on_sm107(monkeypatch, mode, quant_method): - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) +@pytest.mark.parametrize("sm", [103, 107]) +def test_kimi_k3_nvfp4_routes(monkeypatch, mode, quant_method, sm): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (sm // 10, sm % 10)) quantization = _k3_nvfp4_mixed_config() if quant_method == "modelopt_mixed": quantization = dict(quantization["quantization"], quant_method=quant_method) config = _model_config(_k3_config(quantization=quantization), **_TP16_MOETP4EP4) assert _qualified(modeling_v2_resolve(config, _llm_args(mode))) == ( - "models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4.modeling." - "ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4" + f"models.kimi_k3_vl.kimi_k3_nvfp4__sm_{sm}__tp16_moetp4ep4.modeling." + f"ModelingV2KimiK3Nvfp4Sm{sm}Tp16Moetp4ep4" ) @pytest.mark.parametrize("bad_layer", ["missing", "mxfp4", "wrong_group_size"]) -def test_kimi_k3_nvfp4_rejects_incomplete_expert_declarations(monkeypatch, bad_layer): - monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (10, 7)) +@pytest.mark.parametrize("sm", [103, 107]) +def test_kimi_k3_nvfp4_rejects_incomplete_expert_declarations(monkeypatch, bad_layer, sm): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **kw: (sm // 10, sm % 10)) quantization = _k3_nvfp4_mixed_config() layers = quantization["quantization"]["quantized_layers"] key = "language_model.model.layers.92.mlp.experts" @@ -531,14 +546,13 @@ def test_kimi_k3_nvfp4_declines_excluded_experts(monkeypatch: pytest.MonkeyPatch def test_kimi_k3_nvfp4_constructor_uses_per_layer_quantization(): - import importlib from types import SimpleNamespace - target = importlib.import_module( - "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." - "kimi_k3_nvfp4__sm_107__tp16_moetp4ep4.modeling" + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( + KimiK3MoERuntime, ) - moe_runtime = target.KimiK3MoERuntime + + moe_runtime = KimiK3MoERuntime quant, layers = ModelConfig.load_hf_quant_config(_k3_nvfp4_mixed_config(), "TRTLLM") config = SimpleNamespace(quant_config=quant, quant_config_dict=layers) @@ -551,6 +565,22 @@ def test_kimi_k3_nvfp4_constructor_uses_per_layer_quantization(): moe_runtime._resolve_routed_quant_config(config, 92) +@pytest.mark.parametrize("prefix", ["language_model.model.", "model.", ""]) +def test_kimi_k3_nvfp4_experts_honor_excluded_aliases(prefix: str) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( + KimiK3MoERuntime, + ) + from tensorrt_llm.models.modeling_utils import QuantConfig + + name = f"{prefix}layers.92.mlp.experts" + config = ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION, exclude_modules=[name]), + quant_config_dict={name: QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16)}, + ) + with pytest.raises(ValueError, match="excluded from quantization"): + KimiK3MoERuntime._resolve_routed_quant_config(config, 92) + + def test_kimi_k3_nvfp4_keeps_mla_dspark_worker(): """The reference MLA drafter must retain its compressed-context-cache worker.""" import importlib @@ -583,9 +613,7 @@ def test_kimi_k3_nvfp4_keeps_mla_dspark_worker(): @pytest.mark.parametrize("prefix", ["language_model.model.", "model.", ""]) def test_kimi_k3_nvfp4_attention_quantization_preserves_config(prefix: str) -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( - modeling as target, - ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target from tensorrt_llm.models.modeling_utils import QuantConfig projection = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) @@ -605,9 +633,7 @@ def test_kimi_k3_nvfp4_attention_quantization_preserves_config(prefix: str) -> N def test_kimi_k3_nvfp4_attention_honors_exclusions() -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( - modeling as target, - ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target from tensorrt_llm.models.modeling_utils import QuantConfig config = ModelConfig( @@ -631,9 +657,7 @@ def test_kimi_k3_nvfp4_attention_honors_exclusions() -> None: def test_kimi_k3_nvfp4_attention_rejects_unsupported_quantization( algo: QuantAlgo, group_size: int ) -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( - modeling as target, - ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target from tensorrt_llm.models.modeling_utils import QuantConfig config = ModelConfig( @@ -647,9 +671,7 @@ def test_kimi_k3_nvfp4_attention_rejects_unsupported_quantization( def test_kimi_k3_nvfp4_attention_rejects_conflicting_block_sizes() -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( - modeling as target, - ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target from tensorrt_llm.models.modeling_utils import QuantConfig config = ModelConfig( diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index d486111e565f..85d205860ad1 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -1970,13 +1970,11 @@ def test_modeling_v2_nvfp4_streamed_moe_matches_reference(num_tokens: int) -> No """Exercise the target's loader and router without allocating the full K3 model.""" from transformers import PretrainedConfig - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4__sm_107__tp16_moetp4ep4 import ( - modeling as target, - ) + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target from tensorrt_llm.models.modeling_utils import QuantConfig - if torch.cuda.get_device_capability() != (10, 7): - pytest.skip("The NVFP4 modeling_v2 target requires SM107") + if torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)): + pytest.skip("NVFP4 MoE requires SM100, SM103 or SM107") torch.manual_seed(91) cfg = _K3Config( hidden_size=_TP_HIDDEN, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_k3_decode_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_k3_decode_worker.py index 55bdaa0ae8bd..5349482a7279 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_k3_decode_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_k3_decode_worker.py @@ -31,7 +31,7 @@ import torch from torch import nn -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.spec_worker import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.spec_worker import ( # noqa: E501 KimiK3DFlashWorker, KimiK3DSparkWorker, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_capture_view.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_capture_view.py index a78ac6ac0281..04f080e4ba79 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_capture_view.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_capture_view.py @@ -22,7 +22,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.spec_worker import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.spec_worker import ( # noqa: E501 capture_view, ) from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_ctx_rows_start.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_ctx_rows_start.py index 1106e84de438..06505a8ab0c8 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_ctx_rows_start.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_ctx_rows_start.py @@ -26,10 +26,10 @@ import torch from torch import nn -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( # noqa: E501 K3DSparkDrafter, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.spec_worker import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.spec_worker import ( # noqa: E501 KimiK3DFlashWorker, KimiK3DSparkWorker, dflash_ctx_rows_kwargs, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_draft_block_rows.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_draft_block_rows.py index ebc4da336738..e65774e1ba69 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_draft_block_rows.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_draft_block_rows.py @@ -26,7 +26,7 @@ import torch from torch import nn -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.spec_worker import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.spec_worker import ( # noqa: E501 KimiK3DFlashWorker, KimiK3DSparkWorker, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_drafter_comm.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_drafter_comm.py index 4334a1349157..34a2f3ff0201 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_drafter_comm.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_drafter_comm.py @@ -39,10 +39,10 @@ import torch from torch import nn -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 decode_comm, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 modeling as target, ) from tensorrt_llm._torch.modules.rms_norm import RMSNorm @@ -358,16 +358,14 @@ def test_gate_hands_the_state_to_the_k3_drafter_only(): taken = [] drafter = _drafter() drafter.use_decode_comm = taken.append - assert target.ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4._gate_drafter_comm( - SimpleNamespace(draft_model=drafter), comm - ) + assert target.ModelingV2KimiK3._gate_drafter_comm(SimpleNamespace(draft_model=drafter), comm) assert taken == [comm] - assert not target.ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4._gate_drafter_comm( + assert not target.ModelingV2KimiK3._gate_drafter_comm( SimpleNamespace(draft_model=drafter), None ) assert taken == [comm] for other in (None, SimpleNamespace(use_decode_comm=taken.append)): - assert not target.ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4._gate_drafter_comm( + assert not target.ModelingV2KimiK3._gate_drafter_comm( SimpleNamespace(draft_model=other), comm ) assert taken == [comm] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_worker_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_worker_gate.py index 6ce8a931a926..4d92b269770b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_worker_gate.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_spec_worker_gate.py @@ -29,15 +29,15 @@ import torch from torch import nn -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 decode_gemv, spec_worker, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import ( # noqa: E501 modeling as target_modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_mxfp4__sm_100__tp16_moetp4ep4.modeling import ( # noqa: E501 - ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4 as Target, +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( # noqa: E501 + ModelingV2KimiK3 as Target, ) from tensorrt_llm._torch.speculative.dflash import DFlashWorker from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode @@ -63,7 +63,7 @@ def _target(worker, gemvs=True, head_workspace=True): _SPEC_BRANCHES = { "dspark gqa": ("DSPARK", False, False, True, K3DSpark), - "dspark mla": ("DSPARK", False, True, False, K3DSpark), + "dspark mla": ("DSPARK", False, True, False, None), "dspark embedded": ("DSPARK", True, False, False, None), "dflash": ("DFLASH", False, False, False, K3DFlash), "sa": ("SA", False, False, False, None), From c458efecc716f2bc4f3be3b4e8775b544006b7a5 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 19 Oct 2026 15:52:27 +0800 Subject: [PATCH 3/8] Isolate NVFP4 targets and fuse FP8 attention (By Agent) Restore the original PR19991 MXFP4 SM100 target trees. Keep NVFP4 modeling in an independent copy shared by the SM103 and SM107 targets. Preserve FP8 projections while fusing KDA decode and MLA query, cache update, attention, V-B and output gating. Validation: 13 B300 kernel parity cases including CUDA Graph replay and changing KV lengths; 16 source structure checks. SM107 hardware and full-model validation remain outstanding. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 22 +- .../__init__.py | 2 +- .../decode_comm.py | 402 ++ .../decode_gemv.py | 465 +++ .../decode_moe.py | 495 +++ .../modeling.py | 3529 +++++++++++++++- .../spec_worker.py | 1343 ++++++ .../weights.py | 44 +- .../__init__.py | 2 +- .../decode_comm.py | 402 ++ .../decode_gemv.py | 465 +++ .../decode_moe.py | 408 ++ .../modeling.py | 3592 ++++++++++++++++- .../spec_worker.py | 1343 ++++++ .../weights.py | 42 +- .../kimi_k3_vl/kimi_k3_nvfp4/__init__.py | 3 + .../kimi_k3_vl/kimi_k3_nvfp4/attention.py | 433 ++ .../kimi_k3_vl/kimi_k3_nvfp4/decode_comm.py | 402 ++ .../kimi_k3_vl/kimi_k3_nvfp4/decode_gemv.py | 465 +++ .../kimi_k3_vl/kimi_k3_nvfp4/decode_moe.py | 408 ++ .../kimi_k3_nvfp4/implementation.py | 3403 ++++++++++++++++ .../kimi_k3_vl/kimi_k3_nvfp4/spec_worker.py | 1343 ++++++ .../kimi_k3_vl/kimi_k3_nvfp4/weights.py | 42 + .../modeling.py | 4 +- .../weights.py | 4 +- .../modeling.py | 4 +- .../weights.py | 4 +- .../test_lists/test-db/l0_b200.yml | 1 + .../test_modeling_v2_k3_fp8_fused.py | 282 ++ .../modeling_v2/test_modeling_v2_claims.py | 15 +- .../test_modeling_v2_kimi_k3_drift.py | 40 +- .../modeling_v2/test_modeling_v2_routing.py | 20 +- 32 files changed, 19370 insertions(+), 59 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_comm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_moe.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_comm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_moe.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/spec_worker.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/weights.py create mode 100644 tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 6aba851fa95d..ac03f4362c8a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -10,17 +10,17 @@ entries, sharing nothing with its siblings, and trusted through accuracy gates instead of shared abstractions. The one-to-one correspondence between `models//` here and `modeling_.py` there is the point of the exercise. -Kimi K3 keeps concrete checkpoint, SM and parallelism identities in thin target -classes and shares its implementation within `models/kimi_k3_vl/`. Its MXFP4 and -NVFP4 targets use the same decoder, attention residuals, BF16 GEMVs, speculative -worker and checkpoint streaming lifecycle. The expert checkpoint specification -selects the loader and finalization behavior. MXFP4 uses the K3 MXFP4/MXFP8 fused -decode kernels; NVFP4 uses the existing NVFP4 MoE kernels and stock quantized -attention, preserving FP8 projection metadata and the configured KV-cache format. -The MXFP4 fusion's group-32 E8M0 scales and MXFP8 intermediate differ from NVFP4's -group-16 scales and FP4 activation quantization, so the formats keep independent -MoE kernels. Shared model code does not imply shared accuracy or performance -certification: every hardware and quantization target needs its own gates. +Kimi K3 keeps its original MXFP4/SM100 deployments independent. The NVFP4 +implementation is a separate copy under `models/kimi_k3_vl/kimi_k3_nvfp4/`, +shared by thin SM103 and SM107 target classes. NVFP4 retains checkpoint FP8 +projection weights and activations: decode fuses KDA's gate, convolution, +recurrence and output normalization, and MLA's query/absorption/cache writes +and attention/V-B/output gate. The ordinary FP8 GEMMs remain responsible for +input and output projections. MLA's fused decode currently reads a BF16 KV +cache with 64-token pages; other layouts use the native path. The KDA fused +verify path retains FP32 states and the original per-token draft records. +The NVFP4 expert kernels remain separate from the MXFP4/MXFP8 kernels. +Every hardware and quantization target requires its own accuracy gates. ## Using it diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py index e48707ebcc19..2b24a35badfd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 MXFP4 / sm_100 / tp16_moetp16ep1.""" +"""Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 16 x moe_ep 1.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_comm.py new file mode 100644 index 000000000000..f5327daf9474 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_comm.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's residual-update collectives on the catalog's Kimi K3 MNNVL and sandwich entries. + +A layer's post-attention step adds the attention output to the running prefix sum, selects the attention residual +and applies the post-attention RMSNorm. Under TP the attention output is a sum over the TP group, so the step is that +all-reduce followed by the residual update. On a step of at most `AR_ATTN_RES_MAX_TOKENS` tokens (wide decode steps +aside) one collective runs both: + +* `K3DecodeComm.sandwich_oproj`, `comm/k3_sandwich_oproj`: o_proj, its all-reduce and the residual update in one + kernel, at most `SANDWICH_MAX_TOKENS` tokens of an o_proj of the TP16 per-rank shape [7168, 768]. The attention + hands over its gated o_proj input. +* `K3DecodeComm.allreduce_attn_res`, `comm/mnnvl_allreduce_attn_res`, everywhere else: the one-shot MNNVL all-reduce + of the attention's unreduced o_proj partial, with the residual update as its epilogue. + +The pre-attention step after a MoE layer whose row-parallel tail was handed on (a `PendingTail`, `decode_moe.py`) is +the same kind of collective: `K3DecodeComm.sandwich_tail`, `comm/k3_sandwich_tail`, runs the tail GEMV, its +all-reduce and the next layer's residual update (the final norm's, after the last layer) in one kernel. + +The DSpark drafter (`K3DSparkDrafter`) runs its collectives with a plain residual add and RMSNorm on the same state: + +* `K3DecodeComm.sandwich_plain`, `comm/k3_sandwich_plain`: a row-parallel projection of at most `SANDWICH_MAX_TOKENS` + tokens (a layer's attention output projection, or its MLP's SiLU-and-mul and down projection), its all-reduce, the + residual add and the RMSNorm in one kernel. +* `K3DecodeComm.allreduce_norm`, `comm/mnnvl_fusion_allreduce`: the all-reduce of an unreduced projection output of + any token count the MNNVL workspace holds, with the residual add and the RMSNorm in its epilogue (the context + projection's with a zero residual). + +The state is one `MnnvlWorkspace` and one `K3SandwichWorkspace` of the TP group (`K3DecodeComm.create`): collective +over the group and eager, built by the target in `post_load_weights` before any CUDA-graph capture. Every rank must +make the same calls on each in the same order. Which call a step takes is decided from its token count and kind and +from the load-time layout alone, which every rank of the group shares. + +The plain MNNVL all-reduces the decode path keeps (the stock modules') send one-shot up to +`DECODE_AR_ONE_SHOT_MAX_BYTES` (`use_decode_one_shot`), except a wide decode step's (`wide_all_reduce`). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_oproj import ( + K3SandwichWorkspace, + k3_sandwich_oproj, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_plain import ( + k3_sandwich_plain, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_tail import ( + k3_sandwich_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allreduce_attn_res import ( + mnnvl_allreduce_attn_res, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_fusion_allreduce import ( + mnnvl_fusion_allreduce, + required_buffer_bytes, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) + +# The kernel's support predicate: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_sandwich import op as _sandwich_op +from tensorrt_llm._torch.distributed import AllReduceParams + +# The most tokens of a step whose post-attention all-reduce runs one-shot with the residual update in its epilogue. +AR_ATTN_RES_MAX_TOKENS = 16 + +# The most tokens the sandwich kernel takes (one token tile). +SANDWICH_MAX_TOKENS = _sandwich_op.MAX_TOKENS + +# The MNNVL workspace's buffer size: a one-shot call pushes T x 7168 bf16 from each of 16 ranks, 3.5 MiB at T = 16. +MNNVL_BUFFER_BYTES = 4 << 20 + +# The one-shot ceiling of the stock MNNVL all-reduces on the decode path: 8 tokens x 7168 x 16 ranks x 2 B is +# 1.75 MiB, which the stock 1 MiB ceiling would send two-shot. +DECODE_AR_ONE_SHOT_MAX_BYTES = 4 << 20 + +# A wide decode step's ceiling: the stock 1 MiB. At 16 ranks two-shot is faster for every wide step's rows. +WIDE_AR_ONE_SHOT_MAX_BYTES = 1 << 20 + + +def use_decode_one_shot(model: nn.Module) -> None: + """Every stock MNNVL all-reduce of ``model`` sends one-shot up to `DECODE_AR_ONE_SHOT_MAX_BYTES`. Each grows its + workspace on its first eager call of a larger size, before the capture of that size.""" + for module in model.modules(): + mnnvl = getattr(module, "mnnvl_allreduce", None) + if mnnvl is not None: + mnnvl.one_shot_max_bytes = DECODE_AR_ONE_SHOT_MAX_BYTES + mnnvl.early_trigger = True + + +def skip_all_reduce() -> AllReduceParams: + """All-reduce parameters under which a row-parallel module returns this rank's unreduced output (its projection + without its all-reduce): a ``Linear``'s ``all_reduce_params``, a ``GatedMLP``'s ``final_all_reduce_params``.""" + return AllReduceParams(enable_allreduce=False) + + +def wide_all_reduce(all_reduce: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``all_reduce(x)`` of a wide decode step (a stock ``AllReduce`` module, no fusion): its MNNVL all-reduce with + the `WIDE_AR_ONE_SHOT_MAX_BYTES` ceiling, else the module itself.""" + mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None: + out = mnnvl( + x.contiguous(), AllReduceParams(), one_shot_max_bytes=WIDE_AR_ONE_SHOT_MAX_BYTES + ) + if out is not None: + return out + return all_reduce(x) + + +class PendingTail(NamedTuple): + """A MoE layer's row-parallel tail left to its consumer's fused pre-attention step (``K3DecodeComm.sandwich_tail``): + the reduced latent ``[T, 3584]``, the shared experts' activation, the tail weight ``[latent up columns | padding | + shared down]``, this rank's first latent column and the latent norm's epsilon.""" + + latent: torch.Tensor + act: torch.Tensor + weight: torch.Tensor + lo: int + lat_eps: float + + +def _eps(norm: nn.Module) -> float: + """The epsilon of a KimiK3RMSNorm or ``torch.nn.RMSNorm`` (``eps``; for torch's None, its default: the machine + epsilon of the weight's dtype) or a stock RMSNorm (``variance_epsilon``).""" + eps = norm.eps if hasattr(norm, "eps") else norm.variance_epsilon + return float(torch.finfo(norm.weight.dtype).eps if eps is None else eps) + + +def _res_args(res_proj: nn.Module, res_norm: nn.Module, out_norm: nn.Module) -> tuple: + """The residual update's weights and epsilons in the entries' order.""" + return ( + res_proj.weight.reshape(-1), + res_norm.weight, + out_norm.weight, + _eps(res_norm), + _eps(out_norm), + ) + + +@dataclass(eq=False) +class K3DecodeComm: + """The decode path's collective state for the TP group: the `MnnvlWorkspace` and the `K3SandwichWorkspace` the + post-attention steps run on. Built by `create`; owned by the target and shared by every layer.""" + + mnnvl: MnnvlWorkspace + sandwich: K3SandwichWorkspace + + @classmethod + def create(cls, mapping, oproj_weight: Optional[torch.Tensor] = None) -> "K3DecodeComm": + """The state for ``mapping``'s TP group. Collective: every rank of the group calls it at the same point, + eagerly, before any CUDA-graph capture; each workspace fails on every rank or on none. + + ``oproj_weight``: an o_proj weight ``takes_oproj`` holds for. The sandwich kernel compiles here for its + shape, with one call on a zero row of a zero weight, so no capture compiles it; the call advances the + sandwich workspace on every rank alike.""" + state = cls( + MnnvlWorkspace.create(mapping, MNNVL_BUFFER_BYTES), + K3SandwichWorkspace.create(mapping), + ) + if oproj_weight is not None: + weight = torch.zeros_like(oproj_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_oproj( + weight.new_zeros(1, weight.shape[1]), + weight, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + state.sandwich, + ) + torch.cuda.synchronize(weight.device) + return state + + def compile_tail(self, latent_size: int, act_size: int, tail_weight: torch.Tensor) -> None: + """Compile the sandwich tail kernel for a MoE tail of ``latent_size`` latent and ``act_size`` shared columns + and ``tail_weight``'s shape, with one call on a zero row of a zero weight, before any capture. Collective: every + rank of the group makes the call; it advances the sandwich workspace on every rank alike.""" + weight = torch.zeros_like(tail_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_tail( + weight.new_zeros(1, latent_size), + weight.new_zeros(1, act_size), + weight, + 0, + 1e-6, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + self.sandwich, + ) + torch.cuda.synchronize(weight.device) + + def takes_post_attention(self, hidden_states: torch.Tensor, step) -> bool: + """Whether one of this state's collectives runs the post-attention step of a layer whose attention input is + ``hidden_states``: at most `AR_ATTN_RES_MAX_TOKENS` bf16 rows of a hidden size the MNNVL entry takes, on any + step but a wide decode step (``step.wide``), whose post-attention update stays the all-reduce and the fused + add + attn_res + RMSNorm.""" + if step is not None and step.wide: + return False + if hidden_states.dim() != 2 or hidden_states.dtype != torch.bfloat16: + return False + rows, hidden = hidden_states.shape + return ( + hidden % 1024 == 0 + and hidden <= 8192 + and 0 < rows <= min(AR_ATTN_RES_MAX_TOKENS, self.mnnvl.max_one_shot_tokens(hidden)) + ) + + @staticmethod + def takes_oproj(o_proj: nn.Module, max_snapshots: int) -> bool: + """Whether ``sandwich_oproj`` takes the layer of output projection ``o_proj`` on a step of at most + `SANDWICH_MAX_TOKENS` tokens, decided once the weights are final: a bias-free, contiguous bf16 weight of the + kernel's shape (the TP16 per-rank [7168, 768]) and a snapshot bank of at most ``max_snapshots`` rows the + kernel's candidate count holds.""" + weight = getattr(o_proj, "weight", None) + return ( + getattr(o_proj, "bias", None) is None + and isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.is_cuda + and _sandwich_op.supports(weight.new_empty((1, weight.shape[1])), weight, max_snapshots) + ) + + @staticmethod + def takes_tail( + latent: torch.Tensor, act: torch.Tensor, tail_weight: torch.Tensor, max_snapshots: int + ) -> bool: + """Whether ``sandwich_tail`` takes a MoE tail of these tensors' shapes (rows of ``latent`` and ``act``, at most + `SANDWICH_MAX_TOKENS`; ``tail_weight``) with a snapshot bank of at most ``max_snapshots`` rows: the TP16 + per-rank shapes.""" + return _sandwich_op.supports_tail(latent, act, tail_weight, max_snapshots) + + def allreduce_attn_res( + self, + partial: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_allreduce_attn_res`` call: ``updated = prefix_sum + + allreduce(partial)`` (the sum alone without ``prefix_sum``), ``normed = out_norm(attn_res(block_residual..., + updated))``, the attention residual selected with ``res_proj`` and ``res_norm``. ``block_residual`` holds the + valid snapshots, ``[S, T, H]``.""" + return mnnvl_allreduce_attn_res( + partial.contiguous(), + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.mnnvl, + ) + + def sandwich_tail( + self, + pending: PendingTail, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + updated_out: Optional[torch.Tensor] = None, + tap: Optional[torch.Tensor] = None, + tap_updated: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of a MoE layer's row-parallel tail ``pending`` + (``[RMSNorm(latent)[:, lo:lo + 224] | act] @ weight.T``), in one ``comm/k3_sandwich_tail`` call. + ``updated_out``: a bf16 ``[T, H]`` tensor the call stores ``updated`` into (the consumer's snapshot bank row), + returned as ``updated``. ``tap``: a bf16 ``[T, H]`` view (a speculative capture slot) the call also stores + the pre-norm attention-residual mixture into, or ``updated`` with ``tap_updated``.""" + return k3_sandwich_tail( + pending.latent, + pending.act, + pending.weight, + pending.lo, + pending.lat_eps, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + tap=tap, + tap_updated=tap_updated, + updated_out=updated_out, + ) + + def sandwich_oproj( + self, + core: torch.Tensor, + o_weight: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of ``core @ o_weight.T``, in one ``comm/k3_sandwich_oproj`` + call: ``core`` is this rank's gated o_proj input ``[T, 768]``, ``o_weight`` its o_proj slice. Bit for bit + o_proj followed by ``allreduce_attn_res``.""" + return k3_sandwich_oproj( + core.contiguous(), + o_weight, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + ) + + def compile_plain(self, weight: torch.Tensor, swiglu: bool = False) -> bool: + """Compile the plain sandwich (``sandwich_plain``) for a projection of ``weight``'s shape (with ``swiglu``, a + down projection after the SiLU-and-mul) with one call on a zero row of a zero weight, before any capture; + False, with nothing launched, where the kernel does not take that shape. Collective: every rank of the group + makes the call; it advances the sandwich workspace on every rank alike.""" + zero = torch.zeros_like(weight) + hidden, k_in = zero.shape + x = zero.new_zeros(1, 2 * k_in if swiglu else k_in) + residual = zero.new_zeros(1, hidden) + ones = zero.new_ones(hidden) + if not _sandwich_op.supports_plain(x, zero, residual, ones, swiglu): + return False + k3_sandwich_plain(x, zero, residual, ones, 1e-6, self.sandwich, swiglu=swiglu) + torch.cuda.synchronize(zero.device) + return True + + @staticmethod + def takes_plain( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> bool: + """Whether ``sandwich_plain`` takes the call: at most `SANDWICH_MAX_TOKENS` contiguous bf16 rows ``x`` of a + row-parallel slice ``weight`` [7168, K] (K a multiple of 128 up to 896; with ``swiglu`` K 896 and ``x`` the + ``[gate | up]`` rows, 2 K wide), ``residual`` [rows, 7168], and ``norm``'s bf16 [7168] weight.""" + return _sandwich_op.supports_plain(x, weight, residual, norm.weight, swiglu) + + def sandwich_plain( + self, + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/k3_sandwich_plain`` call: ``updated = residual + allreduce(x @ + weight.T)`` (with ``swiglu``, of ``silu_and_mul(x) @ weight.T``), ``normed = norm(updated)`` for a plain + RMSNorm ``norm``. Bit for bit, by the kernel's statement, the projection on ``k3_ctm_gemv`` at split 1 (with + ``swiglu``, ``k3_ctm_gemv_swiglu`` at split 2) followed by ``allreduce_norm``'s one-shot call.""" + return k3_sandwich_plain( + x.contiguous(), + weight, + residual.contiguous(), + norm.weight, + _eps(norm), + self.sandwich, + swiglu=swiglu, + ) + + def takes_allreduce_norm(self, rows: int, hidden: int) -> bool: + """Whether ``allreduce_norm`` takes ``rows`` bf16 rows of ``hidden`` columns: the MNNVL workspace holds the + call at the decode path's one-shot ceiling (`DECODE_AR_ONE_SHOT_MAX_BYTES`, two-shot above it).""" + if rows <= 0 or hidden <= 0 or hidden % 8: + return False + world, buffer_bytes = self.mnnvl.world_size, self.mnnvl.buffer_bytes + need = required_buffer_bytes( + rows, hidden, world, torch.bfloat16, DECODE_AR_ONE_SHOT_MAX_BYTES + ) + two_shot = rows * hidden * world * 2 > DECODE_AR_ONE_SHOT_MAX_BYTES + return need <= buffer_bytes and not (two_shot and buffer_bytes % 32) + + def allreduce_norm( + self, partial: torch.Tensor, residual: torch.Tensor, norm: nn.Module + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_fusion_allreduce`` call: ``updated = residual + + allreduce(partial)``, ``normed = norm(updated)`` for a plain RMSNorm ``norm``, sent one-shot up to + `DECODE_AR_ONE_SHOT_MAX_BYTES`. ``partial`` is this rank's unreduced ``[rows, hidden]`` bf16 output of a + row-parallel projection, of a shape ``takes_allreduce_norm`` holds for.""" + return mnnvl_fusion_allreduce( + partial.contiguous(), + self.mnnvl, + DECODE_AR_ONE_SHOT_MAX_BYTES, + residual.contiguous(), + norm.weight, + _eps(norm), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_gemv.py new file mode 100644 index 000000000000..301b709643b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_gemv.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. + +* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest + at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, + `gemm/k3_ctm_gemv_wide`, `gemm/k3_ctm_gemv_long`, `gemm/k3_ctm_gemv` or `gemm/k3_ctm_gemv_swiglu`, and, where the + site lists it, more rows (up to `WIDE_ROWS`) on `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused + projections (MLA's [W_a; W_g] with the gate rows through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention + output projection, the built-in MLA path's q_a / kv_a, q_b and gate projections, the MoE decode path's projections + (`decode_moe.py`), two of them with fp32 outputs, and the DSpark drafter's q / k / v, output, gate / up and down + projections (`K3DSparkDrafter`). +* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on + `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the + stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the + drafter's logits on the same head take it too. +* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual + bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. +* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: + gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then + runs the down projection's all-reduce. + +Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. + +A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. +`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The +embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call +whose kernel has not run eagerly is refused. + +The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one +stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, Set + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( + k3_ctm_gemv_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( + K3HeadGemvWorkspace, + k3_head_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm + +# The kernels' support predicates: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._utils import mpi_disabled + +# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 +# requests of 8 tokens). +MAX_ROWS = 8 +WIDE_ROWS = 64 + + +@dataclass(frozen=True) +class Site: + """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows + ("decode", "wide", "long", "ctm" for k3_ctm_gemv, or "swiglu" for k3_ctm_gemv_swiglu, whose rows are the + ``[gate | up]`` input of the SiLU-and-mul, 2 ``k`` wide), and k3_ctm_gemv_wide at MAX_ROWS+1..``wide_rows`` rows + where ``wide``. Output columns from ``sig_col0`` on are stored through a sigmoid; ``out_fp32``: an fp32 output + (k3_ctm_gemv_wide only). ``split`` / ``ring`` / ``push``: the CTM kernels' CTAs per 128-row weight tile, + k3_ctm_gemv_long's weight-ring stages, and whether the partial sums are pushed to each row's owner.""" + + n: int + k: int + small: str + wide: bool = False + sig_col0: int = -1 + split: int = 0 + ring: int = 0 + push: bool = False + out_fp32: bool = False + wide_rows: int = WIDE_ROWS + + +SITES: Dict[str, Site] = { + # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), + # the gate rows through a sigmoid. + "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), + # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). + "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), + # The attention output projection (row parallel; 6 heads x 128 in). + "o_proj": Site(7168, 768, "decode", wide=True), + # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. + "kv_a": Site(2112, 7168, "decode"), + "q_b": Site(1152, 1536, "wide"), + "g_proj": Site(768, 7168, "wide"), + # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. + "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), + "dense_down": Site(7168, 2112, "long", split=2, ring=6), + # The MoE decode path (decode_moe.py): this rank's head slice [latent down 224 | router 56] with an fp32 output, + # the shared experts' gate_up, the row-parallel tail [latent up 224 | padding 32 | shared down 384] and the + # replicated tail's latent up projection with an fp32 output. The wide kernel takes the head slice and the tail up + # to 32 rows, where it is faster than the stock GEMM. + "moe_head": Site(280, 7168, "wide", wide=True, out_fp32=True, wide_rows=32), + "moe_shared_gate_up": Site(768, 7168, "wide", wide=True), + "moe_tail": Site(7168, 640, "wide", wide=True, wide_rows=32), + "moe_up": Site(7168, 3584, "wide", out_fp32=True), + # The DSpark drafter's projections (K3DSparkDrafter; 6 query heads, 1 KV head of 64, MLP width 896): the fused + # q / k / v, the output projection (row parallel), gate_up [gate 896 | up 896], and the down projection with the + # SiLU-and-mul (row parallel). + "drafter_qkv": Site(512, 7168, "long", split=8, ring=6, push=True), + "drafter_o": Site(7168, 384, "ctm", split=1, push=True), + "drafter_gate_up": Site(1792, 7168, "long", split=8, ring=6, push=True), + "drafter_down": Site(7168, 896, "swiglu", split=2, push=True), +} + + +def _width(spec: Site) -> int: + """The row width a site's kernel reads: ``k``, or the ``[gate | up]`` input of a SiLU-and-mul site, 2 ``k``.""" + return 2 * spec.k if spec.small == "swiglu" else spec.k + + +def _wide_tile(rows: int) -> int: + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel + + return k3_ctm_gemv_kernel.wide_tile(rows) + + +def _run( + spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor +) -> Optional[torch.Tensor]: + """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" + if kernel == "decode": + if spec.sig_col0 >= 0 or spec.out_fp32 or not _decode_op.supports(x2d, weight): + return None + return k3_decode_gemv(x2d, weight) + if kernel == "wide": + if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, spec.out_fp32): + return None + return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0, out_fp32=spec.out_fp32) + if spec.out_fp32: + return None + if kernel == "ctm": + if spec.sig_col0 >= 0 or not _ctm_op.supports(x2d, weight, spec.split): + return None + return k3_ctm_gemv(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + if kernel == "swiglu": + if spec.sig_col0 >= 0 or not _ctm_op.supports_swiglu(x2d, weight, spec.split): + return None + return k3_ctm_gemv_swiglu(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. + sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count + if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( + x2d, weight, spec.split, spec.ring + ): + return None + return k3_ctm_gemv_long( + x2d, + weight, + sig_col0=spec.sig_col0, + split=spec.split, + ring=spec.ring, + trigger_early=True, + push=spec.push, + ) + + +def _capturing() -> bool: + return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() + + +def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: + """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" + if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: + return x2d.clone(memory_format=torch.contiguous_format) + return x2d + + +def _head_takes_module(lm_head: nn.Module) -> bool: + """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along + the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" + weight = getattr(lm_head, "weight", None) + mapping = getattr(lm_head, "mapping", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and weight.is_contiguous() + and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" + and getattr(lm_head, "gather_output", False) + and getattr(lm_head, "gather_output_sizes", None) is None + and getattr(lm_head, "padding_size", None) == 0 + and getattr(lm_head, "bias", None) is None + and not getattr(lm_head, "has_any_quant", True) + and mapping is not None + and not mapping.enable_attention_dp + ) + + +def _plain_rmsnorm(norm: nn.Module) -> bool: + """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for + bit. An unknown module fails closed.""" + weight = getattr(norm, "weight", None) + return ( + IS_FLASHINFER_AVAILABLE + and isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "use_cuda_tile", True) + and not getattr(norm, "return_hp_output", True) + and getattr(norm, "nvfp4_scale", None) is None + and hasattr(norm, "variance_epsilon") + ) + + +class K3DecodeGemvs: + """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran + eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are + final; owned by the target.""" + + def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: + self.head_workspace = head_workspace + self._ran: Set[tuple] = set() + + @classmethod + def create( + cls, + lm_head: Optional[nn.Module] = None, + sites: Iterable[str] = tuple(SITES), + device: Optional[torch.device] = None, + ) -> "K3DecodeGemvs": + """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: + it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on + zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or + head whose kernel does not take its shape keeps the generic path.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" + ) + head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + if device is None: + device = ( + head_weight.device + if isinstance(head_weight, torch.Tensor) + else torch.device("cuda", torch.cuda.current_device()) + ) + state = cls() + for site in sites: + spec = SITES[site] + weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) + for rows in (1, 16, 32, 64) if spec.wide else (1,): + if rows > spec.wide_rows: + continue + state._project(site, weight.new_zeros(rows, _width(spec)), weight, warm=True) + del weight + if "dense_gate_up" in sites: + gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) + if _ctm_op.supports_situ_mul(gu): + for linear_beta in (None, 1.0): + k3_situ_mul(gu, 1.0, linear_beta) + state._ran.add(("situ_mul", linear_beta is not None)) + if lm_head is not None and _head_takes_module(lm_head): + x = head_weight.new_zeros(1, head_weight.shape[1]) + if _head_op.supports(x, head_weight): + workspace = K3HeadGemvWorkspace.create( + head_weight.shape[0], head_weight.shape[1], head_weight.device + ) + k3_head_gemv(x, head_weight, workspace) + state.head_workspace = workspace + state._ran.add(("lm_head",)) + torch.cuda.synchronize(device) + return state + + def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: + """``x @ weight.T`` (``[..., N]``: bf16, the site's sigmoid columns through the sigmoid; fp32 at an + ``out_fp32`` site; on a SiLU-and-mul site, ``silu(gate) * up`` of ``x = [gate | up]`` first) for ``site``'s + weight on its decode kernel, or None where none takes the call: more rows than the site's kernels take, + another shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its + module.""" + return self._project(site, x, weight, warm=False) + + def _project( + self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool + ) -> Optional[torch.Tensor]: + spec = SITES[site] + width = _width(spec) + if ( + weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (spec.n, spec.k) + or not weight.is_contiguous() + or x.dtype != torch.bfloat16 + or x.dim() < 1 + or x.shape[-1] != width + ): + return None + rows = x.numel() // width + if 0 < rows <= MAX_ROWS: + kernel = spec.small + elif spec.wide and MAX_ROWS < rows <= spec.wide_rows: + kernel = "wide" + else: + return None + key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) + capturing = _capturing() + if capturing and not warm and key not in self._ran: + return None + y = _run(spec, kernel, _dense_rows(x.reshape(rows, width)), weight) + if y is None: + return None + if not capturing: + self._ran.add(key) + return y.view(*x.shape[:-1], spec.n) + + def lm_head_logits( + self, rows: torch.Tensor, lm_head: nn.Module, gather: bool = True + ) -> Optional[torch.Tensor]: + """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on + ``gemm/k3_head_gemv`` (``gather`` False: the shard alone, ``[M, vocab / TP]``); None where it does not take + the call (more than `MAX_ROWS` rows, another head, ...).""" + workspace = self.head_workspace + if workspace is None or not _head_takes_module(lm_head): + return None + weight = lm_head.weight + if ( + tuple(weight.shape) != (workspace.n_out, workspace.k_in) + or weight.device != workspace.partials.device + or rows.dim() != 2 + or rows.dtype != torch.bfloat16 + or not 0 < rows.shape[0] <= MAX_ROWS + or rows.shape[1] != workspace.k_in + ): + return None + if _capturing() and ("lm_head",) not in self._ran: + return None + group = lm_head.mapping.tp_group + if gather and len(group) > 1 and mpi_disabled(): + return None + x = _dense_rows(rows) + if not _head_op.supports(x, weight): + return None + local = k3_head_gemv(x, weight, workspace) + if not gather or len(group) == 1: + return local + gathered = allgather(local, None, group) + return torch.cat(torch.split(gathered, rows.shape[0], dim=0), dim=-1) + + def dense_mlp( + self, + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + beta: float, + linear_beta: Optional[float], + ) -> Optional[torch.Tensor]: + """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on + ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on + ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then + runs the module.""" + gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] + rows = x.shape[0] if x.dim() == 2 else 0 + if ( + not 0 < rows <= MAX_ROWS + or linear_beta == 0.0 + or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) + or tuple(down_weight.shape) != (down.n, down.k) + or gate_up.n != 2 * down.k + ): + return None + if ( + _capturing() + and not { + ("dense_gate_up", "long", 0), + ("dense_down", "long", 0), + ("situ_mul", linear_beta is not None), + } + <= self._ran + ): + return None + gu = self.project("dense_gate_up", x, gate_up_weight) + if gu is None or not _ctm_op.supports_situ_mul(gu): + return None + return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) + + def embed_norm( + self, + input_ids: torch.Tensor, + table: torch.Tensor, + norm: nn.Module, + bank: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one + ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not + apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose + kernel has not run eagerly).""" + if not _plain_rmsnorm(norm): + return None + ids = input_ids.reshape(-1) + if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): + return None + key = ("embed_norm", ids.numel(), ids.dtype) + capturing = _capturing() + if capturing and key not in self._ran: + return None + normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) + if not capturing: + self._ran.add(key) + return normed + + +class _K3Head: + """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" + + __slots__ = ("_gemvs", "_lm_head") + + def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: + self._gemvs = gemvs + self._lm_head = lm_head + + def __call__(self, rows: torch.Tensor) -> torch.Tensor: + logits = self._gemvs.lm_head_logits(rows, self._lm_head) + return self._lm_head(rows) if logits is None else logits + + +class K3LogitsProcessor(nn.Module): + """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. + + It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock + processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. + """ + + def __init__(self, stock: nn.Module) -> None: + super().__init__() + self.stock = stock + self.gemvs: Optional[K3DecodeGemvs] = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) + return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) + + def lm_head_shard(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: + """This rank's bf16 vocabulary shard of ``lm_head(rows)`` on ``gemm/k3_head_gemv``, without the gather, or + None where ``gemvs`` does not take the rows. A speculative worker that keeps its logits vocabulary-sharded + computes them here, so they hold the values of the logits this processor gathers.""" + if self.gemvs is None: + return None + return self.gemvs.lm_head_logits(rows, lm_head, gather=False) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py new file mode 100644 index 000000000000..e39c3ddeb0e7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/decode_moe.py @@ -0,0 +1,495 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# >>> route B: this target's MoE engines (every expert on each rank) +"""The decode path's latent MoE on the catalog's Kimi K3 MoE entries, with every expert on each rank (moe TP16 x EP1). + +**At most 8 tokens** (`MAX_TOKENS`), a MoE layer runs: + +* `moe/k3_moe_front`, one kernel: this rank's slice of the MoE head (its latent-down rows and its router rows) as one + GEMV, the slices' all-gather over the TP group's `K3MoeHeadWorkspace`, the top-16 routing, the MXFP8 latent, and + the shared experts' gate_up + SiTU. The routing is the noaux_tc arithmetic of `moe/kimi_k3_noaux_tc_mxfp8_quant`, + as the generic path's (`KimiK3MoeRoutingMethod`), on the router logits of the front's head GEMV; +* the routed experts of all 896 experts over this rank's intermediate slice: `moe/k3_moe_m1` at one token, + `moe/k3_moe_m2` at two, `moe/k3_moe` (on a `K3MoeState`) at three to eight; +* the latent all-reduce. On a pushing step (`DecodeStep.latent_push`: a pure decode step captured into a CUDA graph + whose attention layers all run the decode kernels) the engine runs its push form, which stores this rank's partial + into every rank's `K3LatentExchange`, and `comm/k3_latent_reduce` sums the partials in the MNNVL one-shot's order, + so with the same bits, while the experts' grid completes. On every other step it is the routed experts' all-reduce + (one-shot, as `decode_comm.use_decode_one_shot` sets); +* the tail. The latent norm's weight is folded into the latent up projection at load, so + `[RMSNorm(latent) slice | shared activation] @ [latent up columns | shared down]` is this rank's share of the MoE + output. Where the next layer's pre-attention step (or the final norm) reduces it, the layer hands it on unreduced + as a `PendingTail`, which the consumer runs with its all-reduce and residual update as one `comm/k3_sandwich_tail` + kernel (`decode_comm.py`). Elsewhere the replicated tail runs: the latent RMS applied to the fp32 output of one + GEMV with the folded latent up weight, plus the shared experts' down projection and its all-reduce. + +**More than 8 tokens** run the generic path: `k3_moe`'s wide build does not fit 896 local experts, so the wide step +below (`_wide`, `tp16_moetp4ep4`'s) never runs here. The engines compile the SiTU caps 4 and 25 in +(`ENGINE_SITU_CAPS`): a checkpoint with other caps keeps the generic path (`layout_gaps`). + +The GEMVs run on the decode GEMV sites of `decode_gemv.py` where they take the call, else on the stock GEMM ops. + +`K3DecodeMoe` holds what every MoE layer shares: the head workspace and the latent exchange (both collective over the +TP group), the engines' workspaces (`K3MoeM1State`, `K3MoeM2State` and the `k3_moe` build's scratch), and the TP +group's MNNVL workspace (`decode_comm.K3DecodeComm`'s). `K3DecodeMoeLayer` holds one layer's decode weights and its +engine handles. The target builds both in `post_load_weights`, before any CUDA-graph capture, and runs every kernel +once there so none compiles under a capture. + +Every MoE layer's push and reduce go to the one exchange, in the stream's order: each push is followed by exactly one +reduce of the same token count before the next push, on every rank in the same order (the exchange's call-order +invariant, `comm/k3_latent_reduce`). Every kernel between a reduce and the next push waits for its predecessor (or +launches without programmatic dependent launch), which the decode kernels a pushing step runs do; pushes run only in +CUDA-graph replays, so no other step's kernels sit between two pushes. +""" +# <<< route B + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_latent_reduce import ( + K3LatentExchange, + k3_latent_reduce, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allgather_split import ( + mnnvl_allgather_split, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe import ( + K3MoeHeadWorkspace, + K3MoeLayer, + K3MoeState, + K3MoeWideState, + is_supported, + k3_moe, + k3_moe_push, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe_front import ( + front_weight, + k3_moe_front, + weight_supported, +) + +# >>> route B: the one- and two-token engines +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe_m1 import ( + K3MoeM1Layer, + K3MoeM1State, + k3_moe_m1, + k3_moe_m1_push, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe_m2 import ( + K3MoeM2Layer, + K3MoeM2State, + k3_moe_m2, + k3_moe_m2_push, +) + +# <<< route B +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_route_quant import k3_route_quant +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel + +from .decode_comm import K3DecodeComm, PendingTail, wide_all_reduce + +# The most tokens of the front and the small k3_moe build (one token tile), and of the wide build. +MAX_TOKENS = 8 +WIDE_MAX_TOKENS = 64 + +# The tail weight's latent columns are zero-padded to whole 128-column k-tiles of the tail kernels +# (comm/k3_sandwich_tail takes the TP16 tail weight [7168, 256 + 384]). +TAIL_K_TILE = 256 + +# >>> route B: the routed experts' SiTU caps (activation_situ_beta, activation_situ_linear_beta) that k3_moe_m1, +# k3_moe_m2 and k3_moe compile in (their kernels' SITU_GATE_CAP and SITU_LINEAR_CAP) +ENGINE_SITU_CAPS = (4.0, 25.0) +# <<< route B + + +@dataclass(eq=False) +class K3DecodeMoe: + # >>> route B: the engines' workspaces; no wide build + """What every MoE layer's decode path shares on one device: the TP group's ``K3MoeHeadWorkspace`` (the front's + all-gather), the engines' workspaces (``k3_moe_m1``, ``k3_moe_m2`` and the ``k3_moe`` build for up to 8 tokens; + the layers run one at a time on one stream), the TP group's ``MnnvlWorkspace``, and the TP group's + ``K3LatentExchange`` for a pushing step's latent all-reduce (None: every step keeps the routed experts' + all-reduce). ``wide`` is None: ``k3_moe``'s wide build does not fit 896 local experts. Built by `create`.""" + + # <<< route B + + head: K3MoeHeadWorkspace + small: K3MoeState + wide: K3MoeWideState + mnnvl: MnnvlWorkspace + exchange: Optional[K3LatentExchange] = None + # >>> route B: the one- and two-token engines' workspaces (``create`` builds both) + m1: Optional[K3MoeM1State] = None + m2: Optional[K3MoeM2State] = None + # <<< route B + + # >>> route B: every expert on each rank: the engines with their push builds for the exchange and the k3_moe build + # for up to 8 tokens; no wide build + @classmethod + def create( + cls, + mapping, + device, + i_tp: int, + num_local: int, + mnnvl: MnnvlWorkspace, + push: bool = True, + *, + i_logical: int, + ) -> "K3DecodeMoe": + """The state for ``mapping``'s TP group on ``device``, for experts of ``i_tp`` intermediate columns per rank + (the loader's zero-padded width; ``i_logical`` of them hold the checkpoint's slice, which k3_moe_m1 and + k3_moe_m2 stream), ``num_local`` of them on this rank. Collective (the head workspace and the latent + exchange): every rank of the group calls it at the same point, eagerly, before any CUDA-graph capture. + ``push``: build the latent exchange (where it takes the TP size: 4, 8 or 16 ranks) and the engines' push builds + for it; without it every step keeps the routed experts' all-reduce.""" + head = K3MoeHeadWorkspace.create(mapping) + exchange = None + if push: + try: + exchange = K3LatentExchange.create(mapping) + # Raised before any collective step, on every rank of the group alike: a TP size the exchange does not take. + except ValueError: + exchange = None + builds = () if exchange is None else ((mapping.tp_size, 1),) + return cls( + head, + K3MoeState(device, i_tp, num_local), + None, + mnnvl, + exchange, + K3MoeM1State.create(device, i_logical, i_tp, num_local, push=builds), + K3MoeM2State.create(device, i_logical, i_tp, num_local, push=builds), + ) + + # <<< route B + + +def _experts(moe: nn.Module) -> tuple: + """The routed experts' TRTLLM-Gen W4A8_MXFP4_MXFP8 buffers ``k3_moe`` reads in place.""" + backend = moe.routed_experts.backend + return ( + backend.w3_w1_weight, + backend.w3_w1_weight_scale, + backend.w2_weight, + backend.w2_weight_scale, + ) + + +def layout_gaps(moe: nn.Module, tp_size: int, max_snapshots: int) -> list: + """Why the decode path does not take MoE layer ``moe`` (empty when it takes it), read once the routed experts' + buffers exist: the checks in order, up to the first that fails. ``max_snapshots``: the model's snapshot bank + rows.""" + gate, backend = moe.gate, moe.routed_experts.backend + shared = moe.shared_experts + gate_up, down = shared.gate_up_proj.weight, shared.down_proj.weight + if not moe._reduce_routed_output: + return ["a routed output the model does not reduce"] + # >>> route B: k3_moe_m1, k3_moe_m2 and k3_moe compile the SiTU caps in + if tuple(moe._situ_betas) != ENGINE_SITU_CAPS: + return [f"SiTU caps {tuple(moe._situ_betas)}; the MoE engines compile {ENGINE_SITU_CAPS}"] + # <<< route B + if (moe.num_experts, moe.top_k, moe.moe_hidden_size, moe.hidden_size) != (896, 16, 3584, 7168): + return [ + f"experts / top-k / latent / hidden {(moe.num_experts, moe.top_k, moe.moe_hidden_size)}" + ] + if (gate.num_expert_group, gate.topk_group) != (1, 1): + return ["grouped routing"] + if moe.moe_hidden_size % (8 * tp_size) or moe.num_experts % (4 * tp_size): + return [f"a head slice of TP {tp_size}"] + latent = (gate.weight, moe.routed_expert_down_proj.weight, moe.routed_expert_up_proj.weight) + if not isinstance(moe.routed_expert_up_proj, nn.Linear) or any( + w.dtype != torch.bfloat16 for w in latent + ): + return ["latent projections or a router other than bf16"] + if ( + gate_up.dtype != torch.bfloat16 + or down.dtype != torch.bfloat16 + or shared.gate_up_proj.bias is not None + or shared.down_proj.bias is not None + or down.shape[0] != moe.hidden_size + ): + return ["a shared expert other than bf16 and unbiased"] + names = ( + "w3_w1_weight", + "w3_w1_weight_scale", + "w2_weight", + "w2_weight_scale", + "expert_size_per_partition", + ) + if ( + not all(hasattr(backend, name) for name in names) + or not is_supported(*_experts(moe), backend.expert_size_per_partition)[0] + ): + return ["routed-expert buffers k3_moe does not read"] + shared_cols, width = gate_up.shape[0] // 2, moe.moe_hidden_size // tp_size + if not weight_supported(tp_size, shared_cols, gate_up.shape[1], gate_up.device): + return ["a MoE front of this TP size and shared width"] + tail_cols = width + (-width % TAIL_K_TILE) + shared_cols + probe = gate_up.new_empty + if not K3DecodeComm.takes_tail( + probe(1, moe.moe_hidden_size), + probe(1, shared_cols), + probe(moe.hidden_size, tail_cols), + max_snapshots, + ): + return ["a row-parallel tail the sandwich tail kernel does not take"] + return [] + + +def fold_latent_norm(moe: nn.Module) -> None: + """Fold the latent RMSNorm's weight into the latent up projection's columns; the norm keeps a weight of ones, so + every step computes the same function and the tails may normalize before slicing.""" + up, norm = moe.routed_expert_up_proj, moe.routed_expert_norm + with torch.no_grad(): + up.weight.mul_(norm.weight.to(up.weight.dtype)[None, :]) + norm.weight.fill_(1) + + +@dataclass(eq=False) +class K3DecodeMoeLayer: + """One MoE layer's decode path: the front weight (this rank's head slice padded to whole tiles, then the shared + gate_up re-ordered), the head slice (a view of it), the row-parallel tail weight + ``[latent up columns lo:lo+width | padding | shared down]``, the routing bias, and its ``k3_moe`` handles on the + shared builds. Built by `create` once the weights are final (and the latent norm folded).""" + + state: K3DecodeMoe + front_weight: torch.Tensor + head_weight: torch.Tensor + tail_weight: torch.Tensor + tail_pad: Optional[torch.Tensor] + # The gate's routing bias detached (a view of the parameter): an op given a tensor that requires grad while + # autograd records, as in post_load_weights, returns outputs that require grad, which k3_moe cannot take. + bias: torch.Tensor + lo: int + width: int + shared_cols: int + small: K3MoeLayer + wide: K3MoeLayer + # >>> route B: the one- and two-token engines' handles + m1: K3MoeM1Layer + m2: K3MoeM2Layer + # <<< route B + + @classmethod + def create( + cls, moe: nn.Module, state: K3DecodeMoe, tp_rank: int, tp_size: int + ) -> "K3DecodeMoeLayer": + """Build MoE layer ``moe``'s decode weights and its handles on ``state``'s builds (``layout_gaps`` empty).""" + width = moe.moe_hidden_size // tp_size + experts = moe.num_experts // tp_size + gate_up = moe.shared_experts.gate_up_proj.weight + shared_down = moe.shared_experts.down_proj.weight + up = moe.routed_expert_up_proj.weight + lo = tp_rank * width + pad = -width % TAIL_K_TILE + with torch.no_grad(): + head = torch.cat( + [ + moe.routed_expert_down_proj.weight[lo : lo + width], + moe.gate.weight[tp_rank * experts : (tp_rank + 1) * experts], + ] + ) + front = front_weight(head, gate_up) + parts = [up[:, lo : lo + width]] + if pad: + parts.append(up.new_zeros(up.shape[0], pad)) + parts.append(shared_down) + tail = torch.cat(parts, dim=1).contiguous() + weights = _experts(moe) + return cls( + state=state, + front_weight=front, + head_weight=front[: head.shape[0]], + tail_weight=tail, + tail_pad=up.new_zeros(WIDE_MAX_TOKENS, pad) if pad else None, + bias=moe.gate.e_score_correction_bias.detach(), + lo=lo, + width=width, + shared_cols=gate_up.shape[0] // 2, + # >>> route B: the engines' handles; no wide build + small=state.small.layer(*weights), + wide=None, + m1=state.m1.layer(*weights), + m2=state.m2.layer(*weights), + # <<< route B + ) + + def warm_up(self, moe: nn.Module) -> None: + # >>> route B: the engines; no wide build + """One call of every kernel of the decode path on zero inputs, so none compiles under a capture: the front + (collective: every rank makes the same call), each engine (``k3_moe_m1``, ``k3_moe_m2``, ``k3_moe``) and, with + the latent exchange, its push build, each push followed by the reduce (collective). Once per model: every + layer's calls compile the same builds.""" + # <<< route B + device = self.front_weight.device + x = torch.zeros(1, moe.hidden_size, dtype=torch.bfloat16, device=device) + ids, weights, x_fp8, x_sf, _ = self._front(moe, x) + offset = moe.routed_experts.backend.slot_start + # >>> route B: every engine returning and, with the latent exchange, pushing (k3_moe's builds compile here), + # each push followed by the reduce; no wide build + exchange = self.state.exchange + for rows in (1, 2, 3): + args = [t.expand(rows, -1).contiguous() for t in (x_fp8, x_sf, ids, weights)] + self._routed(*args, offset) + if exchange is not None: + # Two calls on one layer must not run back to back (moe/k3_moe): the push build starts once the + # call above has ended. + torch.cuda.synchronize(device) + self._routed(*args, offset, push=True) + k3_latent_reduce(rows, exchange) + # <<< route B + torch.cuda.synchronize(device) + + def _front(self, moe: nn.Module, x: torch.Tensor): + return k3_moe_front( + x.contiguous(), + self.front_weight, + self.bias, + float(moe.gate.routed_scaling_factor), + self.shared_cols, + *moe._situ_betas, + self.state.head, + ) + + def takes(self, hidden_states: torch.Tensor, step, partial_tail: bool) -> bool: + # >>> route B: no wide step + """Whether this decode path runs the layer on ``step`` (bf16 rows ``hidden_states``): any step of at most + `MAX_TOKENS` tokens, with or without ``partial_tail``.""" + # <<< route B + rows = hidden_states.shape[0] + if hidden_states.dtype != torch.bfloat16: + return False + # >>> route B: no wide step (k3_moe's wide build does not fit 896 local experts) + return 0 < rows <= MAX_TOKENS + # <<< route B + + def forward( + self, + moe: nn.Module, + hidden_states: torch.Tensor, + gemvs, + partial_tail: bool, + push: bool = False, + ) -> Union[torch.Tensor, PendingTail]: + """The MoE output of ``hidden_states`` (``takes`` holds): with ``partial_tail``, this rank's unreduced share + (a ``PendingTail`` at most `MAX_TOKENS` tokens); else the reduced output. ``gemvs``: the decode GEMVs' state, + or None. ``push`` (the step's ``DecodeStep.latent_push``): at most `MAX_TOKENS` tokens, the latent all-reduce + is the routed experts' push form plus ``comm/k3_latent_reduce`` on the state's exchange (the same bits as the + routed experts' all-reduce); a state without the exchange ignores it.""" + if hidden_states.shape[0] > MAX_TOKENS: + return self._wide(moe, hidden_states, gemvs) + ids, weights, x_fp8, x_sf, shared_act = self._front(moe, hidden_states) + offset = moe.routed_experts.backend.slot_start + exchange = self.state.exchange + # >>> route B: the engine of the token count + if push and exchange is not None: + self._routed(x_fp8, x_sf, ids, weights, offset, push=True) + latent = k3_latent_reduce(hidden_states.shape[0], exchange) + else: + latent = moe.routed_experts.all_reduce(self._routed(x_fp8, x_sf, ids, weights, offset)) + # <<< route B + if partial_tail: + return PendingTail( + latent.contiguous(), + shared_act.contiguous(), + self.tail_weight, + self.lo, + float(moe.routed_expert_norm.variance_epsilon), + ) + # The replicated tail: the latent RMS on the fp32 accumulator of the folded latent up projection, plus the + # shared experts' reduced output, rounded to bf16 once. + shared = moe.shared_experts + shared_out = shared.down_proj(shared_act, layer_idx=shared.layer_idx) + up = _gemv(gemvs, "moe_up", latent, moe.routed_expert_up_proj.weight, out_fp32=True) + scale = torch.rsqrt( + latent.float().pow(2).mean(-1, keepdim=True) + moe.routed_expert_norm.variance_epsilon + ) + return (up * scale + shared_out.float()).bfloat16() + + # >>> route B: the engine of a token count + def _routed(self, x_fp8, x_sf, ids, weights, offset: int, push: bool = False): + """This rank's routed partial of the front's outputs on ``k3_moe_m1`` at one token, ``k3_moe_m2`` at two, else + ``k3_moe`` (up to `MAX_TOKENS`); with ``push``, stored into every rank's latent exchange instead (None).""" + rows = x_fp8.shape[0] + if push: + exchange = self.state.exchange + if rows == 1: + k3_moe_m1_push(x_fp8, x_sf, ids, weights, offset, self.m1, exchange) + elif rows == 2: + k3_moe_m2_push(x_fp8, x_sf, ids, weights, offset, self.m2, exchange) + else: + k3_moe_push(x_fp8, x_sf, ids, weights, offset, self.small, exchange) + return None + if rows == 1: + return k3_moe_m1(x_fp8, x_sf, ids, weights, offset, self.m1) + if rows == 2: + return k3_moe_m2(x_fp8, x_sf, ids, weights, offset, self.m2) + return k3_moe(x_fp8, x_sf, ids, weights, offset, self.small) + + # <<< route B + def _wide(self, moe: nn.Module, hidden_states: torch.Tensor, gemvs) -> torch.Tensor: + """A wide decode step's MoE: this rank's unreduced share of the output, ``[M, hidden]`` bf16.""" + x = hidden_states.contiguous() + head = _gemv(gemvs, "moe_head", x, self.head_weight, out_fp32=True) + routed_in, router_logits = mnnvl_allgather_split(head, self.width, self.state.mnnvl) + shared = moe.shared_experts + + def _routed_partial(): + ids, weights, x_fp8, x_sf = k3_route_quant( + router_logits, self.bias, routed_in.contiguous(), + float(moe.gate.routed_scaling_factor), early_trigger=True, + ) # fmt: skip + return k3_moe( + x_fp8, x_sf, ids, weights, moe.routed_experts.backend.slot_start, self.wide + ) + + def _shared_activation(): + return shared._apply_activation( + _gemv(gemvs, "moe_shared_gate_up", x, shared.gate_up_proj.weight) + ) + + routed, shared_act = maybe_execute_in_parallel( + _routed_partial, + _shared_activation, + moe.moe_main_event, + moe.moe_shared_event, + moe.shared_expert_stream, + disable_on_compile=True, + ) + # The latent norm's weight is folded into the tail weight: normalize the whole latent row, keep this rank's + # columns. + normed = moe.routed_expert_norm(wide_all_reduce(moe.routed_experts.all_reduce, routed)) + parts = [normed[:, self.lo : self.lo + self.width]] + if self.tail_pad is not None: + parts.append(self.tail_pad[: x.shape[0]]) + parts.append(shared_act) + return _gemv(gemvs, "moe_tail", torch.cat(parts, dim=1), self.tail_weight) + + +def _gemv( + gemvs, site: str, x: torch.Tensor, weight: torch.Tensor, out_fp32: bool = False +) -> torch.Tensor: + """``x @ weight.T`` on ``site``'s decode GEMV where it takes the call (``decode_gemv.K3DecodeGemvs.project``), + else on the stock GEMM: cuBLAS through ``trtllm::dsv3_router_gemm_op`` for an fp32 output, ``F.linear`` else.""" + y = None if gemvs is None else gemvs.project(site, x, weight) + if y is not None: + return y + if out_fp32: + # The op reads the weight with a leading dimension of K: a strided weight would give wrong values. + if not weight.is_contiguous(): + raise ValueError( + f"the fp32 GEMM needs a contiguous weight, got strides {weight.stride()}" + ) + return torch.ops.trtllm.dsv3_router_gemm_op( + x.contiguous(), weight.t(), bias=None, out_dtype=torch.float32 + ) + return torch.nn.functional.linear(x, weight) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index 575c6a578a97..b53cdb0e8482 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -1,21 +1,3532 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 MXFP4, SM100, attention TP16 and MoE TP16 x EP1.""" +# >>> route B: this target's layout, and decoding without speculation +"""ModelingV2 target: Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 16 x moe_ep 1. +Kimi K3's language model: 93 layers, hidden 7168. Every fourth layer (3, 7, ..., 91) is MLA attention with 96 +query heads; the others are Kimi Delta Attention (KDA), a gated linear-attention recurrence with per-request state. +Layer 0's MLP is dense. Every other layer's MLP is a latent MoE: 896 routed experts, top-16, expert width 3584. +Attention residuals mix each sublayer's input from a bank of earlier outputs. The routed experts are MXFP4 (the +checkpoint's compressed-tensors default, read as W4A8_MXFP4_MXFP8); everything else is bf16, and so is the KV pool. + +`tp16_moetp16ep1` is `tensor_parallel_size: 16` with `moe_tensor_parallel_size: 16` and +`moe_expert_parallel_size: 1`, both set explicitly, no attention data parallelism, all 16 GPUs in one NVLink domain. +Attention is head-split (6 MLA query heads per rank), and every rank holds all 896 experts at a sixteenth of their +intermediate width (192 of 3072 values, zero-padded to 256 when they load). With the expert split left unset the +engine resolves the same sizes, but Kimi K3 then runs the experts expert-parallel over the 16 ranks, a layout no +target serves. This is Kimi K3's layout for decoding without speculation; `tp16_moetp4ep4` serves speculative +decoding. + +**This module is `tp16_moetp4ep4`'s, copied** (targets share no files) and changed only inside its marked route B +blocks; test_modeling_v2_kimi_k3_drift.py checks that the two modules match everywhere else. + +**Each step is classified once, on the host, from its shape** (`decode_step`, a `DecodeStep` or None), and the +classification decides which kernels each module runs: + +* **small**: at most 8 tokens (`DECODE_MAX_TOKENS`, one token tile), context requests included. The token-count + kernels take it: decode GEMVs, the sandwiches, the embedding and residual epilogues. +* **decode**: a pure decode step of R <= 8 generation requests and no context request; without speculation each + request has one token. The request-aware kernels take it: MLA attention and its KV store. + +Every other step (prefill, mixed steps, decode steps above those bounds) runs the **generic path**: this target's +text model (`KimiLinearModel` below: decoder layers, attention residuals, the MLA / KDA / MoE runtimes), computed as +the built-in Kimi K3 text model computes it, on stock modules and ops that have no catalog entries yet +(`UNCERTIFIED_GENERIC_CALLS` names them), except that the routed experts' top-16 is computed outside the TRTLLM-Gen +kernel (`KimiK3MoeRoutingMethod`), with the kernel's arithmetic. + +The text model hands each step's classification to its attention modules, which run a **decode step** on the K3 +decode kernels' catalog entries: + +* KDA (`K3DecodeKDA`): one token per request, the fused input projection and the plain decode in one + `ssm/k3_kda_decode_attn` launch. +* MLA (`K3DecodeMLA`): `attention/k3_mla_qkv` (the query path and the step's latent KV rows into the paged cache), + then `attention/k3_mla_attn_vb_out` (the attention, v_b and the output gate in one launch). + +The projections around them run on the decode GEMV sites of `decode_gemv.py` (the [W_a; W_g] projection, `o_proj` +on every classified step), as do the LM head, the embedding and layer 0's dense MLP. A classified step's +attention-residual epilogues (the selection and the RMSNorm after it) take the fused kernels up to one token tile. On +any step of at most 16 tokens, the post-attention all-reduce carries the residual update (`decode_comm.py`): at most 8 +tokens on an attention decode branch, o_proj, the all-reduce and the update are one `comm/k3_sandwich_oproj` kernel; +otherwise the attention's unreduced o_proj output goes through `comm/mnnvl_allreduce_attn_res`. The stock MNNVL +all-reduces send one-shot up to 4 MiB. The state those kernels share (the KDA projection's Lamport buffers, the MLA +attention workspace, the decode GEMVs' state, the TP group's MNNVL and sandwich workspaces) lives in typed objects this +target creates in `post_load_weights`, before any graph capture. + +A MoE layer on a step of at most 8 tokens runs `decode_moe.py` with every expert on each rank: `moe/k3_moe_front` +(this rank's head slice, its all-gather, the top-16 routing, the MXFP8 latent and the shared experts' gate_up + SiTU), +then the routed experts as `moe/k3_moe_m1` at one token, `moe/k3_moe_m2` at two and `moe/k3_moe` over all 896 experts +up to 8, the latent all-reduce and the row-parallel tail. On a pushing step (`DecodeStep.latent_push`) the engines +push their partials into the TP group's `K3LatentExchange` and `comm/k3_latent_reduce` sums them; other steps use the +routed experts' all-reduce, which sums in the same order. Its routing is the front's: the noaux_tc arithmetic, as on +the generic path, on the router logits of the front's head GEMV. Wider steps, and every step of a checkpoint whose +SiTU caps differ from the ones the engines compile in, run the generic path: `k3_moe`'s wide build does not fit 896 +local experts. + +**What this target asserts rather than adapts**: SM 10.0; the topology above, with the expert split set explicitly; +no speculative decoding; the MXFP4 checkpoint's quantization (W4A16_MXFP4 with no per-layer declarations, so the +routed experts run the W4A8_MXFP4_MXFP8 default and the excluded modules stay bf16); bf16 weights and a bf16 KV pool; +tokens_per_block 64 (the MLA generation kernels K3's 96 heads reach exist only at 64); the V2 hybrid KV / state +manager, which holds the KDA states, with block reuse off and fp32 recurrent states; an all-reduce strategy of AUTO or +MNNVL. The construction-time ones fail in `__init__`, the +per-engine ones on the first forward, each naming the setting. A layer the decode kernels do not take fails the +weight load. + +**Text only.** The checkpoint is the vision-language wrapper. This target builds and loads no vision tower (its +weights are a predicted non-load, `weights.py`), and a step carrying multimodal input raises. + +**No speculative decoding.** A speculative decoding config fails construction, naming `tp16_moetp4ep4`'s expert +split. The causal LM is still the stock one-engine shell the built-in model builds on, which without that config +builds no drafter and no worker. This target drops `tp16_moetp4ep4`'s KDA verify kernels and the per-token verify +states they keep; the text model's speculative hidden-state taps stay as `tp16_moetp4ep4` has them and never fire. +The DSpark drafter (`K3DSparkDrafter`) and the speculative worker (`spec_worker.py`), which `_own_spec_branch` puts in +place of the shell's stock ones, also stay as `tp16_moetp4ep4` has them and are never built here; the drafter's decode +GEMV sites are warmed with the others and never called. +""" +# <<< route B + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( + k3_drafter_attn_qknorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( + k3_mla_attn_vb_out, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( + K3MlaAttnWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import k3_mla_qkv + +# >>> route B: no KDA verify kernels, and no trtllm::kda_mtp_decode (the built-in KDA verify's) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import K3KdaBuffers +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( + k3_kda_decode_attn, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.decode_view import k3_mla_decode_view + +# <<< route B +from tensorrt_llm._torch.distributed import AllReduce from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.models.modeling_dspark import ( + GQADSparkForCausalLM, + draft_is_embedded_in_target, +) +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.moe.fused_moe import ( + ConfigurableMoE, + SiTuActivation, + TRTLLMGenFusedMoE, + create_moe, +) +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.config_utils import is_mla +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from . import decode_comm as _decode_comm +from . import decode_gemv as _decode_gemv +from . import decode_moe as _decode_moe +from . import spec_worker as _spec_worker +from . import weights as _weights + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, +# and the certification is per arch. +_SM = (10, 0) + +#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in +#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router +#: and fused-A GEMMs), then the decode kernels'. +REQUIRED_TRTLLM_OPS = ( + "attn_res_fwd", + "attn_res_rmsnorm_fwd", + "attn_res_add_rmsnorm_fwd", + "attn_res_add_rmsnorm_persistent_fwd", + "kda_prefill", + "kda_decode", + # >>> route B: no KDA verify (kda_mtp_decode, k3_kda_attn, k3_kda_verify) + "dsv3_router_gemm_op", + "dsv3_fused_a_gemm_op", + "k3_kda_decode_attn", + # <<< route B + "k3_mla_qkv", + "k3_mla_attn_vb_out", + # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). + "k3_decode_gemv", + "k3_ctm_gemv_wide", + "k3_ctm_gemv_long", + "k3_situ_mul", + "k3_head_gemv", + "k3_embed_norm", + "allgather", + # The decode path's collectives (decode_comm.py) and MoE (decode_moe.py). + "mnnvl_allreduce_attn_res", + "k3_sandwich_oproj", + "k3_sandwich_tail", + "k3_moe_front", + "k3_moe", + "k3_latent_reduce", + "k3_route_quant", + "mnnvl_allgather_split", + # >>> route B: the MoE engines of one and two tokens + "k3_moe_m1", + "k3_moe_m2", + # <<< route B + # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces + # with the residual add and RMSNorm (the split context projection's with hidden_norm). + "k3_ctm_gemv", + "k3_ctm_gemv_swiglu", + "k3_drafter_attn_qknorm", + "k3_sandwich_plain", + "mnnvl_fusion_allreduce", +) + +#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. +#: A renamed field upstream fails here, loudly, instead of reading as a default. +REQUIRED_ENGINE_FIELDS = { + "attn_metadata": ( + "num_contexts", + "num_generations", + "num_seqs", + "num_tokens", + "seq_lens", + "tokens_per_block", + "kv_cache_manager", + "mamba_metadata", + ), + "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), +} + +#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every +#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry +#: leaves the list when a catalog entry replaces it. +UNCERTIFIED_GENERIC_CALLS = ( + # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. + "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", + "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", + "tensorrt_llm._torch.models.modeling_utils.DecoderModel", + # The text model's stock modules. + "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", + "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", + # >>> route B: no trtllm::kda_mtp_decode registration (the built-in KDA verify's) + # <<< route B + "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", + "tensorrt_llm._torch.moe.fused_moe.create_moe", + "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", + "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", + "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", + "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", + "tensorrt_llm._torch.modules.situ.SituAndMul", + "tensorrt_llm._torch.modules.rms_norm.RMSNorm", + "tensorrt_llm._torch.distributed.AllReduce", + "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", + # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries + # do not take a block, its context projection without the TP group's collective state, its context k / v, its + # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. + "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", + "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", + "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", + "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", +) + +# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware +# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). +DECODE_MAX_TOKENS = 8 +MAX_REQUESTS = 8 +MAX_TOKENS_PER_REQUEST = 8 + +# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own +# get_model_defaults sets it for the same reason). +_TOKENS_PER_BLOCK = 64 + +_LANG_PREFIX = "language_model." + + +# ---------------------------------------------------------------------------------------------------------------------- +# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and +# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, +# `K3DecodeMLA` below). +# ---------------------------------------------------------------------------------------------------------------------- + +# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections +# instead of the min-latency fused GEMM op (read once at import). +_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( + os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" +) + + +# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table +# holds cos=1/sin=0), but the chunked-context path indexes the table by +# absolute position, so it must cover max_position_embeddings (~512MB per +# backend for the 1M-position checkpoint); a smaller table is read out of +# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context +# deployments. +_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" + + +class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): + """DeepSeek-V3 routing, computed outside the TRTLLM-Gen MoE kernel. + + The kernel's own top-16 of 896 experts is the slow part of a generic MoE call of a few tokens. Outside it, the + scheduler routes with the backend's fused route + MXFP8 quantize (``trtllm::k3_route_quant``) up to 64 tokens and + with ``noaux_tc_op`` above, and the kernel takes the expert ids and weights. Both compute the kernel's arithmetic + and break ties the same way; a bf16 weight can differ by one ulp where it lies within fp32 rounding of a bf16 + rounding boundary. From about a thousand tokens on, ``noaux_tc_op`` takes longer than the kernel's routing. + """ + + @property + def requires_separated_routing(self) -> bool: + return True + + +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + @property + def routing_method(self) -> KimiK3MoeRoutingMethod: + """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return KimiK3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + +def _get_text_config(pretrained_config: "PretrainedConfig"): + """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" + if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( + not hasattr(pretrained_config, "linear_attn_config") + and hasattr(pretrained_config, "text_config") + ): + return pretrained_config.text_config + return pretrained_config + + +def _is_kda_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] + + +def _is_mla_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" + + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" + + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" + + +KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" + + +_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") + + +# Turning the port on should not require knowing which of the three topologies +# is the right one: "1" selects the measured policy (``split``), so enabling the +# feature and choosing the policy are one action through one variable. +_ATTN_RES_TOPOLOGY_ON = "1" + + +def _read_attn_res_topology() -> str: + """Default ``per_token``: the persistent kernel is opt-in. + + ``1`` is the only accepted on-value and resolves to ``split``. The named + topologies stay available for measurement: ``persistent`` uses the + persistent kernel at every shape it implements, ``per_token`` at none. + """ + raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") + if raw == _ATTN_RES_TOPOLOGY_ON: + return "split" + if raw not in _ATTN_RES_TOPOLOGIES: + # Loudly, for the same reason as the token ceiling below: a mistyped A/B + # arm that silently fell back to the default would measure one side + # twice and report no difference. + raise ValueError( + f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " + f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " + f"(which means 'split'), got {raw!r}" + ) + return raw + + +def _read_fused_attn_res_max_tokens() -> int: + """Resolved after the topology, because its default follows it. + + With the persistent port off -- the default -- the ceiling is 1, the + pre-existing gate: the fused epilogue is taken at the single-token decode + shape and nowhere else. Enabling the port raises it to 32, the top of the + measured range, so that "off" keeps meaning "unchanged". + """ + default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" + raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) + try: + value = int(raw) + except ValueError: + # Failing loudly matters more than usual here: a mistyped A/B arm that + # silently fell back to the default would measure the candidate twice + # and report no difference. + raise ValueError( + f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value < 1: + raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") + return value + + +_ATTN_RES_TOPOLOGY = _read_attn_res_topology() + + +_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() + + +def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: + """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. + + No token ceiling: the persistent grid is sized by the SM count, not by the + token count, so prefill is the case it exists for. + """ + del M # deliberately unused; see above + return H == 7168 and 2 <= N <= 9 + + +def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: + """Pick between the two fused kernels for this call site. + + ``persistent`` takes the persistent kernel at every shape it implements; + ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in + for the prefill/decode boundary. Shapes the persistent kernel does not + implement fall through to the caller's existing gate and land on the + unfused path. + """ + if not _persistent_attn_res_applicable(M, H, N): + return False + if _ATTN_RES_TOPOLOGY == "persistent": + return True + return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> Optional[torch.Tensor]: + """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. + + Returns ``None`` when the call falls outside the fused kernel's + contract (dtype/shape/arch) so the caller can use the exact fp32 reference + instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` + layout. Candidate order matches the reference: snapshots first, the + running prefix sum last. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + try: + attn_res_op = torch.ops.trtllm.attn_res_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output, _rsigma, _probs, _logits = attn_res_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + ) + return output.reshape(M, H) + + +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that + the shape gate let the call through, and a rejected call looks exactly like + a disabled one in the logs. Emitted at debug level: it is once per distinct + shape, not once per process, so it is a diagnostic rather than a summary. + """ + logger.debug_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, + default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, + # which is the measured window; larger token counts have not been measured + # and fall back to the unfused add + attn_res_fwd + RMSNorm path. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + _, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + None, + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) + return output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in + ``_apply_attn_res_rmsnorm_fused``). + + The production residual add produces a BF16 tensor that remains live across + the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual output, + while avoiding a separate add launch and a re-read of the intermediate by + attention-residual selection. + """ + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # Same measured window as _apply_attn_res_rmsnorm_fused above. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + updated_prefix_sum, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + addend.reshape(M, 1, H).contiguous(), + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = attn_res_add_rmsnorm_op( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + +def _apply_attn_res( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> torch.Tensor: + """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). + + prefix_sum: ``[num_tokens, hidden_size]`` + block_residual: ``[num_snapshots, num_tokens, hidden_size]`` + + Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's + contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. + Only the fallback boundary restores the HF ``[M, K, H]`` layout. + """ + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + block_residual_hf = block_residual.transpose(0, 1) + v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.eps) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the + fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_rmsnorm_fused( + prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in + ``_apply_attn_res_and_rmsnorm``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + + +# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where +# they are faster than the add -> attn_res -> RMSNorm chain. +_WIDE_ATTN_RES_MAX_TOKENS = 32 + + +def _attn_res_max_tokens(step: DecodeStep) -> int: + """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step + ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" + return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS + + +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + + +# The routed experts' module path prefixes. ``exclude_modules`` matches with +# wildcards and walks ancestor prefixes, so an empty prefix would widen what +# matches instead of just missing. +_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") + + +# Routed-expert quantization used when the checkpoint declares nothing per +# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors +# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that +# checkpoint is what this default has always served. +_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 + + +def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_packed_mxfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), + w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), + w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), + w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), + w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), + w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( + kinds=("weight_packed", "weight_scale"), + loader=_load_packed_mxfp4_expert, + loaded_slots_attr="_packed_mxfp4_loaded_slots", + needs_layer_finalize=False, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + self.num_experts = cfg.num_experts + self.top_k = cfg.num_experts_per_token + self.moe_hidden_size = cfg.routed_expert_hidden_size + # ValueError (not assert): these guard unsupported checkpoint + # configurations and must stay active under ``python -O``. + if self.moe_hidden_size is None: + raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") + if not getattr(cfg, "latent_moe_use_norm", False): + raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") + + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) + self._situ_betas = (situ_beta, situ_linear_beta) + dtype = torch.bfloat16 + + # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 + # accumulate/output (checkpoint stores the gate weight in bf16; saves a + # per-layer input cast + fp32 splitK pair on the bs1 decode path). + # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True + self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) + + routed_moe_model_config = self._routed_moe_model_config(model_config) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) + routed_moe_kwargs = dict( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_intermediate_size, + dtype=dtype, + # Kimi owns the latent reduction so it can order that collective + # after the shared expert's auxiliary-stream reduction. + reduce_results=False, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + # Let CommunicationFactory select the best available strategy. + communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. As measured once: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver + # instead, which reports the rejection trail. + # + # CUTLASS is absent on purpose: it is the fallback target, so + # "degraded to CUTLASS" is not a thing that can happen to it. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + ) + self._check_trtllm_situ_quant( + routed_moe_model_config.moe_backend, routed_quant_config.quant_algo + ) + + self.routed_experts = create_moe(**routed_moe_kwargs) + if not isinstance(self.routed_experts, ConfigurableMoE): + raise RuntimeError( + "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." + ) + if self.routed_experts.layer_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "dynamic EPLB or replicated expert slots." + ) + local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) + if local_expert_ids != list( + range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) + ): + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming currently requires a " + "contiguous static expert partition." + ) + self.local_expert_ids = tuple(local_expert_ids) + self.experts_per_rank = len(local_expert_ids) + self.expert_lo = local_expert_ids[0] + self.expert_hi = self.expert_lo + self.experts_per_rank + + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = model_config.mapping.tp_size > 1 + self._reduce_routed_output = ( + use_shared_tp + and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + ) + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self.shared_experts = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=shared_model_config, + reduce_output=use_shared_tp, + layer_idx=layer_idx, + is_shared_expert=True, + ) + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.moe_main_event = torch.cuda.Event() + self.moe_shared_event = torch.cuda.Event() + self.routed_expert_down_proj = nn.Linear( + cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype + ) + # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer + # fallback is the same fp32-variance eager math as KimiK3RMSNorm). + self.routed_expert_norm = RMSNorm( + hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. + self.decode_moe: Optional[_decode_moe.K3DecodeMoeLayer] = None + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, projection.weight.t(), None, None + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + # >>> route B: this target's split + """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit + ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (16 x 1, asserted at + construction).""" + # <<< route B + return mapping.moe_tp_size, mapping.moe_ep_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Routed-expert quantization for ``layer_idx``: the MXFP4 checkpoint declares nothing per layer (asserted at + construction), so the experts keep the ``W4A8_MXFP4_MXFP8`` default. + + An exclusion outranks the default: ``create_weights`` treats an override + as authoritative over anything ``__post_init__`` wrote, so this return + value stands in for both quantization passes and exclusion is the one + that runs second. It is matched as a pattern, so it is asked only about + real module names. + """ + quant_config = model_config.quant_config + if quant_config is not None and any( + quant_config.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + logger.debug( + "Kimi K3 layer %d routed experts: excluded from quantization, " + "keeping them unquantized", + layer_idx, + ) + return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) + + logger.debug( + "Kimi K3 layer %d routed experts: no per-layer quant config in the " + "checkpoint, defaulting to %s", + layer_idx, + _K3_DEFAULT_ROUTED_QUANT_ALGO, + ) + return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) + + @staticmethod + def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: + """Reject a routed-expert format trtllm-gen has no fused SiTu cubin for. + + trtllm-gen has fused SiTu FC1 cubins for two input formats and no + standalone SiTu activation kernel, so anything else has to die here + rather than in a cubin lookup deep inside the runner. Checked against + the resolved backend, not the K3 architecture branch, because the + generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also + land on TRTLLM. + + The admitted set is read off the backend rather than restated here, + because restating it is what broke. This guard was written in #17865 + when MXFP4 was the only fused SiTu drop; #17940 then added the NVFP4 + (group-16 ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``) cubins and updated + ``TRTLLMGenFusedMoE``'s set without touching this copy. For the week + in between, an NVFP4 K3 checkpoint could not start at all -- and not + only when TRTLLM was asked for by name, because + ``ModelConfig.resolve_moe_backend`` sends every K3 architecture to + TRTLLM, so the default AUTO configuration hit this raise too. The unit + tests did not catch it: they call ``create_moe`` directly and never + reach this guard, so the kernel path stayed green while the model path + was closed. + + A staticmethod, not an inline block, so that the invariant is + reachable from a test without constructing the whole runtime. + """ + situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() + if moe_backend != "TRTLLM" or quant_algo in situ_supported: + return + supported = ", ".join(sorted(algo.name for algo in situ_supported)) + raise ValueError( + f"Kimi K3 routed experts are quantized as {quant_algo}, which the " + "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " + f"exist only for {supported}. Set moe_config.backend to CUTLASS " + "or MEGAMOE_CUTEDSL." + ) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config, with the split of ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. + supported_backends = { + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + f"got {model_config.moe_backend!r}." + ) + if model_config.moe_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "EPLB or replicated expert slots." + ) + mapping = model_config.mapping + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + logger.info_once( + f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " + f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", + key="kimi_k3_moe_tp_ep_split", + ) + + mapping_dict = mapping.to_dict() + mapping_dict["moe_cluster_size"] = 1 + mapping_dict["moe_tp_size"] = moe_tp + mapping_dict["moe_ep_size"] = moe_ep + routed_mapping = Mapping.from_dict(mapping_dict) + + routed_model_config = copy.copy(model_config) + routed_model_config._frozen = False + routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) + routed_model_config.mapping = routed_mapping + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend in { + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + }: + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) + if configured_moe_max_num_tokens < default_moe_max_num_tokens: + logger.info_once( + "Kimi K3 MegaMoE raises moe_max_num_tokens from " + f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " + "because the global DP SymmBuffer requires capacity for " + "max_num_tokens * dp_size.", + key=( + "kimi_k3_megamoe_capacity_override_" + f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" + ), + ) + routed_model_config.moe_max_num_tokens = max( + configured_moe_max_num_tokens, + default_moe_max_num_tokens, + ) + routed_model_config._frozen = True + return routed_model_config + + def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: + """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel + share (``partial_tail``): the decode path takes the step with that tail.""" + return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + partial_tail: bool = False, + *, + step: DecodeStep, + ) -> Union[torch.Tensor, _decode_comm.PendingTail]: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification + (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where + ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at + most 8 tokens, a tensor on a wide decode step.""" + decode = self.decode_moe + if decode is not None and decode.takes(hidden_states, step, partial_tail): + return decode.forward( + self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push + ) + if partial_tail: + raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") + return self.forward_generic(hidden_states, all_rank_num_tokens) + + def forward_generic( + self, hidden_states: torch.Tensor, all_rank_num_tokens=None + ) -> torch.Tensor: + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None -from ..modeling import K3TargetSpec, ModelingV2KimiK3 -from ..modeling import _check_construction as _check_target_construction + def _routed_output(): + # Latent down/up projections via the min-latency fused GEMM op: + # at <=16 tokens (decode graphs) it runs a single pipelined + # bf16 kernel per projection instead of cuBLAS's split-K GEMV + + # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); + # for larger token counts the op falls back to cuBLAS internally. + # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear + # (A/B escape hatch). When the FP8 weight-read conversion has + # replaced the projection module, call it directly: its weight is + # an e4m3 buffer the bf16 dsv3 op must not read, and its forward + # is already a single fused GEMM (fp8_swap_ab_gemm). + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + if self._reduce_routed_output: + return y + # Communication-backed paths return a complete routed result. + y = self.routed_expert_norm(y) + return self._routed_projection(y, self.routed_expert_up_proj) -_TARGET = K3TargetSpec(sm=(10, 0), quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, moe_tp=16, moe_ep=1) + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. The shared GatedMLP includes its output all-reduce on + # the auxiliary stream. The join below must precede the routed + # all-reduce: concurrent collectives on different streams can corrupt + # SYMM_MEM all-reduce state. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + disable_on_compile=True, + ) + if self._reduce_routed_output: + routed_latent = moe_all_reduce(routed_out) + routed_latent = self.routed_expert_norm(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + return routed_out + shared_out + + +def resolve_attention_quant_config( + config: ModelConfig | None, layer_idx: int, projection: str +) -> QuantConfig: + """Resolve a checkpoint projection, including mixed-precision exclusions.""" + if config is None: + return QuantConfig() + global_config = config.quant_config or QuantConfig() + names = [ + f"{prefix}layers.{layer_idx}.self_attn.{projection}" + for prefix in ("language_model.model.", "model.", "") + ] + if any(global_config.is_module_excluded_from_quantization(name) for name in names): + return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) + declarations = config.quant_config_dict or {} + matches = [declarations[name] for name in names if name in declarations] + if matches: + selected = matches[0] + if any(match.quant_algo != selected.quant_algo for match in matches[1:]): + raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") + elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: + selected = QuantConfig() + else: + selected = global_config + if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " + f"quantization {selected.quant_algo}" + ) + if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): + raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") + result = copy.copy(selected) + result.kv_cache_quant_algo = global_config.kv_cache_quant_algo + return result + + +class KimiMLARuntime(nn.Module): + """Wraps K3 MLA and applies its external TP output reduction.""" + + def __init__( + self, + cfg: "PretrainedConfig", + layer_idx: int, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the + # final output reduction in this wrapper so the output gate remains + # between attention and the row-parallel o_proj. + mapping = model_config.mapping + reduce_output = mapping.tp_size > 1 + self._o_allreduce = ( + AllReduce( + mapping=mapping, + strategy=model_config.allreduce_strategy, + dtype=torch.bfloat16, + ) + if reduce_output + else None + ) + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in ( + "q_a_proj", + "kv_a_proj_with_mqa", + "q_b_proj", + "kv_b_proj", + "g_proj", + "o_proj", + ) + } + attention_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + attention_config._frozen = model_config._frozen + self.mixer = K3DecodeMLA( + hidden_size=cfg.hidden_size, + num_heads=cfg.num_attention_heads, + q_lora_rank=cfg.q_lora_rank, + kv_lora_rank=cfg.kv_lora_rank, + qk_nope_head_dim=cfg.qk_nope_head_dim, + qk_rope_head_dim=cfg.qk_rope_head_dim, + v_head_dim=cfg.v_head_dim, + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + use_output_gate=cfg.mla_use_output_gate, + max_position_embeddings=max_positions, + model_config=attention_config, + aux_stream_dict=aux_stream_dict, + ) + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" + return self.mixer.will_run_decode_branch(attn_metadata, step) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the + mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer( + None, hidden_states, attn_metadata, step=step, project_output=project_output + ) + return self._reduce(out) if project_output and reduce_output else out + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + out = self.mixer.forward_generic(None, hidden_states, attn_metadata) + return self._reduce(out) if reduce_output else out + + def _reduce(self, out: torch.Tensor) -> torch.Tensor: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + return self._o_allreduce(out) if self._o_allreduce is not None else out + + +class KimiLinearDecoderLayer(nn.Module): + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + dtype = torch.bfloat16 + + self.is_kda = _is_kda_layer(cfg, layer_idx) + is_mla = _is_mla_layer(cfg, layer_idx) + if self.is_kda == is_mla: + raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") + + if self.is_kda: + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in projection_names + } + attention_config._frozen = model_config._frozen + self.linear_attn = K3DecodeKDA( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + ) + + self.is_moe = ( + cfg.num_experts is not None + and layer_idx >= cfg.first_k_dense_replace + and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 + ) + if self.is_moe: + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's + # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take + # (decode_gemv.SITES); otherwise it stays within one node. + spans_nodes = self._mnnvl_allreduce() is not None + if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. + self.mlp = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, + ) + self._situ = (situ_beta, situ_linear_beta) + # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics + # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied + # after downcast, use_gemma=False). + self.input_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3). The res norms stay + # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by + # _apply_attn_res and the fused attn_res op, never called as + # modules. + self.attn_res_block_size = cfg.attn_res_block_size + assert self.attn_res_block_size is not None, ( + "Kimi K3 runtime expects attn_res_block_size to be set" + ) + self.self_attention_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set + # by the target's post_load_weights. + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + self.sandwich_oproj = False + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + step: DecodeStep, + capture: Optional[Tuple[Any, int]] = None, + prenormed: bool = False, + pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, + defer_moe_tail: bool = False, + ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + ``block_residual`` is a preallocated snapshot bank in kernel-native + ``[K_max, M, H]`` layout. Returns the running prefix sum and the + number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. + + ``step`` is the step's classification (``decode_step``), handed to the + attention module. + + ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is + this layer's input norm, and the layer's input, the step's embedding, + already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). + + The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens + runs the attention's all-reduce and the residual update in one + collective once the target has built its ``decode_comm``: with o_proj + (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, + the step and the attention's decode branch, else on the attention's + unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). + + ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the + previous layer's MoE output, unreduced (``defer_moe_tail``); then + ``hidden_states`` is the prefix sum without it, and this layer's + pre-attention step reduces and adds it: a ``PendingTail`` in one + ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its + all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer + the ``sandwich_tail`` call also stores the tap into the capture slot + where the speculative metadata exposes it (``capture_view``); other + tapped layers take the tap after the step. + + ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` + instead, ``partial`` this layer's MoE output unreduced, for the next + consumer to reduce. + """ + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + attn_res_max_tokens = _attn_res_max_tokens(step) + tail = ( + pending_moe_partial + if isinstance(pending_moe_partial, _decode_comm.PendingTail) + else None + ) + # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum + # straight into the bank row it snapshots. + snapshot_row = None + if tail is not None and self.layer_idx % self.attn_res_block_size == 0: + snapshot_row = block_residual[num_snapshots] + # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the + # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). + tap_view = None + if tail is not None and capture is not None: + tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) + + if prenormed: + assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 + assert pending_moe_partial is None and capture is None + elif tail is not None: + hidden_states, prefix_sum = self.decode_comm.sandwich_tail( + tail, + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + updated_out=snapshot_row, + tap=tap_view, + tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, + ) + elif pending_moe_partial is not None: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + elif capture is not None: + # The mixture tap needs the PRE-norm value, which the fused + # attn-res + RMSNorm kernel does not expose. Keep the two steps + # split on captured layers only and fuse everywhere else. + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if capture is not None and pending_moe_partial is not None and tap_view is None: + # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output + # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. + tapped = ( + _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + if _AUX_ATTN_RES_STREAM_ENABLED + else prefix_sum + ) + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + + if self.layer_idx % self.attn_res_block_size == 0: + if not prenormed and snapshot_row is None: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, step): + if ( + self.sandwich_oproj + and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS + and attention.will_run_decode_branch(attn_metadata, step) + ): + core = attention(hidden_states, attn_metadata, step=step, project_output=False) + hidden_states, prefix_sum = comm.sandwich_oproj( + core, + self._o_proj().weight, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + if comm is not None and step.wide: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) + else: + attention_output = attention(hidden_states, attn_metadata, step=step) + if prefix_sum is None: + prefix_sum = attention_output + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention_output, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + if self.is_moe and defer_moe_tail: + partial = self.block_sparse_moe( + hidden_states, all_rank_num_tokens, partial_tail=True, step=step + ) + return prefix_sum, num_snapshots, partial + if self.is_moe: + hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) + else: + hidden_states = self._dense_mlp(hidden_states, step) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def forward_generic( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + ) -> Tuple[torch.Tensor, int]: + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + if capture is not None: + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, None): + partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + elif prefix_sum is None: + prefix_sum = attention.forward_generic(hidden_states, attn_metadata) + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention.forward_generic(hidden_states, attn_metadata), + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + if self.is_moe: + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + hidden_states = self.block_sparse_moe.forward_generic( + hidden_states, all_rank_num_tokens + ) + else: + hidden_states = self.mlp(hidden_states) + return prefix_sum + hidden_states, num_snapshots + + def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: + """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode + kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" + gemvs = self.decode_gemvs + if gemvs is not None and step.small: + out = gemvs.dense_mlp( + hidden_states, + self.mlp.gate_up_proj.weight, + self.mlp.down_proj.weight, + *self._situ, + ) + if out is not None: + return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out + return self.mlp(hidden_states) + + def _mnnvl_allreduce(self): + """The MNNVL all-reduce of this layer's attention output, or None.""" + attention = self.linear_attn if self.is_kda else self.self_attn + return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) + + def _o_proj(self) -> nn.Module: + """This layer's attention output projection (row parallel).""" + return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj + + def _o_allreduce(self) -> AllReduce: + """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" + return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce + + def accepts_moe_partial(self, num_snapshots: int) -> bool: + """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target + built its collective state, and the snapshot bank is not empty.""" + return num_snapshots > 0 and self.decode_comm is not None + + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + + +class KimiLinearModel(DecoderModel): + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = _get_text_config(model_config.pretrained_config) + self._text_cfg = cfg + dtype = torch.bfloat16 + + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream() for _ in range(4)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + } + + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + + # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) + # by _apply_attn_res and the fused attn_res op. + self.output_attn_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.num_attn_res_snapshots = ( + cfg.num_hidden_layers + cfg.attn_res_block_size - 1 + ) // cfg.attn_res_block_size + # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the + # weights are final. None: every step embeds on the generic path. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + + # >>> route B: no per-token KDA verify states (kda_draft_records): no speculative decoding + @property + def kda_draft_records(self) -> bool: + """False: this target decodes without speculation, so the hybrid cache manager keeps no KDA state after + each verify token, and a captured step's latent-push decision (``_latent_push``) reads that.""" + return False + + # <<< route B + + def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, + from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from + load-time state, so every rank and every replay of the graph decides alike.""" + if not torch.cuda.is_current_stream_capturing(): + return False + kda = [layer.linear_attn for layer in self.layers if layer.is_kda] + mla = [layer.self_attn for layer in self.layers if not layer.is_kda] + return latent_push( + step, + capturing=True, + breakable=is_in_breakable_cuda_graph(), + kda_draft_records=self.kda_draft_records, + kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), + mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), + ) + + def _defer_moe_tail( + self, + i: int, + hidden_states: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + step: DecodeStep, + ) -> bool: + """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE + takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after + the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps + its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated + tail.""" + layer = self.layers[i] + if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): + return False + if spec_metadata is not None and ( + capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) + ): + return False + return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] + step = decode_step(attn_metadata, num_tokens) + if step is None: + return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) + if self._latent_push(attn_metadata, step): + step = replace(step, latent_push=True) + # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. + prenormed = None + if ( + inputs_embeds is None + and self.decode_gemvs is not None + and len(self.layers) > 0 + and self.num_attn_res_snapshots > 0 + ): + table = self.embed_tokens.weight + block_residual = table.new_empty( + self.num_attn_res_snapshots, num_tokens, table.shape[1] + ) + prenormed = self.decode_gemvs.embed_norm( + input_ids, table, self.layers[0].input_layernorm, block_residual + ) + if prenormed is not None: + hidden_states = prenormed + else: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) + # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. + pending_moe_partial = None + for i, layer in enumerate(self.layers): + defer_moe_tail = self._defer_moe_tail( + i, hidden_states, num_snapshots, spec_metadata, capture_set, step + ) + outputs = layer( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + step=step, + capture=self._capture(spec_metadata, capture_set, i), + prenormed=i == 0 and prenormed is not None, + pending_moe_partial=pending_moe_partial, + defer_moe_tail=defer_moe_tail, + ) + if defer_moe_tail: + hidden_states, num_snapshots, pending_moe_partial = outputs + else: + (hidden_states, num_snapshots), pending_moe_partial = outputs, None + + if isinstance(pending_moe_partial, _decode_comm.PendingTail): + normed, _ = self.layers[-1].decode_comm.sandwich_tail( + pending_moe_partial, + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + return normed + if pending_moe_partial is not None: + _, normed = _apply_attn_res_add_and_rmsnorm( + hidden_states, + _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + _attn_res_max_tokens(step), + ) + return normed + return self._final_norm( + hidden_states, + block_residual, + num_snapshots, + spec_metadata, + capture_set, + _attn_res_max_tokens(step), + ) + + def _generic_forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor], + inputs_embeds: Optional[torch.FloatTensor], + spec_metadata, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + for i, layer in enumerate(self.layers): + hidden_states, num_snapshots = layer.forward_generic( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + capture=self._capture(spec_metadata, capture_set, i), + ) + return self._final_norm( + hidden_states, block_residual, num_snapshots, spec_metadata, capture_set + ) + + def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + return (spec_metadata, self.layers[i - 1].layer_idx) + return None + + def _final_norm( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + max_fused_tokens: Optional[int] = None, + ) -> torch.Tensor: + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) + + return _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + max_fused_tokens, + ) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. +# ---------------------------------------------------------------------------------------------------------------------- + + +class K3DecodeKDA(KimiKDALinearAttention): + # >>> route B: no verify kernels + """Kimi K3's KDA attention: the built-in module, with the decode kernels on the steps they take. + + * A decode step of one token per request runs the fused input projection and the plain decode in one + ``ssm/k3_kda_decode_attn`` launch. + * On every step ``decode_step`` classifies, ``o_proj`` runs on the ``o_proj`` decode GEMV site where it takes the + rows, then the module's all-reduce. + + Every other step runs the built-in module. The kernels read one ``[q | k | v | g | f_a | b]`` weight, built at the + checkpoint load from the module's own, and the device's ``K3KdaBuffers`` and decode GEMVs' state, which the + target sets in ``post_load_weights``. + """ + + # <<< route B + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # The fused [q | k | v | g | f_a | b | pad] projection weight; the six projections' weights and the built-in + # fused [q | k | v | g] and [f_a | b | pad] ones are views of it. + self.k3_proj_weight: Optional[torch.Tensor] = None + # The fused projection's Lamport buffers: one set per device, shared by every KDA layer. + self.k3_buffers: Optional[K3KdaBuffers] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @property + def takes_k3_kernels(self) -> bool: + """Whether the decode kernels can run this layer: its fused projection weight is built.""" + return self.k3_proj_weight is not None + + def finalize_decode_weights(self) -> None: + """The built-in fused weights, then one ``[q | k | v | g | f_a | b | pad]`` weight of both.""" + super().finalize_decode_weights() + assert ( + self.use_full_rank_gate + and self.gate_lower_bound is not None + and self._qkvg_proj_weight is not None + and self._bfa_proj_weight is not None + and self._qkvg_proj_weight.dtype == self._bfa_proj_weight.dtype == torch.bfloat16 + ), ( + f"Kimi K3 KDA layer {self.layer_idx}: the decode kernels read the bf16 fused projections the built-in " + "module builds on CUDA at head dim 128, with a full-rank output gate and a gate lower bound" + ) + rows = self._qkvg_proj_weight.shape[0] + with torch.no_grad(): + fused = self._merge_projection_weights( + (self.q_proj, self.k_proj, self.v_proj, self.g_proj, self.f_a_proj, self.b_proj), + pad_rows_to=8, + ) + self.k3_proj_weight = fused + self._qkvg_proj_weight, self._bfa_proj_weight = fused[:rows], fused[rows:] + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode branch: outside a breakable CUDA graph.""" + return not is_in_breakable_cuda_graph() + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic`` where ``will_run_decode_branch`` does not hold. On the decode branch: the plain decode + on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the built-in dispatch; then + ``o_proj`` on its decode GEMV site and the TP all-reduce. + + ``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce), and, on the decode branch only, + ``project_output=False`` the post-o_norm core ``[N, H * 128]`` (no ``o_proj``).""" + if not self.will_run_decode_branch(attn_metadata, step): + return self.forward_generic(hidden_states, attn_metadata, reduce_output) + if ( + step.decode + and step.tokens_per_request == 1 + and self.takes_k3_kernels + and self.k3_buffers is not None + ): + core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) + else: + core = self._forward_impl(hidden_states, attn_metadata) + if not project_output: + return core.reshape(-1, self.proj_size) + return self._k3_project_output(core, reduce_output) + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + if reduce_output: + return super().forward(hidden_states, attn_metadata) + return self.o_proj( + self._builtin_core(hidden_states, attn_metadata).reshape(-1, self.proj_size) + ) + + def _builtin_core( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """The built-in forward's core ``[N, H, 128]``: inside a breakable CUDA graph from the eager + ``maybe_bcg_kda_core_inplace``, else from ``_forward_impl``.""" + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + return core + return self._forward_impl(hidden_states, attn_metadata) + + def _k3_project_output(self, core: torch.Tensor, reduce_output: bool = True) -> torch.Tensor: + """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then, with + ``reduce_output``, the TP all-reduce.""" + core2d = core.reshape(-1, self.proj_size) + out = None + if self.decode_gemvs is not None: + out = self.decode_gemvs.project("o_proj", core2d, self.o_proj.weight) + if out is None: + if reduce_output: + return self._project_output(core) + out = self.o_proj(core2d) + if reduce_output and self._o_allreduce is not None: + out = self._o_allreduce(out) + return out + + def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: + """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; + each slot's conv window and state advance in place.""" + mamba_metadata = attn_metadata.mamba_metadata + slots = getattr(mamba_metadata, "generation_state_indices", None) + if slots is None: + slots = mamba_metadata.state_indices[: x.shape[0]] + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + w_q, w_k, w_v = self._get_mtp_conv_weights() + core = k3_kda_decode_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.conv, + layer_cache.temporal, + slots, + self.k3_buffers, + float(self.gate_lower_bound), + self.head_k_dim**-0.5, + float(self.o_norm.eps), + ) + # Speculative decoding's replay caches keep their committed conv window in step with the pool's. + self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) + return core + + # >>> route B: no verify kernels (a verify step needs speculative decoding) + # <<< route B + + +class K3DecodeMLA(KimiK3MLAAttention): + """Kimi K3's MLA attention: the built-in module, with a decode step's attention on the decode kernels. + + A decode step runs ``x [W_a; W_g]^T`` with the gate columns through a sigmoid (the ``mla_ag`` decode GEMV site + where it takes the rows, else one GEMM), then ``attention/k3_mla_qkv`` (the q_a / kv_a RMSNorms, q_b and the k_b + absorption into the fused query, the step's latent rows stored into the paged cache) and + ``attention/k3_mla_attn_vb_out`` (the attention over the paged cache, v_b and the output gate in one launch), then + ``o_proj`` (the ``o_proj`` site where it takes the rows, else the module). Every other step, and a decode step whose + cache the kernels do not read (``k3_mla_decode_view`` says why), runs the built-in module. + + ``[W_a; W_g]`` is built at load from the module's weights, which become views of it. The attention workspace is + the device's ``K3MlaAttnWorkspace``; it and the decode GEMVs' state are set by the target in + ``post_load_weights``. + """ + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + # [W_a; W_g]: the fused q_a / kv_a projection's rows, then the output gate's. + self.k3_ag_weight: Optional[torch.Tensor] = None + # The decode attention's workspace: one per device, shared by every MLA layer. + self.k3_workspace: Optional[K3MlaAttnWorkspace] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + def post_load_weights(self) -> None: + """The built-in post-load, then ``[W_a; W_g]`` (once: CUDA graphs captured since read it).""" + super().post_load_weights() + if self.k3_ag_weight is not None: + return + gaps = self._k3_layout_gaps() + assert not gaps, ( + f"Kimi K3 MLA layer {self.layer_idx}: the decode kernels do not take {'; '.join(gaps)}" + ) + qkv_a, gate = self.kv_a_proj_with_mqa, self.g_proj + rows = qkv_a.weight.shape[0] + with torch.no_grad(): + fused = torch.cat([qkv_a.weight, gate.weight]) + qkv_a.weight = nn.Parameter(fused[:rows], requires_grad=False) + gate.weight = nn.Parameter(fused[rows:], requires_grad=False) + self.k3_ag_weight = fused + + def _k3_layout_gaps(self) -> list: + """What of this layer the decode kernels do not take (empty when they take all of it).""" + if not (self.use_output_gate and self.fuse_qkv_a_proj and not self.is_lite): + return ["a layer without the output gate or the fused q_a / kv_a projection"] + linears = (self.kv_a_proj_with_mqa, self.g_proj, self.q_b_proj, self.o_proj) + checks = ( + (not self.mapping.has_cp_helix(), "helix context parallelism"), + (not self.apply_rotary_emb and not self.llama_4_scaling, "RoPE or llama-4 scaling"), + (self.sparse_attn_hooks is None, "sparse attention"), + ( + self.kv_cache_dtype != "fp8_ds_mla" + and not getattr(self.mqa, "has_fp8_kv_cache", False) + and not getattr(self.mqa, "has_fp4_kv_cache", False), + "a quantized KV cache", + ), + ( + all(m.weight.dtype == torch.bfloat16 and m.bias is None for m in linears) + and self.k_b_proj_trans.dtype == self.v_b_proj.dtype == torch.bfloat16, + "projections other than bf16 and unbiased", + ), + ( + not getattr(self.q_a_layernorm, "is_nvfp4", False) + and not getattr(self.kv_a_layernorm, "use_gemma", False), + "an NVFP4 q_a norm or a Gemma kv_a norm", + ), + ( + self.num_heads_tp % 6 == 0 + and self.kv_lora_rank == 512 + and self.qk_rope_head_dim == 64 + and self.q_lora_rank == 1536, + f"{self.num_heads_tp} heads, latent {self.kv_lora_rank}, rope {self.qk_rope_head_dim}, " + f"q_lora {self.q_lora_rank}", + ), + ) + return [why for ok, why in checks if not ok] + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode kernels (see ``_k3_step_view``). Only there does it return + the gated attention output before ``o_proj``; the KV writes forbid running the attention twice, so a caller + that needs it asks first.""" + return self._k3_step_view(attn_metadata, step) is not None + + def _k3_step_view(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> Optional[dict]: + """The paged-cache view the decode kernels read on ``step``: a decode step, this layer's ``[W_a; W_g]`` and + workspace built, outside a breakable CUDA graph, and a cache ``k3_mla_decode_view`` takes. Else None.""" + if not ( + step.decode + and self.k3_ag_weight is not None + and self.k3_workspace is not None + and not is_in_breakable_cuda_graph() + ): + return None + return self._k3_decode_view(attn_metadata, step.num_tokens) + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic``, except on a decode step whose cache the decode kernels read + (``will_run_decode_branch``). There only, ``project_output=False`` returns the gated attention output + ``[M, H * 128]``, the input of ``o_proj``.""" + view = self._k3_step_view(attn_metadata, step) + if view is None: + if not project_output: + raise ValueError("project_output=False needs a step will_run_decode_branch takes") + return self.forward_generic(position_ids, hidden_states, attn_metadata) + x = hidden_states[: step.num_tokens].contiguous() + rows = self.kv_a_proj_with_mqa.weight.shape[0] + gemvs = self.decode_gemvs + ag = None if gemvs is None else gemvs.project("mla_ag", x, self.k3_ag_weight) + if ag is None: + ag = torch.nn.functional.linear(x, self.k3_ag_weight) + ag[:, rows:].sigmoid_() + fused_q = k3_mla_qkv( + ag, + self.q_a_layernorm.weight, + float(self.q_a_layernorm.variance_epsilon), + self.q_b_proj.weight, + self.k_b_proj_trans, + self.kv_a_layernorm.weight, + float(self.kv_a_layernorm.variance_epsilon), + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + ) + attn_output = self.create_output(x, 0) + k3_mla_attn_vb_out( + fused_q, + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + view["softmax_scale"], + self.v_b_proj, + attn_output, + self.k3_workspace, + gate=ag, + gate_col0=rows, + ) + if not project_output: + return attn_output + out = None if gemvs is None else gemvs.project("o_proj", attn_output, self.o_proj.weight) + if out is None: + out = self._project_output([attn_output], position_ids, attn_metadata, None) + return out + + def forward_generic( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + return super().forward(position_ids, hidden_states, attn_metadata) + + def _k3_decode_view(self, attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[dict]: + """The paged latent cache as the decode kernels read it this step, or None when they do not read it (the + reason is logged once).""" + view = k3_mla_decode_view(self.mqa, attn_metadata, num_tokens) + if isinstance(view, str): + logger.info_once( + f"Kimi K3 MLA: the built-in path for a decode step the decode kernels do not read ({view})", + key=f"k3_mla_decode_view_{view}", + ) + return None + return view + + +# ---------------------------------------------------------------------------------------------------------------------- +# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. +# ---------------------------------------------------------------------------------------------------------------------- + +#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at +#: max_draft_len 7. +DRAFTER_ATTN_SPLITS = frozenset( + {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} +) +# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k +# RMSNorm's epsilon and the NeoX RoPE's base. +_DRAFTER_HEADS = (6, 1) +_DRAFTER_HEAD_DIM = 64 +_DRAFTER_PAGE = 64 +_DRAFTER_EPS = 1e-5 +_DRAFTER_ROPE_BASE = 10000.0 + + +def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: + """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input + feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" + if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: + return None + width = in_features // tp_size + return tp_rank * width, (tp_rank + 1) * width + + +class K3FcSlice(nn.Module): + """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is + ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's + ranks is ``fc``'s output.""" + + def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: + super().__init__() + self.weight = nn.Parameter(weight, requires_grad=False) + self.start = start + self.end = end + + @classmethod + def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: + """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns + (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" + columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) + if columns is None: + return None + start, end = columns + return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the + GEMM reads the column block in place, through the rows' stride.""" + return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) + + +class K3DSparkDrafter(GQADSparkForCausalLM): + """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. + + A block the entries take runs, per layer: + + * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; + * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to + its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the + worker writes the accepted rows' context k / v before a block reads them; + * the output projection on ``drafter_o``, then the module's all-reduce; + * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the + module's all-reduce. + + The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state + (below); a projection whose site does not take its rows runs its module. Every other block runs the stock + ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or + normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under + CUDA-graph capture. + + Where the target hands the drafter the TP group's collective state (``use_decode_comm``): + + * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and + ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; + * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the + post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the + block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the + projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection + after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is + followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers + whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. + + Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the + stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); + the context k / v and the Markov head stay upstream code. + """ + + def __init__( + self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" + ) -> None: + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. + self._k3_take: Optional[bool] = None + # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes + # only these. + self._k3_attn_ran: set = set() + # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + # The zero residual of the split context projection's all-reduce, up to a decode step's rows. + self._k3_zero_rows: Optional[torch.Tensor] = None + # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms + # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. + self._k3_norms_fuse = False + self._k3_sandwich_forms: frozenset = frozenset() + + def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: + """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the + target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are + loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. + + * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they + split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` + sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. + * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and + RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose + shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not + compile is not used.""" + self.decode_comm = comm + self._k3_split_fc() + self._k3_norms_fuse = self._k3_norms_take_comm() + forms = set() + if self._k3_norms_fuse: + layers = self.model.layers + for form, swiglu, weights in ( + ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), + ("down", True, [layer.mlp.down_proj.weight for layer in layers]), + ): + if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( + weights[0], swiglu=swiglu + ): + forms.add(form) + self._k3_sandwich_forms = frozenset(forms) + logger.info( + "Kimi K3 DSpark drafter: residual adds and RMSNorms " + + ( + f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " + "mnnvl_fusion_allreduce)" + if self._k3_norms_fuse + else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" + ) + ) + + def _k3_norms_take_comm(self) -> bool: + """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and + post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no + quantized output), and every output and down projection is row parallel with its all-reduce.""" + hidden = self.config.hidden_size + + def plain(norm: nn.Module) -> bool: + weight = getattr(norm, "weight", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (hidden,) + and hasattr(norm, "variance_epsilon") + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "return_hp_output", True) + ) + + def reduces(linear: nn.Module) -> bool: + return ( + getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" + and getattr(linear, "reduce_output", False) + and getattr(linear, "all_reduce", None) is not None + ) + + return plain(self.model.norm) and all( + plain(layer.input_layernorm) + and plain(layer.post_attention_layernorm) + and reduces(layer.self_attn.o_proj) + and reduces(layer.mlp.down_proj) + for layer in self.model.layers + ) + + def _k3_tp_all_reduce(self) -> Optional[nn.Module]: + """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" + o_proj = self.model.layers[0].self_attn.o_proj + if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": + return None + return getattr(o_proj, "all_reduce", None) + + def _k3_split_fc(self) -> None: + """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero + residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is + a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP + all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" + fc = getattr(self, "fc", None) + if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): + return + mapping = self.model_config.mapping + weight = getattr(fc, "weight", None) + sliced = None + if ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and getattr(fc, "bias", None) is None + and self._k3_tp_all_reduce() is not None + ): + sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) + if sliced is None: + logger.info( + "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " + f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" + ) + return + self.fc = sliced + self._k3_zero_rows = weight.new_zeros( + MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] + ) + logger.info( + f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " + f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" + ) + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" + result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + self._k3_split_fc() + return result + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. + + With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group + with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode + step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then + ``hidden_norm``. Otherwise the stock projection.""" + fc = self.fc + if not isinstance(fc, K3FcSlice): + return super().project_target_hidden(hidden_states) + partial = fc(hidden_states.to(fc.weight.dtype)) + rows = partial.shape[0] + zeros = self._k3_zero_rows + if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): + normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) + return normed + if rows == 0: + return self.hidden_norm(partial) + return self.hidden_norm(self._k3_tp_all_reduce()(partial)) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class + docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows + ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows + as a view instead of gathering them.""" + keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) + if keys is None: + return super().dflash_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ) + out = self._k3_block_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ctx_rows_start, + ) + if not torch.cuda.is_current_stream_capturing(): + self._k3_attn_ran |= keys + return out + + def _k3_layers_take(self) -> bool: + """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head + layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, + non-causal unwindowed block attention, and no block convolution or post-attention gate.""" + if self._k3_take is None: + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + take = ( + self._use_fused_qk_norm_rope + and not self.has_block_conv + and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate + ) + for layer_idx, layer in enumerate(self.model.layers): + attn, mlp = layer.self_attn, layer.mlp + pos = getattr(attn, "pos_embd_params", None) + rope = getattr(pos, "rope", None) + norms = (attn.q_norm, attn.k_norm) + take = take and ( + (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS + and attn.head_dim == _DRAFTER_HEAD_DIM + and getattr(attn, "is_qk_norm", False) + and not getattr(attn, "use_gemma_rms_norm", True) + and rope is not None + and pos.is_neox + and getattr(pos, "mrope_section", None) is None + and rope.theta == _DRAFTER_ROPE_BASE + and all( + norm.variance_epsilon == _DRAFTER_EPS + and norm.weight.dtype == torch.bfloat16 + and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) + for norm in norms + ) + and attn.qkv_proj.bias is None + and attn.o_proj.bias is None + and mlp.activation is torch.nn.functional.silu + and mlp.swiglu_limit is None + and mlp.swiglu_alpha in (None, 1.0) + and mlp.swiglu_beta in (None, 0.0) + and mlp.gate_up_proj.bias is None + and mlp.down_proj.bias is None + and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) + ) + self._k3_take = bool(take) + logger.info( + f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " + + ( + "on k3_drafter_attn_qknorm and the drafter GEMV sites" + if self._k3_take + else "on the stock block forward (a layer the drafter entries do not take)" + ) + ) + return self._k3_take + + def _k3_block_keys( + self, + noise_embedding: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor], + ctx_page_table: Optional[torch.Tensor], + ) -> Optional[frozenset]: + """The attention's compile keys for this block when the drafter entries take it, else None.""" + if ( + self.dflash_attention_backend != "TRTLLM" + or ctx_kv_cache is None + or ctx_page_table is None + or noise_embedding.dtype != torch.bfloat16 + or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS + or is_in_breakable_cuda_graph() + or not self._k3_layers_take() + ): + return None + num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads + page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) + dense = ( + num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + ) + strides = set() + for layer_idx in range(len(self.model.layers)): + cache = ctx_kv_cache[layer_idx] + if ( + cache.dtype != torch.bfloat16 + or cache.dim() != 5 + or tuple(cache.shape[1:]) != page + or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) + or cache.stride(0) % 8 + ): + return None + strides.add(cache.stride(0)) + keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) + if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: + return None + return keys + + def _k3_block_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: torch.Tensor, + ctx_page_table: torch.Tensor, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + batch, block = noise_embedding.shape[:2] + rows = batch * block + ctx_len = num_ctx_per_req[:batch].to(torch.int32) + if ctx_rows_start is None: + page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) + else: + # The batch's page-table rows are one contiguous run: a view, no gather. + page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] + positions = query_positions.reshape(-1).contiguous() + hidden = noise_embedding.reshape(rows, -1) + layers = self.model.layers + comm = self._k3_fused_comm(hidden) + residual = None + if comm is not None: + # The fused all-reduces read the residual and return the updated one: the block's input serves as the + # first residual, uncopied. + residual = hidden + normed = layers[0].input_layernorm(hidden) + for layer_idx, layer in enumerate(layers): + attn = layer.self_attn + if comm is None: + if residual is None: + residual = hidden.clone() + normed = layer.input_layernorm(hidden) + else: + normed, residual = layer.input_layernorm(hidden, residual) + qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) + out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) + k3_drafter_attn_qknorm( + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + positions, + attn.q_norm.variance_epsilon, + attn.pos_embd_params.rope.theta, + ctx_kv_cache[layer_idx], + page_table, + ctx_len, + attn.num_heads, + attn.num_key_value_heads, + out, + ) + if comm is not None: + # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after + # the last layer the final norm. + next_norm = ( + layers[layer_idx + 1].input_layernorm + if layer_idx + 1 < len(layers) + else self.model.norm + ) + normed, residual = self._k3_project_norm( + comm, out, attn.o_proj, residual, layer.post_attention_layernorm + ) + normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) + continue + hidden = self._k3_project("drafter_o", out, attn.o_proj) + hidden, residual = layer.post_attention_layernorm(hidden, residual) + hidden = self._k3_mlp(layer.mlp, hidden) + if comm is not None: + return normed + out, _ = self.model.norm(hidden, residual) + return out + + def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: + """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its + all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace + holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" + comm = self.decode_comm + if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): + return None + return comm + + def _k3_project_norm( + self, + comm: _decode_comm.K3DecodeComm, + x: torch.Tensor, + linear: nn.Module, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, + the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the + projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` + where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" + if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( + x, linear.weight, residual, norm + ): + return comm.sandwich_plain(x, linear.weight, residual, norm) + gemvs = self.decode_gemvs + partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) + if partial is None: + partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_mlp_norm( + self, + comm: _decode_comm.K3DecodeComm, + mlp: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down + projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul + form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down + projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; + then ``comm/mnnvl_fusion_allreduce``.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + partial = None + if gate_up is not None: + if "down" in self._k3_sandwich_forms and comm.takes_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ): + return comm.sandwich_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ) + partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + if partial is None: + partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: + """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel + projection's all-reduce.""" + y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) + if y is None: + return linear(x) + return _row_parallel_reduce(linear, y) + + def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on + ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + down = ( + None + if gate_up is None + else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + ) + if down is None: + return mlp(x) + return _row_parallel_reduce(mlp.down_proj, down) + + +def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: + """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" + all_reduce = getattr(linear, "all_reduce", None) + if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: + return all_reduce(partial) + return partial + + +# ---------------------------------------------------------------------------------------------------------------------- +# The target: step classification, the construction checks and the class route returns. +# ---------------------------------------------------------------------------------------------------------------------- + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. + + The checkpoint names its language-model modules `language_model.`; the text model's are ``, with + `layers.*` under `model.`. + """ + config = model_config.pretrained_config + text = copy.copy(model_config) + text._frozen = False + text.pretrained_config = config.text_config + excluded = text.quant_config.exclude_modules + if excluded: + text.quant_config = copy.copy(text.quant_config) + renamed = [] + for name in excluded: + if name.startswith(_LANG_PREFIX): + name = name[len(_LANG_PREFIX) :] + if name.startswith("layers."): + name = "model." + name + renamed.append(name) + text.quant_config.exclude_modules = renamed + text.skip_create_weights_in_init = True + text._frozen = True + return text + + +@dataclass(frozen=True) +class DecodeStep: + """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` + generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" + + num_tokens: int + num_requests: Optional[int] = None + tokens_per_request: Optional[int] = None + #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` + #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. + latent_push: bool = False + + @property + def small(self) -> bool: + """Whether the step fits one token tile of the token-count kernels.""" + return self.num_tokens <= DECODE_MAX_TOKENS + + @property + def decode(self) -> bool: + """Whether the step is a pure decode step the request-aware kernels take.""" + return self.num_requests is not None + + @property + def wide(self) -> bool: + """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode + layout's MoE head and tail, on M-general ops.""" + return self.decode and not self.small + + +def latent_push( + step: DecodeStep, + *, + capturing: bool, + breakable: bool, + kda_draft_records: bool, + kda_decode_kernels: bool, + mla_decode_branch: bool, +) -> bool: + """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at + most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a + breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels + (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA + layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). + Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel + between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a + graph launch orders every pushing replay behind whatever ran before it.""" + if not (step.decode and step.small) or not capturing or breakable: + return False + if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): + return False + return mla_decode_branch + + +def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: + """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). + + ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers + and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input + here is fixed by that shape, so a captured step and its replays are classified alike. + """ + if num_tokens <= 0: + return None + requests = _decode_requests(attn_metadata, num_tokens) + if requests is not None: + return DecodeStep(num_tokens, requests, num_tokens // requests) + if num_tokens <= DECODE_MAX_TOKENS: + return DecodeStep(num_tokens) + return None + + +def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: + """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" + if attn_metadata.num_contexts != 0: + return None + requests = attn_metadata.num_generations + if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: + return None + tokens = num_tokens // requests + if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: + return None + seq_lens = getattr(attn_metadata, "seq_lens", None) + if seq_lens is not None and seq_lens.device.type == "cpu": + lens = seq_lens[:requests] + if lens.numel() != requests or bool((lens != tokens).any()): + return None + return requests def _check_construction(model_config: ModelConfig) -> None: - _check_target_construction(model_config, _TARGET) + """The settings this target is built for that are fixed before the first step.""" + capability = torch.cuda.get_device_capability() + assert capability == _SM, ( + f"this target is certified on sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" + ) + mapping = model_config.mapping + topology = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.enable_attention_dp, + ) + # >>> route B: this target's topology, its expert split set explicitly, and no speculative decoding + assert topology == (16, 16, 1, 16, 1, False), ( + "the tp16_moetp16ep1 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " + "moe_tensor_parallel_size 16, moe_expert_parallel_size 1 and enable_attention_dp false; the engine built " + f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" + ) + assert mapping.moe_tp_ep_user_specified, ( + "the tp16_moetp16ep1 target needs moe_tensor_parallel_size 16 and moe_expert_parallel_size 1 set " + "explicitly; with the split unset the engine resolves the same sizes, but Kimi K3 then runs the experts " + "expert-parallel over the 16 ranks" + ) + assert getattr(model_config, "spec_config", None) is None, ( + "the tp16_moetp16ep1 target decodes without speculation; for DSpark, DFlash or SA set " + "moe_tensor_parallel_size 4 and moe_expert_parallel_size 4 (the tp16_moetp4ep4 target)" + ) + # <<< route B + assert model_config.torch_dtype == torch.bfloat16, ( + f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" + ) + kv_algo = model_config.quant_config.kv_cache_quant_algo + assert kv_algo is None, ( + f"this target's MLA kernels read a bf16 KV pool; kv_cache_config.dtype resolved to {kv_algo}" + ) + quant_algo = model_config.quant_config.quant_algo + assert quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( + "this target loads the MXFP4 checkpoint, whose compressed-tensors config the model config reads as " + "W4A16_MXFP4 with no per-layer declarations (its routed experts run the W4A8_MXFP4_MXFP8 default); the " + f"engine read {quant_algo} with {len(model_config.quant_config_dict or {})} per-layer declarations" + ) + strategy = model_config.allreduce_strategy + assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( + f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" + ) + + +# >>> route B: this target's name +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(KimiLinearForCausalLM): + # <<< route B + """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. + + It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through + weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model + keeps the built-in one's. + """ + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The V2 hybrid manager holds the KDA states; the step contract requires it.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + config = model_config.pretrained_config + assert getattr(config, "text_config", None) is not None, ( + "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " + "text_config" + ) + _check_construction(model_config) + text = _text_model_config(model_config) + # >>> route B: no speculative decoding (_check_construction), so the one-engine shell builds no drafter + # <<< route B + # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 + # weight-read conversion of the shared / latent MLPs. + self._fp8_weight_read_moe_mlp = False + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + cfg = text.pretrained_config + SpecDecOneEngineForCausalLM.__init__( + self, + KimiLinearModel(text), + text, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self._own_spec_branch(text) + self._step_checked = False + # >>> route B: no drafter to hand the LM head to + # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. + self.logits_processor = _decode_gemv.K3LogitsProcessor(self.logits_processor) + # <<< route B + # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which + # must therefore be the text config, as the built-in wrapper leaves it. + model_config._frozen = False + model_config.pretrained_config = self.config + model_config._frozen = True + + def _own_spec_branch(self, model_config: ModelConfig) -> None: + """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a + standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is + `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" + spec_config = model_config.spec_config + if spec_config is None: + return + mode = spec_config.spec_dec_mode + standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) + if ( + standalone_dspark + and self.draft_config is not None + and not is_mla(self.draft_config.pretrained_config) + ): + drafter = K3DSparkDrafter( + self.draft_config, dflash_attention_backend=spec_config.attention_backend + ) + if mode.is_parallel_draft(): + drafter.logits_processor = self.logits_processor + self.epilogue[self.epilogue.index(self.draft_model)] = drafter + self.draft_model = drafter + if mode.is_dflash(): + worker_cls = _spec_worker.KimiK3DFlashWorker + elif standalone_dspark: + worker_cls = _spec_worker.KimiK3DSparkWorker + else: + return + worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) + worker.set_draft_model(self.draft_model) + self.epilogue[self.epilogue.index(self.spec_worker)] = worker + self.spec_worker = worker + + def load_weights(self, weights, *args, **kwargs): + _weights.load(self, weights) + + def cache_derived_state(self) -> None: + """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of + every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: + a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites + are among them, and the drafter gets the state too.""" + super().cache_derived_state() + if self.model.decode_gemvs is not None: + return + gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head) + self.model.decode_gemvs = gemvs + self.logits_processor.gemvs = gemvs + for layer in self.model.layers: + if not layer.is_moe: + layer.decode_gemvs = gemvs + if isinstance(self.draft_model, K3DSparkDrafter): + self.draft_model.decode_gemvs = gemvs + + def post_load_weights(self) -> None: + """The state the decode kernels share, built once per device before any CUDA-graph capture and handed to + every layer of its kind: the KDA projection's Lamport buffers and the MLA decode attention's workspace; the + decode GEMVs' state (built by ``cache_derived_state``) handed to every attention module; and, where every + attention all-reduce runs over MNNVL, the TP group's collective state (``K3DecodeComm``, collective: every + rank builds it here) handed to every layer, with whether its sandwich takes the layer's o_proj, the decode + path's one-shot ceiling on every stock MNNVL all-reduce (``use_decode_one_shot``), and the MoE decode path + (``_build_decode_moe``); then the speculative worker's decode kernels (``_gate_spec_worker_kernels``) and the + DSpark drafter's collectives (``_gate_drafter_comm``).""" + super().post_load_weights() + kda = [layer.linear_attn for layer in self.model.layers if layer.is_kda] + mla = [layer.self_attn.mixer for layer in self.model.layers if not layer.is_kda] + if kda[0].k3_buffers is not None: + return # built by an earlier call; CUDA graphs captured since hold it + device = self.model.embed_tokens.weight.device + buffers = K3KdaBuffers.create(device) + workspace = K3MlaAttnWorkspace.create(device, mla[0].num_heads_tp // 6) + for module in kda: + module.k3_buffers = buffers + for module in mla: + module.k3_workspace = workspace + for module in kda + mla: + module.decode_gemvs = self.model.decode_gemvs + layers = self.model.layers + comm = None + if all(layer._mnnvl_allreduce() is not None for layer in layers): + for layer in layers: + layer.sandwich_oproj = _decode_comm.K3DecodeComm.takes_oproj( + layer._o_proj(), self.model.num_attn_res_snapshots + ) + oproj = next((layer._o_proj().weight for layer in layers if layer.sandwich_oproj), None) + comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, oproj) + for layer in layers: + layer.decode_comm = comm + _decode_comm.use_decode_one_shot(self) + moe_layers = 0 if comm is None else self._build_decode_moe(comm) + logger.info( + # >>> route B: no verify kernels + "Kimi K3 decode kernels: KDA on k3_kda_decode_attn " + # <<< route B + f"({sum(m.takes_k3_kernels for m in kda)} / {len(kda)} layers take them), MLA on k3_mla_qkv and " + f"k3_mla_attn_vb_out ({len(mla)} layers); the post-attention all-reduce of at most " + f"{_decode_comm.AR_ATTN_RES_MAX_TOKENS} tokens " + + ( + "with the residual update: k3_sandwich_oproj with o_proj at most " + f"{_decode_comm.SANDWICH_MAX_TOKENS} decode tokens " + f"({sum(layer.sandwich_oproj for layer in layers)} / {len(layers)} layers), else " + "mnnvl_allreduce_attn_res" + if comm is not None + else "unfused (an attention all-reduce does not run over MNNVL)" + ) + # >>> route B: this target's MoE engines + + f"; MoE on k3_moe_front, k3_moe_m1 / k3_moe_m2 / k3_moe and the row-parallel tail ({moe_layers} layers)" + # <<< route B + ) + self._gate_spec_worker_kernels(comm) + self._gate_drafter_comm(comm) + + def _build_decode_moe(self, comm: _decode_comm.K3DecodeComm) -> int: + """The MoE decode path (``decode_moe.py``) on every MoE layer it takes: the shared state (collective: the + front's head workspace), each layer's decode weights with its latent norm folded into the latent up + projection, the decode GEMVs' state; then one call of every kernel of the path (collective: the front and + the sandwich tail). Every rank builds it here. Returns the number of MoE layers it takes.""" + mapping = self.model_config.mapping + moes = [layer.block_sparse_moe for layer in self.model.layers if layer.is_moe] + gaps = { + id(moe): _decode_moe.layout_gaps( + moe, mapping.tp_size, self.model.num_attn_res_snapshots + ) + for moe in moes + } + takes = [moe for moe in moes if not gaps[id(moe)]] + for reason in sorted({gap for moe in moes for gap in gaps[id(moe)]}): + logger.info_once( + f"Kimi K3 MoE decode path off on some layers: {reason}", + key=f"k3_decode_moe_off_{reason}", + ) + if not takes: + return 0 + backend = takes[0].routed_experts.backend + state = _decode_moe.K3DecodeMoe.create( + mapping, + backend.w3_w1_weight.device, + backend.w3_w1_weight.shape[1] // 2, + backend.expert_size_per_partition, + comm.mnnvl, + # >>> route B: the experts' logical slice (192 of the loader's 256), which k3_moe_m1 / k3_moe_m2 stream + i_logical=( + getattr(backend.quant_method, "intermediate_size_per_partition_lean", None) + or backend.intermediate_size_per_partition + ), + # <<< route B + ) + for moe in takes: + _decode_moe.fold_latent_norm(moe) + moe.decode_moe = _decode_moe.K3DecodeMoeLayer.create( + moe, state, mapping.tp_rank, mapping.tp_size + ) + moe.decode_gemvs = self.model.decode_gemvs + first = takes[0].decode_moe + first.warm_up(takes[0]) + logger.info( + "Kimi K3 MoE decode path: the latent all-reduce at <= 8 tokens on " + + ( + "the routed experts' push form and k3_latent_reduce on pushing steps (a pure decode step captured " + "into a CUDA graph on the decode kernels), else the routed experts' all-reduce" + if state.exchange is not None + else f"the routed experts' all-reduce (no latent exchange at TP {mapping.tp_size})" + ) + ) + comm.compile_tail(takes[0].moe_hidden_size, first.shared_cols, first.tail_weight) + return len(takes) + + def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, + ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input + that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's + ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the + vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks + each step's own conditions. Returns the setting (False without such a worker).""" + worker = getattr(self, "spec_worker", None) + if not hasattr(worker, "k3_decode"): + return False + gemvs = self.model.decode_gemvs + worker.k3_decode = ( + comm is not None and gemvs is not None and gemvs.head_workspace is not None + ) + logger.info( + f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " + + ( + "on" + if worker.k3_decode + else "off (no MNNVL decode state or no k3_head_gemv LM head)" + ) + ) + return worker.k3_decode + def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built + it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context + projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms + in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on + the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate + requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` + workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the + drafter took the state (False without such a drafter).""" + drafter = getattr(self, "draft_model", None) + if comm is None or not isinstance(drafter, K3DSparkDrafter): + return False + drafter.use_decode_comm(comm) + return True -class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(ModelingV2KimiK3): - """Concrete deployment of the shared Kimi K3 text model.""" + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """First-forward checks of the engine surface and the per-engine settings.""" + objects = { + "attn_metadata": attn_metadata, + "kv_cache_manager": attn_metadata.kv_cache_manager, + } + missing = [ + f"{owner}.{name}" + for owner, names in REQUIRED_ENGINE_FIELDS.items() + for name in names + if not hasattr(objects[owner], name) + ] + assert not missing, f"engine fields this target reads are missing: {missing}" + assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( + f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " + f"{attn_metadata.tokens_per_block}" + ) + manager = attn_metadata.kv_cache_manager + assert isinstance(manager, MambaHybridCacheManagerV2), ( + "this target needs the V2 hybrid KV / state cache manager " + "(kv_cache_config.use_kv_cache_manager_v2); the engine built " + f"{type(manager).__name__}" + ) + assert not manager.enable_block_reuse, ( + "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" + ) + kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) + state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype + assert state_dtype == torch.float32, ( + "this target's KDA kernels keep fp32 recurrent states (kv_cache_config.mamba_ssm_cache_dtype float32 or " + f"auto); the engine built a {state_dtype} state pool" + ) + self._step_checked = True - target_spec = _TARGET + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + if kwargs.pop("multimodal_params", None): + raise ValueError( + "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" + ) + if not self._step_checked: + self._check_step_contract(attn_metadata) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py new file mode 100644 index 000000000000..abc1a1388213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py @@ -0,0 +1,1343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import functools +import inspect +from typing import Optional + +import torch +from torch import nn + +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.speculative.dflash import ( + DFlashSpecMetadata, + DFlashWorker, + dflash_allocated_ctx_limit, + dflash_noise_block_embedding, +) +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase +from tensorrt_llm.logger import logger + + +@functools.lru_cache(maxsize=None) +def _takes_ctx_rows_start_cls(cls) -> bool: + return "ctx_rows_start" in inspect.signature(cls.dflash_forward).parameters + + +def _takes_ctx_rows_start(draft_model) -> bool: + """Whether the draft model's ``dflash_forward`` accepts ``ctx_rows_start``.""" + return _takes_ctx_rows_start_cls(type(draft_model)) + + +def dflash_ctx_rows_kwargs(draft_model, ctx_rows_start: Optional[int]) -> dict: + """The ``ctx_rows_start`` keyword (from ``prepare_1st_drafter_inputs``) for a draft model whose ``dflash_forward`` + takes it; no keyword for another draft model or without a start.""" + if ctx_rows_start is None or not _takes_ctx_rows_start(draft_model): + return {} + return {"ctx_rows_start": ctx_rows_start} + + +def capture_view( + spec_metadata: Optional[SpecMetadata], layer_id: int, num_tokens: int +) -> Optional[torch.Tensor]: + """``layer_id``'s slot of the DFlash capture buffer for ``num_tokens`` rows (a strided view a kernel can write the + tap into), or None when ``spec_metadata`` is not DFlash's or the layer is not captured.""" + if ( + not isinstance(spec_metadata, DFlashSpecMetadata) + or spec_metadata.captured_hidden_states is None + ): + return None + i = spec_metadata._layer_to_idx.get(layer_id) + if i is None: + return None + return spec_metadata.captured_hidden_states[ + :num_tokens, i * spec_metadata.hidden_size : (i + 1) * spec_metadata.hidden_size + ] + + +class KimiK3DFlashWorker(DFlashWorker): + # Set by a Kimi K3 target running its decode kernels: decode steps then run the acceptance and the drafter's + # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context + # K/V with trtllm::k3_ctx_kv. + k3_decode = False + + def _mask_token_id(self, draft_model) -> int: + """The drafter's mask token id, resolved once.""" + if self._resolved_mask_token_id is None: + if ( + hasattr(self.spec_config, "mask_token_id") + and self.spec_config.mask_token_id is not None + ): + self._resolved_mask_token_id = self.spec_config.mask_token_id + elif hasattr(draft_model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.mask_token_id + elif hasattr(draft_model.model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.model.mask_token_id + else: + raise ValueError( + "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " + "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." + ) + return self._resolved_mask_token_id + + @staticmethod + def _trained_mask_embedding(draft_model, mask_token_id: int) -> Optional[torch.Tensor]: + """The drafter's own trained mask row, if it kept one for the mask id in use.""" + if mask_token_id != getattr(draft_model, "mask_token_id", None): + return None + return getattr(draft_model, "mask_token_embedding", None) + + def _k3_accept_applies( + self, logits, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Whether this step's acceptance and the drafter's inputs run as ``trtllm::k3_spec_accept``. + + That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no + rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target + logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block + table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. + """ + if not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return False + K = spec_metadata.runtime_draft_len + if ( + logits.dim() != 2 + or logits.dtype != torch.float32 + or logits.shape[0] != num_gens * (K + 1) + ): + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + embed = draft_model.draft_model_full.model.embed_tokens + return accept_op.supports( + logits.shape[1], num_gens, self._compute_block_size, K, embed.weight.shape[1] + ) + + def _k3_accept_step_applies( + self, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Everything ``_k3_accept_applies`` checks but the target logits.""" + if not self.k3_decode: + return False + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + + if num_contexts != 0 or not 0 < num_gens <= 8 or self.guided_decoder is not None: + return False + if ( + not spec_metadata.is_all_greedy_sample + or self._can_use_rejection_sampling(spec_metadata) + or getattr(spec_metadata, "enable_penalty", False) + or type(self)._reshape_draft_tokens_for_accept + is not SpecWorkerBase._reshape_draft_tokens_for_accept + or type(self)._reshape_logits_for_accept + is not SpecWorkerBase._reshape_logits_for_accept + ): + return False + K = spec_metadata.runtime_draft_len + draft_tokens = spec_metadata.draft_tokens + if ( + K <= 0 + or draft_tokens is None + or draft_tokens.dtype != torch.int32 + or draft_tokens.numel() != num_gens * K + ): + return False + mgr = attn_metadata.kv_cache_manager + mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) + if ( + not isinstance(mgr, MambaHybridCacheManagerV2) + or type(mgr).update_mamba_states is not MambaHybridCacheManagerV2.update_mamba_states + or not mgr.use_kda_replay_update + or getattr(mgr, "prev_num_accepted_tokens", None) is None + or getattr(mgr, "_dummy_request_mask", None) is None + or mamba_metadata is None + or mamba_metadata.state_indices.dtype != torch.int32 + ): + return False + if ( + self._ctx_block_tables is None + or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None + or getattr(attn_metadata, "kv_lens_cuda", None) is None + ): + return False + embed = draft_model.draft_model_full.model.embed_tokens + return getattr(embed, "tp_size", 1) == 1 and embed.weight.dtype == torch.bfloat16 + + def target_logits( + self, hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) -> torch.Tensor: + """This rank's bf16 vocabulary shard of the target logits, [rows, vocab / TP] (``_k3_head_shard``: the head + GEMM without its all-gather and fp32 cast), for a step whose acceptance runs as the sharded + ``trtllm::k3_spec_accept`` (see ``_k3_logits_shard``); the gathered fp32 logits otherwise. At such a step the + acceptance is their only reader: the one-model spec sampler takes the worker's tokens and rejects requests for + logits or log probabilities, so the shard is also what the step returns as ``"logits"``.""" + self._k3_step_shard = None + shard = self._k3_logits_shard( + hidden_states, lm_head, attn_metadata, spec_metadata, draft_model + ) + if shard is None: + return super().target_logits( + hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) + logits = self._k3_head_shard(logits_processor, lm_head, hidden_states) + self._k3_step_shard = (logits, shard) + return logits + + @staticmethod + def _k3_head_shard(logits_processor, lm_head, rows: torch.Tensor) -> torch.Tensor: + """This rank's bf16 vocabulary shard of ``lm_head(rows)``, without the head's all-gather: from the logits + processor's own head kernel where it has one that takes the rows (``lm_head_shard``; the Kimi K3 target's + ``gemm/k3_head_gemv``), so the shard holds the values of the processor's gathered logits; else from the head's + own GEMM.""" + head_shard = getattr(logits_processor, "lm_head_shard", None) + logits = None if head_shard is None else head_shard(rows, lm_head) + if logits is None: + logits = lm_head.apply_linear(rows, lm_head.bias) + return logits + + def _k3_logits_shard(self, hidden_states, lm_head, attn_metadata, spec_metadata, draft_model): + """``(workspace, first column)`` when this step's target logits stay vocabulary-sharded, else None. + + That needs a step ``_k3_accept_step_applies`` takes, plain TP whose all-reduces own an MNNVL workspace for + this mapping, and an unquantized, bias-free, unpadded, evenly split column-parallel bf16 lm_head with shards + the kernel splits. The exchange is collective, so every rank must decide alike: every condition is the + batch's or the configuration's. The workspace is allocated (collectively) and the kernel compiled outside + CUDA-graph capture only, so a graph whose warmup allocated no workspace is captured on the gathered path. + """ + if not self.k3_decode: + return None + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + mapping = self.mapping + if getattr(self, "_k3_shard_head", None) is not lm_head: + reason = self._k3_logits_shard_reason(lm_head) + if reason is None and MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None: + # Not kept: re-checked on later steps, since the model's MNNVL workspace may not exist yet. + logger.info_once( + "DFlash: the target logits are all-gathered (the TP all-reduces are not MNNVL)", + key="dflash_k3_logits_no_mnnvl", + ) + return None + if mapping is not None and mapping.tp_size > 1: + if reason is None: + logger.info( + f"DFlash: decode-step target logits stay vocabulary-sharded ({lm_head.weight.shape[0]} " + f"columns per rank); trtllm::k3_spec_accept exchanges the row maxima, and a request's " + f"logits post-processors get them gathered" + ) + else: + logger.info_once( + f"DFlash: the target logits are all-gathered ({reason})", key=reason + ) + self._k3_shard_head, self._k3_shard_ok = lm_head, reason is None + if not self._k3_shard_ok: + return None + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + K = spec_metadata.runtime_draft_len + if hidden_states.shape[0] != num_gens * (K + 1) or not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return None + embed = draft_model.draft_model_full.model.embed_tokens + columns = lm_head.weight.shape[0] + if not accept_op.supports( + columns, num_gens, self._compute_block_size, K, embed.weight.shape[1], mapping.tp_size + ): + return None + if torch.cuda.is_current_stream_capturing(): + ws = accept_op.existing_workspace(mapping) + if ws is None: + return None + else: + ws = accept_op.workspace(mapping) + return ws, mapping.tp_rank * columns + + def _k3_logits_shard_reason(self, lm_head): + """Why the target logits cannot stay vocabulary-sharded under this configuration and head (None if they can, + given an MNNVL workspace for the mapping, which ``_k3_logits_shard`` checks).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + if mapping is None or mapping.tp_size <= 1: + return "no TP" + if mapping.enable_attention_dp or mapping.has_cp() or mapping.has_pp(): + return "attention DP, CP or PP" + if ( + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or getattr(lm_head, "gather_output_sizes", None) is not None + ): + return "the lm_head is not an evenly split, unpadded column-parallel head" + if ( + lm_head.bias is not None + or lm_head.has_any_quant + or lm_head.weight.dtype != torch.bfloat16 + ): + return "the lm_head has a bias, is quantized or is not bf16" + if not accept_op.supports_columns(lm_head.weight.shape[0], mapping.tp_size): + columns = lm_head.weight.shape[0] + return f"trtllm::k3_spec_accept cannot split {columns} columns over {mapping.tp_size} ranks" + return None + + def _k3_static_arange(self, lo: int, hi: int) -> torch.Tensor: + """``arange(lo, hi)`` (int64), built once outside CUDA-graph capture.""" + cache = self.__dict__.setdefault("_k3_aranges", {}) + t = cache.get((lo, hi)) + if t is None: + t = torch.arange(lo, hi, dtype=torch.long, device="cuda") + if not torch.cuda.is_current_stream_capturing(): + cache[(lo, hi)] = t + return t + + def _k3_mask_row(self, draft_model) -> torch.Tensor: + """The noise block's mask embedding, bf16 [hidden], cached: the drafter's own trained mask row if it kept one + for the mask id in use, else the embedding's row (``dflash_noise_block_embedding`` reads the same row).""" + row = getattr(self, "_k3_mask_embedding", None) + if row is None: + embed = draft_model.draft_model_full.model.embed_tokens + mask_token_id = self._mask_token_id(draft_model) + trained = self._trained_mask_embedding(draft_model, mask_token_id) + if trained is not None: + row = trained.to(embed.weight.dtype).reshape(-1).contiguous() + else: + row = embed.weight[mask_token_id].contiguous() + if torch.cuda.is_current_stream_capturing(): + return row + self._k3_mask_embedding = row + return row + + def _k3_ctx_pool_key(self): + """The identity of the pool the ctx cache is bound to now (KV cache estimation rebinds it to a new one).""" + return ( + tuple(t.data_ptr() for t in self._ctx_kv_buf) if self._ctx_kv_buf is not None else None + ) + + def _k3_ctx_kv_applies(self, draft_model, projected: torch.Tensor) -> bool: + """Whether ``trtllm::k3_ctx_kv`` writes this step's context K/V: the manager-bound paged pool with K and V, + bf16, the fused K/V weight without bias or context input norm, k_norm, NeoX RoPE from flashinfer's fp32 + cache, a context the op can split (up to 64 tokens: B <= 8 requests of K + 1 <= 8, ``ctx_op.pick_split``), one + allocation for every layer's pool. Checked once per token count and + bound pool; the pool view is kept for the bound pool only, so a rebind never leaves the kernel writing the + replaced pool.""" + pool_key = self._k3_ctx_pool_key() + if getattr(self, "_k3_ctx_pool_bound", None) != pool_key: + self._k3_ctx_pool_bound = pool_key + self._k3_ctx_pool = None + self._k3_ctx_kv_ok = {} + cache = self._k3_ctx_kv_ok + n = projected.shape[0] + ok = cache.get(n) + if ok is not None: + return ok + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctx_kv import op as ctx_op + from tensorrt_llm._torch.models import modeling_dflash + + reason = None + if not (self._ctx_paged and self._ctx_block_tables is not None): + reason = "the context pool is not the manager-bound paged pool" + elif projected.dtype != torch.bfloat16 or self._ctx_kv_buf[0].dtype != torch.bfloat16: + reason = "not bf16" + else: + if draft_model._fused_kv_weight is None: + draft_model._build_fused_kv_buffers() + view = ctx_op.pool_view(list(self._ctx_kv_buf)) + if ( + draft_model._fused_kv_bias is not None + or getattr(draft_model, "_input_ln_eps", None) is not None + ): + reason = "K/V bias or context input norm" + elif draft_model._k_norm_stacked is None or not draft_model._is_neox: + reason = "no k_norm or not NeoX RoPE" + elif ( + modeling_dflash._flashinfer_rope is None + or draft_model._get_cos_sin_cache().shape[1] != 64 + ): + reason = "RoPE is not flashinfer's full-rotary fp32 cache" + elif view is None or self._ctx_kv_buf[0].size(1) != 2: + reason = "the layers' pools are not K/V views of one allocation" + elif not ctx_op.pick_split( + draft_model._fused_kv_weight.shape[0], + projected.shape[1], + draft_model._num_kv_heads, + self._draft_tokens_per_req, + n, + projected.device, + ): + reason = f"{n} context tokens (at most 64) or the K/V weight shape" + else: + self._k3_ctx_pool = view + ok = reason is None + if not ok: + logger.info(f"DFlash: context K/V on the Python path ({reason})") + if not torch.cuda.is_current_stream_capturing(): + cache[n] = ok + return ok + + def _k3_ctx_kv( + self, draft_model, projected, ctx_positions, num_accepted, slots, rows + ) -> torch.Tensor: + """``trtllm::k3_ctx_kv``: every drafter layer's context K/V of this step into the paged pool and + ``_ctx_len += num_accepted`` (clamped); returns the context length each request may advertise.""" + if self._k3_ctx_pool is None or self._k3_ctx_pool_bound != self._k3_ctx_pool_key(): + raise RuntimeError("k3_ctx_kv: the context pool was rebound after the path was chosen") + flat, layer_off, page_stride, kv_stride, head_stride = self._k3_ctx_pool + return torch.ops.trtllm.k3_ctx_kv( + projected.contiguous(), + draft_model._fused_kv_weight, + draft_model._k_norm_stacked, + draft_model._get_cos_sin_cache(), + ctx_positions, + num_accepted, + self._ctx_len, + slots, + rows, + self._ctx_block_tables, + self._ctx_block_counts, + flat, + layer_off, + page_stride, + kv_stride, + head_stride, + draft_model._k_norm_eps, + self._max_ctx, + self._ctx_page_size, + self._compute_block_size, + draft_model._num_kv_heads, + ) + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Called with this step's acceptance when ``k3_spec_accept`` produced it (a hook for drafter families).""" + + def _k3_accept(self, logits, attn_metadata, spec_metadata, draft_model, shard=None): + """``trtllm::k3_spec_accept`` for a step ``_k3_accept_applies`` accepted: the acceptance, the block table, + the KDA replay record, kv_lens + 1 and the drafter's inputs in one launch. ``shard``: the exchange + workspace and first column when ``logits`` are this rank's vocabulary shard (``_k3_logits_shard``). + Returns (accepted tokens, num accepted, {rewind, bonus, qpos, cpos, noise}).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + num_gens = attn_metadata.num_seqs + K = spec_metadata.runtime_draft_len + mgr = attn_metadata.kv_cache_manager + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + force = 0.0 if is_warmup else float(self.force_num_accepted_tokens) + mode, _, _ = accept_op.force_mode(force, K) + if mode == accept_op._kernel_module().FORCE_FRAC: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + else: + dummy_rng = getattr(self, "_k3_dummy_rng", None) + if dummy_rng is None: + dummy_rng = self._k3_dummy_rng = ( + torch.zeros(4, dtype=torch.float32, device=logits.device), + torch.zeros(2, dtype=torch.int64, device=logits.device), + ) + pool, counter = dummy_rng + embed = draft_model.draft_model_full.model.embed_tokens + accepted, num_acc, rewind, bonus, qpos, cpos, noise = torch.ops.trtllm.k3_spec_accept( + logits, + spec_metadata.draft_tokens.reshape(num_gens, K), + attn_metadata.draft_kv_cache_block_offsets, + self._ctx_pool_idx, + self._ctx_block_divisor, + self._ctx_block_counts, + self._ctx_block_tables, + mgr.prev_num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices, + mgr._dummy_request_mask, + attn_metadata.kv_lens_cuda, + self._batch_to_slot, + self._ctx_len, + self._max_ctx, + embed.weight, + self._k3_mask_row(draft_model), + pool, + counter, + force, + self._compute_block_size, + *self._k3_shard_args(shard), + ) + self._on_acceptance(accepted, num_acc, attn_metadata, spec_metadata) + return ( + accepted, + num_acc, + dict(rewind=rewind, bonus=bonus, qpos=qpos, cpos=cpos, noise=noise), + ) + + @staticmethod + def _k3_shard_args(shard) -> tuple: + """``trtllm::k3_spec_accept``'s exchange arguments for a vocabulary-sharded step (none otherwise).""" + if shard is None: + return () + ws, first_column = shard + return ( + ws["uc"], + ws["mc"], + ws["flags"], + ws["rank"], + ws["slots"], + ws["push_copies"], + first_column, + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + # Set when target_logits kept this step's logits vocabulary-sharded (only k3_spec_accept reads them). + step_shard, self._k3_step_shard = getattr(self, "_k3_step_shard", None), None + shard = None + if step_shard is not None: + if step_shard[0] is not logits: + raise RuntimeError( + "DFlash: the target logits are not the vocabulary shard target_logits returned" + ) + shard = step_shard[1] + + raw_logits = logits + K = spec_metadata.runtime_draft_len + + if K == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + # Lazy init buffers and attach worker reference for prepare() + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + # A decode step of the supported mode runs the acceptance and the drafter's inputs as one kernel + # (trtllm::k3_spec_accept), which also decodes the block table below. + if shard is None: + fused_accept = self._k3_accept_applies( + logits, attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ) + elif self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + fused_accept = True + else: + raise RuntimeError( + "DFlash: the target logits are vocabulary-sharded, but the fused acceptance no longer applies" + ) + # Before any store: prefill and decode both address pages through it. + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + if not fused_accept: + self._refresh_ctx_block_tables(attn_metadata, batch_size) + + # Save context lengths so both warmup and a failed forward can roll + # back the in-place _ctx_len updates made during drafting. + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # Never allocate the snapshot while capturing a CUDA graph: the + # clone would live in the graph memory pool, and its replay-time + # writes could alias blocks reused by later captures. Rollback is + # only meaningful for eager/warmup forwards anyway; a failure + # during capture aborts the graph itself, and captured ops do not + # mutate _ctx_len until replay. + self._saved_ctx_len = self._ctx_len.clone() + self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_req_ctx_pos = dict(self._req_ctx_pos) + self._ctx_len_restore_pending = True + + self._execute_guided_decoder_if_present(logits) + + prep = None + if fused_accept: + # Saved before the kernel updates kv_lens (warmup also snapshots kv_lens_cuda here). + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + accepted_tokens, num_accepted_tokens, prep = self._k3_accept( + logits, attn_metadata, spec_metadata, draft_model, shard + ) + else: + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + # Opt-in acceptance recording (env-gated; eager-mode measurement + # runs only). Skipped for CUDA-graph batches (capture/replay/warmup + # use synthetic requests and forbid the host sync). + if ( + self._accept_stats is not None + and num_gens > 0 + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + ): + self._accept_stats.on_accept( + spec_metadata.request_ids[num_contexts:batch_size], + num_accepted_tokens[num_contexts:batch_size].tolist(), + ) + + if fused_accept: + # The kernel recorded the KDA replay acceptance and advanced kv_lens_cuda. + self._kv_rewind_amount = prep["rewind"] + self._kv_rewind_nc = num_contexts + self._kv_rewind_bs = batch_size + self._kv_rewind_pending = True + attn_metadata.update_for_spec_dec() + else: + # Update GDN/Mamba recurrent states to the accepted token's state. + if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) + + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] + + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + ) + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + self._rebuild_batch_to_slot(spec_metadata.request_ids) + + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + prep=prep, + ) + + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + ctx_kv_cache=inputs["ctx_kv_cache"], + ctx_page_table=inputs["ctx_page_table"], + **dflash_ctx_rows_kwargs(draft_model, inputs["ctx_rows_start"]), + ) + + # Gather K logits per gen request from the block outputs. + # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. + block_size = self._compute_block_size + gen_hidden_states = self._draft_block_hidden_states( + draft_model, hidden_states_out, num_gens, block_size, K + ) + gen_logits = self._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) + + # DFlash 2: replace the independent per-position picks with one + # coherent path through the block (absent for plain DFlash). + if getattr(draft_model, "has_candidate_selector", False): + gen_logits = self._apply_dflash2_selector( + draft_model, + gen_logits, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + spec_metadata, + ) + vocab_size = gen_logits.shape[-1] + + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + ) + + # Opt-in confidence-calibration recording (env-gated). The + # confidence values come from an optional provider on the + # recorder (None here -> no calibration rows collected; the + # DSpark confidence-scheduled verification MR supplies the + # real provider, see accept_stats.py). Guarded off for + # CUDA-graph batches and d2t vocab-mapped drafters. + if ( + self._accept_stats is not None + and self._accept_stats.confidence_provider is not None + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + and self._d2t is None + ): + self._accept_stats.record_draft_confidence( + spec_metadata.request_ids[num_contexts:batch_size], + draft_model, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + gen_draft_tokens, + ) + + else: + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + gen_vocab = vocab_size if num_gens > 0 else None + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + self._rollback_guided_decoder_after_verify(num_accepted_tokens) + + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + # Restore context lengths after warmup; real runs keep the updates. + if is_warmup: + self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_host() + self._ctx_len_restore_pending = False + + outputs = { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + if shard is not None: + # "logits" is this rank's vocabulary shard: the engine gathers it for any logits post-processor. + outputs["logits_vocab_shard"] = True + return outputs + + def _draft_block_hidden_states( + self, + draft_model, + hidden_states_out: torch.Tensor, + num_gens: int, + block_size: int, + num_draft_tokens: int, + ) -> torch.Tensor: + """The block-output rows that produce the K draft logits per gen request (``_draft_slot_ids``). + + The slot ids depend only on the shapes, so they are built once per shape, outside CUDA-graph + capture, and cached. When they are one contiguous run of rows (a single request, or the + shift_label convention with a block of K slots, whose requests' slots follow each other) the + rows are returned as a view, without a gather; otherwise they are gathered. + """ + rows = hidden_states_out.shape[0] + key = (num_gens, block_size, num_draft_tokens, rows) + cache = self.__dict__.setdefault("_draft_block_rows", {}) + entry = cache.get(key) + if entry is None: + ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. + ids = ids.clamp(max=rows - 1) + if torch.cuda.is_current_stream_capturing(): + return hidden_states_out[ids] + host = ids.tolist() + if host and host == list(range(host[0], host[0] + len(host))): + entry = host[0] + else: + entry = ids + cache[key] = entry + if isinstance(entry, int): + return hidden_states_out[entry : entry + num_gens * num_draft_tokens] + return hidden_states_out[entry] + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata, + ) -> torch.Tensor: + """fp32 logits of the block positions over the full vocabulary (the TP lm_head all-gathers its + shards). A drafter family whose refinement and draft sampler reduce across ranks themselves may + return this rank's vocabulary shard instead.""" + return draft_model.logits_processor( + gen_hidden_states, draft_model.lm_head, attn_metadata, True + ) + + def _ctx_rows_start(self, num_contexts: int, num_gens: int) -> Optional[int]: + """The first row of the gen requests' page-table rows where those rows are one contiguous run, else None. + + The manager's block table is keyed by batch position, so the gen requests' rows are [num_contexts, + num_contexts + num_gens): a drafter can view them instead of gathering ``ctx_cache_batch_idx``'s rows. The + private arena's table is keyed by slot, which need not be contiguous. + """ + if num_gens > 0 and self._ctx_block_tables is not None: + return num_contexts + return None + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DFlashSpecMetadata, + draft_model: nn.Module, + total_target_tokens: int = 0, + prep: Optional[dict] = None, + ): + """Prepare inputs for DFlash's draft forward. + + ``prep``: ``trtllm::k3_spec_accept``'s bonus tokens, positions and noise embedding for a decode step it + took (see ``_k3_accept``); computed here otherwise. + + For gen requests, builds: + - noise_embedding: token embeddings for [accepted + mask] tokens + - query_positions: position IDs for the query tokens + - num_ctx_per_req: per-request context length in the pool + - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed + views of the persistent per-layer K/V pool. + - ctx_rows_start: where ctx_cache_batch_idx is the contiguous rows + [ctx_rows_start, ctx_rows_start + num_gens) of ctx_page_table, their + start (``_ctx_rows_start``), else None. + """ + num_contexts = attn_metadata.num_contexts + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + + mask_token_id = self._mask_token_id(draft_model) + + # Get the embed_tokens layer from the draft model + embed_tokens = draft_model.draft_model_full.model.embed_tokens + hidden_dim = ( + spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] + ) + + if num_gens > 0: + gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] + gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] + + total_tokens_per_req = self._draft_tokens_per_req # K+1 + K = spec_metadata.runtime_draft_len + + # Get captured multi-layer hidden states from spec_metadata + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + has_target_features = ( + captured_hs is not None + and hasattr(draft_model, "fc") + and hasattr(draft_model, "hidden_norm") + ) + + # Use cached block_size (resolved once on first call) + block_size = self._compute_block_size + query_tokens_per_req = block_size + + # Get slots for gen requests from pre-computed mapping + slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] + K_plus_1 = K + 1 + + if prep is not None: + gen_rows_out = self._k3_static_arange(num_contexts, num_contexts + num_gens) + offsets_kp1 = self._k3_static_arange(0, K_plus_1) + bonus = prep["bonus"] + query_position_ids = prep["qpos"] + ctx_position_ids = prep["cpos"] + noise_embed_2d = prep["noise"] + else: + gen_rows_out = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + + bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) + bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() + + ctx_len_gen = self._ctx_len[slots] + j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") + offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") + + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, + bonus, + mask_token_id, + query_tokens_per_req, + self._trained_mask_embedding(draft_model, mask_token_id), + ) + + # Accumulate new accepted features into context buffers + fused_ctx_kv = False + if has_target_features: + gen_start = attn_metadata.num_ctx_tokens + # Target now processes exactly K+1 tokens per gen req, so the + # captured slice is already the full set we need to project. + gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] + gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) + projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) + fused_ctx_kv = prep is not None and self._k3_ctx_kv_applies( + draft_model, projected_to_store + ) + if fused_ctx_kv: + num_ctx_fused = self._k3_ctx_kv( + draft_model, + projected_to_store, + ctx_position_ids, + gen_num_accepted, + slots, + gen_rows_out, + ) + # k3_ctx_kv adds the accepted tokens in the kernel; empty the dummy slot after + # it, as _advance_ctx_len does after the add on the torch path. + self._ctx_len[self._dummy_slot].zero_() + if has_target_features and not fused_ctx_kv: + gen_num_accepted_long = gen_num_accepted.long() + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) + if self._ctx_block_tables is not None: + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) + else: + col_idx = col_idx.clamp(max=self._max_ctx - 1) + + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) + col_flat = col_idx.reshape(-1) + proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] + pos_flat = ctx_position_ids.long().reshape(-1) + mask_1d = write_mask.reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + k_new, v_new = draft_model.precompute_context_kv( + proj_flat.to(cache_dtype), pos_flat + ) + mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + slot_long = slot_flat.long() + col_long = col_flat.long() + if self._ctx_paged: + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + self._advance_ctx_len(slots, gen_num_accepted_long) + + if fused_ctx_kv: + num_ctx_per_req_t = num_ctx_fused + else: + num_ctx_per_req_t = self._ctx_len[slots] + if not fused_ctx_kv and self._ctx_block_tables is not None: + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + noise_embedding = noise_embed_2d + query_positions = query_position_ids.long() + + # Update seq_lens for gen requests to K+1 (the number of tokens + # the target forward actually processed). + attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( + total_tokens_per_req + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req + else: + noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) + query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") + num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") + slots = torch.empty(0, dtype=torch.long, device="cuda") + gen_rows_out = slots + bonus = torch.empty(0, dtype=torch.long, device="cuda") + + return { + "noise_embedding": noise_embedding, + "query_positions": query_positions, + "num_ctx_per_req": num_ctx_per_req_t, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": bonus, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), + "ctx_rows_start": self._ctx_rows_start(num_contexts, num_gens), + } + + +class KimiK3DSparkWorker(KimiK3DFlashWorker, DSparkWorker): + # Whether this step's block logits are the Kimi K3 decode path's vocabulary shard (see ``_draft_block_logits``). + _k3_sharded_block_logits = False + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Block logits the Kimi K3 decode path kept vocab-sharded (see ``_draft_block_logits``) run the + whole chain as ``trtllm::k3_markov``, which also returns the greedy draft tokens and + next_new_tokens. Other TP-sharded logits slice markov_w2's rows to this rank's contiguous + shard and chain through the TP-aware global argmax. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # Unlike the d2t guard this one cannot move to set_draft_model: it + # keys on the runtime logits width, and reproducing that at init would + # duplicate the draft head's sharding rules. A standalone drafter + # borrows the target lm_head, whose gather_output defaults to True, so + # the logits normally arrive full-vocab and this branch is skipped. + self._k3_markov = None + k3_sharded, self._k3_sharded_block_logits = self._k3_sharded_block_logits, False + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + if k3_sharded: + return self._k3_markov_chain( + draft_model, gen_logits, first_prev_tokens, vocab_slice + ) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) + + def _keep_draft_logits_sharded(self, draft_model, spec_metadata, num_gens: int) -> bool: + """Keep the draft logits vocab-sharded: ``trtllm::k3_markov`` reduces across the ranks itself. + + The chain's global argmaxes are the greedy draft tokens, so the lm_head's all-gather of the + full logits and the draft sampler's gather are both skipped. Needs a Kimi K3 target running its + decode kernels (``k3_decode``), plain TP, greedy drafting (rejection sampling reads + full-vocabulary probabilities), an unquantized bias-free column-parallel bf16 lm_head whose + shards tile the Markov vocabulary, bf16 Markov weights of the kernel's rank, MNNVL (the model's + all-reduces own an MNNVL workspace for this mapping), and a block, shard and batch of + ``num_gens`` requests the kernel can split. + """ + if not self.k3_decode: + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + lm_head = getattr(draft_model, "lm_head", None) + if ( + lm_head is None + or not getattr(draft_model, "has_markov_head", False) + or mapping is None + or mapping.tp_size <= 1 + or mapping.enable_attention_dp + or spec_metadata.wants_advanced_draft_sampling + or getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "bias", None) is not None + or lm_head.weight.dtype != torch.bfloat16 + or draft_model.markov_w1.dtype != torch.bfloat16 + or draft_model.markov_w2.dtype != torch.bfloat16 + or lm_head.weight.shape[0] * mapping.tp_size != draft_model.markov_w2.shape[0] + or MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None + ): + return False + shard = lm_head.weight.shape[0] + block = spec_metadata.runtime_draft_len + rank = k3_markov_op._kernel_module().MARKOV_RANK + if ( + not 0 < block <= k3_markov_op.WORKSPACE_MAX_BLOCK + or draft_model.markov_w1.dim() != 2 + or draft_model.markov_w1.shape[1] != rank + or tuple(draft_model.markov_w2.shape[1:]) != (rank,) + ): + return False + if k3_markov_op.pick_grid(shard, block, num_gens) == 0: + logger.warning_once( + f"DSpark Markov head: trtllm::k3_markov cannot split a {shard}-row vocab shard for " + f"{num_gens} requests; the draft logits are all-gathered and the chain runs unfused.", + key=f"dspark_k3_markov_shard_{num_gens}", + ) + return False + return True + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + ) -> torch.Tensor: + """This rank's bf16 shard of the block logits when ``k3_markov`` takes them (it converts them to fp32 + exactly, as ``.float()`` would), from the drafter's logits processor's head kernel where it takes the rows, + else the head's own GEMM (``_k3_head_shard``); otherwise the base class' fp32 logits.""" + num_gens = gen_hidden_states.shape[0] // max(spec_metadata.runtime_draft_len, 1) + self._k3_sharded_block_logits = self._keep_draft_logits_sharded( + draft_model, spec_metadata, num_gens + ) + if self._k3_sharded_block_logits: + return self._k3_head_shard( + draft_model.logits_processor, draft_model.lm_head, gen_hidden_states + ) + return super()._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): + """The base acceptance; under ``k3_decode`` its outputs are kept for ``k3_markov``'s next_new_tokens.""" + accepted_tokens, num_accepted_tokens = super().sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + if self.k3_decode: + self._on_acceptance(accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata) + return accepted_tokens, num_accepted_tokens + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Keeps this step's acceptance for ``k3_markov``'s next_new_tokens (and the metadata whose KV lengths it + rewinds).""" + self._k3_acceptance = ( + accepted_tokens, num_accepted_tokens, attn_metadata.num_contexts, spec_metadata, attn_metadata + ) # fmt: skip + + def _k3_markov_chain( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + vocab_slice: slice, + ) -> torch.Tensor: + """The Markov chain as ``trtllm::k3_markov`` on this rank's shard of the block logits. + + Returns the corrected logits (fp32); keeps the kernel's greedy tokens for + ``sample_draft_tokens`` and its next_new_tokens (built from this step's acceptance) for + ``_prepare_next_new_tokens``. + """ + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + + acceptance = getattr(self, "_k3_acceptance", None) + if acceptance is None: + raise RuntimeError("DSpark Markov head: k3_markov needs this step's acceptance first") + accepted_tokens, num_accepted_tokens, num_contexts, spec_metadata, attn_metadata = ( + acceptance + ) + num_gens = gen_logits.shape[0] + # The draft forward's pending KV-length rewind (see _apply_kv_rewind_after_draft) runs in the kernel: every + # reader of kv_lens_cuda in the draft forward precedes it. In warmup the restore of kv_lens_cuda that + # follows overwrites it, as it would the Python rewind. + rewind = getattr(self, "_kv_rewind_amount", None) + fold = ( + getattr(self, "_kv_rewind_pending", False) + and rewind is not None + and getattr(attn_metadata, "kv_lens_cuda", None) is not None + and rewind.dtype == torch.int32 + and self._kv_rewind_bs - self._kv_rewind_nc == rewind.numel() <= num_gens + ) + corrected, tokens, next_new = k3_markov_op.markov_chain( + self.mapping, + gen_logits.contiguous(), + first_prev_tokens.long(), + draft_model.markov_w1, + draft_model.markov_w2[vocab_slice], + vocab_slice.start, + accepted_tokens, + num_accepted_tokens[num_contexts : num_contexts + num_gens], + spec_metadata.batch_indices_cuda[num_contexts : num_contexts + num_gens], + kv_lens=attn_metadata.kv_lens_cuda if fold else None, + rewind=rewind if fold else None, + rewind_first=self._kv_rewind_nc if fold else 0, + ) + if fold: + self._kv_rewind_amount = None + self._kv_rewind_pending = False + self._k3_markov = (corrected, tokens, next_new, accepted_tokens, num_accepted_tokens) + return corrected + + def sample_draft_tokens( + self, + logits, + spec_metadata, + batch_size, + *, + num_contexts=0, + draft_step=None, + mapping_lm_head_tp=None, + ): + """Greedy block drafts straight from ``k3_markov``. + + When ``logits`` are the corrected logits the kernel just returned, its per-position global + argmax (first maximum, lowest vocabulary index among equal values) is the token the + TP-gathered greedy sampler would pick. Anything else goes to the base sampler. + """ + chain = getattr(self, "_k3_markov", None) + if ( + chain is not None + and chain[0] is logits + and mapping_lm_head_tp is None + and not spec_metadata.wants_advanced_draft_sampling + ): + self._k3_markov_next = (chain[1], chain[2], chain[3], chain[4]) + return chain[1] + self._k3_markov_next = None + return super().sample_draft_tokens( + logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + draft_step=draft_step, + mapping_lm_head_tp=mapping_lm_head_tp, + ) + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """``k3_markov``'s next_new_tokens when the drafts are its tokens for the whole batch (no context + requests); otherwise the base assembly.""" + chain = getattr(self, "_k3_markov_next", None) + # The step ends here: release its acceptance (which holds the step's attention and spec metadata) and the + # kernel's outputs, so they do not keep a captured graph's pool or metadata alive after the step. + self._k3_markov_next = None + self._k3_acceptance = None + self._k3_markov = None + if ( + chain is not None + and chain[0] is next_draft_tokens + and chain[2] is accepted_tokens + and chain[3] is num_accepted_tokens + and next_draft_tokens.shape[0] == batch_size + ): + return chain[1] + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py index 24233d194709..66a1f4f72049 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/weights.py @@ -1,6 +1,44 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The shared Kimi K3 text checkpoint loader.""" +# >>> route B: this target's expert slice +"""Weight loading: Kimi K3 (MXFP4) / sm_100 / tp16_moetp16ep1. -from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD -from ..weights import load as load +The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and +`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: + +* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the + routed experts one at a time, keeps this rank's slice (a sixteenth of the width of every expert under moe_tp 16 x + moe_ep 1: 192 of 3072, zero-padded to 256 for the MoE kernels' tiles), and checks its own key coverage. +* The vision tower and the projector are a predicted non-load: listed here, never read. +* Any other key fails the load, naming it, rather than being dropped. +""" +# <<< route B + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_utils import filter_weights + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") + + +def load(model, weights) -> None: + """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) + assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = _LANG_PREFIX + KimiLinearForCausalLM.load_weights(model, lm_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py index e1e4bb7413de..474e865afce4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 MXFP4 / sm_100 / tp16_moetp4ep4.""" +"""Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py new file mode 100644 index 000000000000..f5327daf9474 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_comm.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's residual-update collectives on the catalog's Kimi K3 MNNVL and sandwich entries. + +A layer's post-attention step adds the attention output to the running prefix sum, selects the attention residual +and applies the post-attention RMSNorm. Under TP the attention output is a sum over the TP group, so the step is that +all-reduce followed by the residual update. On a step of at most `AR_ATTN_RES_MAX_TOKENS` tokens (wide decode steps +aside) one collective runs both: + +* `K3DecodeComm.sandwich_oproj`, `comm/k3_sandwich_oproj`: o_proj, its all-reduce and the residual update in one + kernel, at most `SANDWICH_MAX_TOKENS` tokens of an o_proj of the TP16 per-rank shape [7168, 768]. The attention + hands over its gated o_proj input. +* `K3DecodeComm.allreduce_attn_res`, `comm/mnnvl_allreduce_attn_res`, everywhere else: the one-shot MNNVL all-reduce + of the attention's unreduced o_proj partial, with the residual update as its epilogue. + +The pre-attention step after a MoE layer whose row-parallel tail was handed on (a `PendingTail`, `decode_moe.py`) is +the same kind of collective: `K3DecodeComm.sandwich_tail`, `comm/k3_sandwich_tail`, runs the tail GEMV, its +all-reduce and the next layer's residual update (the final norm's, after the last layer) in one kernel. + +The DSpark drafter (`K3DSparkDrafter`) runs its collectives with a plain residual add and RMSNorm on the same state: + +* `K3DecodeComm.sandwich_plain`, `comm/k3_sandwich_plain`: a row-parallel projection of at most `SANDWICH_MAX_TOKENS` + tokens (a layer's attention output projection, or its MLP's SiLU-and-mul and down projection), its all-reduce, the + residual add and the RMSNorm in one kernel. +* `K3DecodeComm.allreduce_norm`, `comm/mnnvl_fusion_allreduce`: the all-reduce of an unreduced projection output of + any token count the MNNVL workspace holds, with the residual add and the RMSNorm in its epilogue (the context + projection's with a zero residual). + +The state is one `MnnvlWorkspace` and one `K3SandwichWorkspace` of the TP group (`K3DecodeComm.create`): collective +over the group and eager, built by the target in `post_load_weights` before any CUDA-graph capture. Every rank must +make the same calls on each in the same order. Which call a step takes is decided from its token count and kind and +from the load-time layout alone, which every rank of the group shares. + +The plain MNNVL all-reduces the decode path keeps (the stock modules') send one-shot up to +`DECODE_AR_ONE_SHOT_MAX_BYTES` (`use_decode_one_shot`), except a wide decode step's (`wide_all_reduce`). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_oproj import ( + K3SandwichWorkspace, + k3_sandwich_oproj, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_plain import ( + k3_sandwich_plain, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_tail import ( + k3_sandwich_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allreduce_attn_res import ( + mnnvl_allreduce_attn_res, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_fusion_allreduce import ( + mnnvl_fusion_allreduce, + required_buffer_bytes, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) + +# The kernel's support predicate: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_sandwich import op as _sandwich_op +from tensorrt_llm._torch.distributed import AllReduceParams + +# The most tokens of a step whose post-attention all-reduce runs one-shot with the residual update in its epilogue. +AR_ATTN_RES_MAX_TOKENS = 16 + +# The most tokens the sandwich kernel takes (one token tile). +SANDWICH_MAX_TOKENS = _sandwich_op.MAX_TOKENS + +# The MNNVL workspace's buffer size: a one-shot call pushes T x 7168 bf16 from each of 16 ranks, 3.5 MiB at T = 16. +MNNVL_BUFFER_BYTES = 4 << 20 + +# The one-shot ceiling of the stock MNNVL all-reduces on the decode path: 8 tokens x 7168 x 16 ranks x 2 B is +# 1.75 MiB, which the stock 1 MiB ceiling would send two-shot. +DECODE_AR_ONE_SHOT_MAX_BYTES = 4 << 20 + +# A wide decode step's ceiling: the stock 1 MiB. At 16 ranks two-shot is faster for every wide step's rows. +WIDE_AR_ONE_SHOT_MAX_BYTES = 1 << 20 + + +def use_decode_one_shot(model: nn.Module) -> None: + """Every stock MNNVL all-reduce of ``model`` sends one-shot up to `DECODE_AR_ONE_SHOT_MAX_BYTES`. Each grows its + workspace on its first eager call of a larger size, before the capture of that size.""" + for module in model.modules(): + mnnvl = getattr(module, "mnnvl_allreduce", None) + if mnnvl is not None: + mnnvl.one_shot_max_bytes = DECODE_AR_ONE_SHOT_MAX_BYTES + mnnvl.early_trigger = True + + +def skip_all_reduce() -> AllReduceParams: + """All-reduce parameters under which a row-parallel module returns this rank's unreduced output (its projection + without its all-reduce): a ``Linear``'s ``all_reduce_params``, a ``GatedMLP``'s ``final_all_reduce_params``.""" + return AllReduceParams(enable_allreduce=False) + + +def wide_all_reduce(all_reduce: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``all_reduce(x)`` of a wide decode step (a stock ``AllReduce`` module, no fusion): its MNNVL all-reduce with + the `WIDE_AR_ONE_SHOT_MAX_BYTES` ceiling, else the module itself.""" + mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None: + out = mnnvl( + x.contiguous(), AllReduceParams(), one_shot_max_bytes=WIDE_AR_ONE_SHOT_MAX_BYTES + ) + if out is not None: + return out + return all_reduce(x) + + +class PendingTail(NamedTuple): + """A MoE layer's row-parallel tail left to its consumer's fused pre-attention step (``K3DecodeComm.sandwich_tail``): + the reduced latent ``[T, 3584]``, the shared experts' activation, the tail weight ``[latent up columns | padding | + shared down]``, this rank's first latent column and the latent norm's epsilon.""" + + latent: torch.Tensor + act: torch.Tensor + weight: torch.Tensor + lo: int + lat_eps: float + + +def _eps(norm: nn.Module) -> float: + """The epsilon of a KimiK3RMSNorm or ``torch.nn.RMSNorm`` (``eps``; for torch's None, its default: the machine + epsilon of the weight's dtype) or a stock RMSNorm (``variance_epsilon``).""" + eps = norm.eps if hasattr(norm, "eps") else norm.variance_epsilon + return float(torch.finfo(norm.weight.dtype).eps if eps is None else eps) + + +def _res_args(res_proj: nn.Module, res_norm: nn.Module, out_norm: nn.Module) -> tuple: + """The residual update's weights and epsilons in the entries' order.""" + return ( + res_proj.weight.reshape(-1), + res_norm.weight, + out_norm.weight, + _eps(res_norm), + _eps(out_norm), + ) + + +@dataclass(eq=False) +class K3DecodeComm: + """The decode path's collective state for the TP group: the `MnnvlWorkspace` and the `K3SandwichWorkspace` the + post-attention steps run on. Built by `create`; owned by the target and shared by every layer.""" + + mnnvl: MnnvlWorkspace + sandwich: K3SandwichWorkspace + + @classmethod + def create(cls, mapping, oproj_weight: Optional[torch.Tensor] = None) -> "K3DecodeComm": + """The state for ``mapping``'s TP group. Collective: every rank of the group calls it at the same point, + eagerly, before any CUDA-graph capture; each workspace fails on every rank or on none. + + ``oproj_weight``: an o_proj weight ``takes_oproj`` holds for. The sandwich kernel compiles here for its + shape, with one call on a zero row of a zero weight, so no capture compiles it; the call advances the + sandwich workspace on every rank alike.""" + state = cls( + MnnvlWorkspace.create(mapping, MNNVL_BUFFER_BYTES), + K3SandwichWorkspace.create(mapping), + ) + if oproj_weight is not None: + weight = torch.zeros_like(oproj_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_oproj( + weight.new_zeros(1, weight.shape[1]), + weight, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + state.sandwich, + ) + torch.cuda.synchronize(weight.device) + return state + + def compile_tail(self, latent_size: int, act_size: int, tail_weight: torch.Tensor) -> None: + """Compile the sandwich tail kernel for a MoE tail of ``latent_size`` latent and ``act_size`` shared columns + and ``tail_weight``'s shape, with one call on a zero row of a zero weight, before any capture. Collective: every + rank of the group makes the call; it advances the sandwich workspace on every rank alike.""" + weight = torch.zeros_like(tail_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_tail( + weight.new_zeros(1, latent_size), + weight.new_zeros(1, act_size), + weight, + 0, + 1e-6, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + self.sandwich, + ) + torch.cuda.synchronize(weight.device) + + def takes_post_attention(self, hidden_states: torch.Tensor, step) -> bool: + """Whether one of this state's collectives runs the post-attention step of a layer whose attention input is + ``hidden_states``: at most `AR_ATTN_RES_MAX_TOKENS` bf16 rows of a hidden size the MNNVL entry takes, on any + step but a wide decode step (``step.wide``), whose post-attention update stays the all-reduce and the fused + add + attn_res + RMSNorm.""" + if step is not None and step.wide: + return False + if hidden_states.dim() != 2 or hidden_states.dtype != torch.bfloat16: + return False + rows, hidden = hidden_states.shape + return ( + hidden % 1024 == 0 + and hidden <= 8192 + and 0 < rows <= min(AR_ATTN_RES_MAX_TOKENS, self.mnnvl.max_one_shot_tokens(hidden)) + ) + + @staticmethod + def takes_oproj(o_proj: nn.Module, max_snapshots: int) -> bool: + """Whether ``sandwich_oproj`` takes the layer of output projection ``o_proj`` on a step of at most + `SANDWICH_MAX_TOKENS` tokens, decided once the weights are final: a bias-free, contiguous bf16 weight of the + kernel's shape (the TP16 per-rank [7168, 768]) and a snapshot bank of at most ``max_snapshots`` rows the + kernel's candidate count holds.""" + weight = getattr(o_proj, "weight", None) + return ( + getattr(o_proj, "bias", None) is None + and isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.is_cuda + and _sandwich_op.supports(weight.new_empty((1, weight.shape[1])), weight, max_snapshots) + ) + + @staticmethod + def takes_tail( + latent: torch.Tensor, act: torch.Tensor, tail_weight: torch.Tensor, max_snapshots: int + ) -> bool: + """Whether ``sandwich_tail`` takes a MoE tail of these tensors' shapes (rows of ``latent`` and ``act``, at most + `SANDWICH_MAX_TOKENS`; ``tail_weight``) with a snapshot bank of at most ``max_snapshots`` rows: the TP16 + per-rank shapes.""" + return _sandwich_op.supports_tail(latent, act, tail_weight, max_snapshots) + + def allreduce_attn_res( + self, + partial: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_allreduce_attn_res`` call: ``updated = prefix_sum + + allreduce(partial)`` (the sum alone without ``prefix_sum``), ``normed = out_norm(attn_res(block_residual..., + updated))``, the attention residual selected with ``res_proj`` and ``res_norm``. ``block_residual`` holds the + valid snapshots, ``[S, T, H]``.""" + return mnnvl_allreduce_attn_res( + partial.contiguous(), + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.mnnvl, + ) + + def sandwich_tail( + self, + pending: PendingTail, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + updated_out: Optional[torch.Tensor] = None, + tap: Optional[torch.Tensor] = None, + tap_updated: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of a MoE layer's row-parallel tail ``pending`` + (``[RMSNorm(latent)[:, lo:lo + 224] | act] @ weight.T``), in one ``comm/k3_sandwich_tail`` call. + ``updated_out``: a bf16 ``[T, H]`` tensor the call stores ``updated`` into (the consumer's snapshot bank row), + returned as ``updated``. ``tap``: a bf16 ``[T, H]`` view (a speculative capture slot) the call also stores + the pre-norm attention-residual mixture into, or ``updated`` with ``tap_updated``.""" + return k3_sandwich_tail( + pending.latent, + pending.act, + pending.weight, + pending.lo, + pending.lat_eps, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + tap=tap, + tap_updated=tap_updated, + updated_out=updated_out, + ) + + def sandwich_oproj( + self, + core: torch.Tensor, + o_weight: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of ``core @ o_weight.T``, in one ``comm/k3_sandwich_oproj`` + call: ``core`` is this rank's gated o_proj input ``[T, 768]``, ``o_weight`` its o_proj slice. Bit for bit + o_proj followed by ``allreduce_attn_res``.""" + return k3_sandwich_oproj( + core.contiguous(), + o_weight, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + ) + + def compile_plain(self, weight: torch.Tensor, swiglu: bool = False) -> bool: + """Compile the plain sandwich (``sandwich_plain``) for a projection of ``weight``'s shape (with ``swiglu``, a + down projection after the SiLU-and-mul) with one call on a zero row of a zero weight, before any capture; + False, with nothing launched, where the kernel does not take that shape. Collective: every rank of the group + makes the call; it advances the sandwich workspace on every rank alike.""" + zero = torch.zeros_like(weight) + hidden, k_in = zero.shape + x = zero.new_zeros(1, 2 * k_in if swiglu else k_in) + residual = zero.new_zeros(1, hidden) + ones = zero.new_ones(hidden) + if not _sandwich_op.supports_plain(x, zero, residual, ones, swiglu): + return False + k3_sandwich_plain(x, zero, residual, ones, 1e-6, self.sandwich, swiglu=swiglu) + torch.cuda.synchronize(zero.device) + return True + + @staticmethod + def takes_plain( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> bool: + """Whether ``sandwich_plain`` takes the call: at most `SANDWICH_MAX_TOKENS` contiguous bf16 rows ``x`` of a + row-parallel slice ``weight`` [7168, K] (K a multiple of 128 up to 896; with ``swiglu`` K 896 and ``x`` the + ``[gate | up]`` rows, 2 K wide), ``residual`` [rows, 7168], and ``norm``'s bf16 [7168] weight.""" + return _sandwich_op.supports_plain(x, weight, residual, norm.weight, swiglu) + + def sandwich_plain( + self, + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/k3_sandwich_plain`` call: ``updated = residual + allreduce(x @ + weight.T)`` (with ``swiglu``, of ``silu_and_mul(x) @ weight.T``), ``normed = norm(updated)`` for a plain + RMSNorm ``norm``. Bit for bit, by the kernel's statement, the projection on ``k3_ctm_gemv`` at split 1 (with + ``swiglu``, ``k3_ctm_gemv_swiglu`` at split 2) followed by ``allreduce_norm``'s one-shot call.""" + return k3_sandwich_plain( + x.contiguous(), + weight, + residual.contiguous(), + norm.weight, + _eps(norm), + self.sandwich, + swiglu=swiglu, + ) + + def takes_allreduce_norm(self, rows: int, hidden: int) -> bool: + """Whether ``allreduce_norm`` takes ``rows`` bf16 rows of ``hidden`` columns: the MNNVL workspace holds the + call at the decode path's one-shot ceiling (`DECODE_AR_ONE_SHOT_MAX_BYTES`, two-shot above it).""" + if rows <= 0 or hidden <= 0 or hidden % 8: + return False + world, buffer_bytes = self.mnnvl.world_size, self.mnnvl.buffer_bytes + need = required_buffer_bytes( + rows, hidden, world, torch.bfloat16, DECODE_AR_ONE_SHOT_MAX_BYTES + ) + two_shot = rows * hidden * world * 2 > DECODE_AR_ONE_SHOT_MAX_BYTES + return need <= buffer_bytes and not (two_shot and buffer_bytes % 32) + + def allreduce_norm( + self, partial: torch.Tensor, residual: torch.Tensor, norm: nn.Module + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_fusion_allreduce`` call: ``updated = residual + + allreduce(partial)``, ``normed = norm(updated)`` for a plain RMSNorm ``norm``, sent one-shot up to + `DECODE_AR_ONE_SHOT_MAX_BYTES`. ``partial`` is this rank's unreduced ``[rows, hidden]`` bf16 output of a + row-parallel projection, of a shape ``takes_allreduce_norm`` holds for.""" + return mnnvl_fusion_allreduce( + partial.contiguous(), + self.mnnvl, + DECODE_AR_ONE_SHOT_MAX_BYTES, + residual.contiguous(), + norm.weight, + _eps(norm), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py new file mode 100644 index 000000000000..301b709643b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_gemv.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. + +* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest + at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, + `gemm/k3_ctm_gemv_wide`, `gemm/k3_ctm_gemv_long`, `gemm/k3_ctm_gemv` or `gemm/k3_ctm_gemv_swiglu`, and, where the + site lists it, more rows (up to `WIDE_ROWS`) on `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused + projections (MLA's [W_a; W_g] with the gate rows through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention + output projection, the built-in MLA path's q_a / kv_a, q_b and gate projections, the MoE decode path's projections + (`decode_moe.py`), two of them with fp32 outputs, and the DSpark drafter's q / k / v, output, gate / up and down + projections (`K3DSparkDrafter`). +* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on + `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the + stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the + drafter's logits on the same head take it too. +* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual + bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. +* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: + gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then + runs the down projection's all-reduce. + +Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. + +A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. +`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The +embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call +whose kernel has not run eagerly is refused. + +The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one +stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, Set + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( + k3_ctm_gemv_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( + K3HeadGemvWorkspace, + k3_head_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm + +# The kernels' support predicates: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._utils import mpi_disabled + +# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 +# requests of 8 tokens). +MAX_ROWS = 8 +WIDE_ROWS = 64 + + +@dataclass(frozen=True) +class Site: + """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows + ("decode", "wide", "long", "ctm" for k3_ctm_gemv, or "swiglu" for k3_ctm_gemv_swiglu, whose rows are the + ``[gate | up]`` input of the SiLU-and-mul, 2 ``k`` wide), and k3_ctm_gemv_wide at MAX_ROWS+1..``wide_rows`` rows + where ``wide``. Output columns from ``sig_col0`` on are stored through a sigmoid; ``out_fp32``: an fp32 output + (k3_ctm_gemv_wide only). ``split`` / ``ring`` / ``push``: the CTM kernels' CTAs per 128-row weight tile, + k3_ctm_gemv_long's weight-ring stages, and whether the partial sums are pushed to each row's owner.""" + + n: int + k: int + small: str + wide: bool = False + sig_col0: int = -1 + split: int = 0 + ring: int = 0 + push: bool = False + out_fp32: bool = False + wide_rows: int = WIDE_ROWS + + +SITES: Dict[str, Site] = { + # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), + # the gate rows through a sigmoid. + "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), + # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). + "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), + # The attention output projection (row parallel; 6 heads x 128 in). + "o_proj": Site(7168, 768, "decode", wide=True), + # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. + "kv_a": Site(2112, 7168, "decode"), + "q_b": Site(1152, 1536, "wide"), + "g_proj": Site(768, 7168, "wide"), + # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. + "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), + "dense_down": Site(7168, 2112, "long", split=2, ring=6), + # The MoE decode path (decode_moe.py): this rank's head slice [latent down 224 | router 56] with an fp32 output, + # the shared experts' gate_up, the row-parallel tail [latent up 224 | padding 32 | shared down 384] and the + # replicated tail's latent up projection with an fp32 output. The wide kernel takes the head slice and the tail up + # to 32 rows, where it is faster than the stock GEMM. + "moe_head": Site(280, 7168, "wide", wide=True, out_fp32=True, wide_rows=32), + "moe_shared_gate_up": Site(768, 7168, "wide", wide=True), + "moe_tail": Site(7168, 640, "wide", wide=True, wide_rows=32), + "moe_up": Site(7168, 3584, "wide", out_fp32=True), + # The DSpark drafter's projections (K3DSparkDrafter; 6 query heads, 1 KV head of 64, MLP width 896): the fused + # q / k / v, the output projection (row parallel), gate_up [gate 896 | up 896], and the down projection with the + # SiLU-and-mul (row parallel). + "drafter_qkv": Site(512, 7168, "long", split=8, ring=6, push=True), + "drafter_o": Site(7168, 384, "ctm", split=1, push=True), + "drafter_gate_up": Site(1792, 7168, "long", split=8, ring=6, push=True), + "drafter_down": Site(7168, 896, "swiglu", split=2, push=True), +} + + +def _width(spec: Site) -> int: + """The row width a site's kernel reads: ``k``, or the ``[gate | up]`` input of a SiLU-and-mul site, 2 ``k``.""" + return 2 * spec.k if spec.small == "swiglu" else spec.k + + +def _wide_tile(rows: int) -> int: + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel + + return k3_ctm_gemv_kernel.wide_tile(rows) + + +def _run( + spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor +) -> Optional[torch.Tensor]: + """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" + if kernel == "decode": + if spec.sig_col0 >= 0 or spec.out_fp32 or not _decode_op.supports(x2d, weight): + return None + return k3_decode_gemv(x2d, weight) + if kernel == "wide": + if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, spec.out_fp32): + return None + return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0, out_fp32=spec.out_fp32) + if spec.out_fp32: + return None + if kernel == "ctm": + if spec.sig_col0 >= 0 or not _ctm_op.supports(x2d, weight, spec.split): + return None + return k3_ctm_gemv(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + if kernel == "swiglu": + if spec.sig_col0 >= 0 or not _ctm_op.supports_swiglu(x2d, weight, spec.split): + return None + return k3_ctm_gemv_swiglu(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. + sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count + if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( + x2d, weight, spec.split, spec.ring + ): + return None + return k3_ctm_gemv_long( + x2d, + weight, + sig_col0=spec.sig_col0, + split=spec.split, + ring=spec.ring, + trigger_early=True, + push=spec.push, + ) + + +def _capturing() -> bool: + return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() + + +def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: + """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" + if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: + return x2d.clone(memory_format=torch.contiguous_format) + return x2d + + +def _head_takes_module(lm_head: nn.Module) -> bool: + """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along + the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" + weight = getattr(lm_head, "weight", None) + mapping = getattr(lm_head, "mapping", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and weight.is_contiguous() + and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" + and getattr(lm_head, "gather_output", False) + and getattr(lm_head, "gather_output_sizes", None) is None + and getattr(lm_head, "padding_size", None) == 0 + and getattr(lm_head, "bias", None) is None + and not getattr(lm_head, "has_any_quant", True) + and mapping is not None + and not mapping.enable_attention_dp + ) + + +def _plain_rmsnorm(norm: nn.Module) -> bool: + """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for + bit. An unknown module fails closed.""" + weight = getattr(norm, "weight", None) + return ( + IS_FLASHINFER_AVAILABLE + and isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "use_cuda_tile", True) + and not getattr(norm, "return_hp_output", True) + and getattr(norm, "nvfp4_scale", None) is None + and hasattr(norm, "variance_epsilon") + ) + + +class K3DecodeGemvs: + """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran + eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are + final; owned by the target.""" + + def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: + self.head_workspace = head_workspace + self._ran: Set[tuple] = set() + + @classmethod + def create( + cls, + lm_head: Optional[nn.Module] = None, + sites: Iterable[str] = tuple(SITES), + device: Optional[torch.device] = None, + ) -> "K3DecodeGemvs": + """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: + it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on + zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or + head whose kernel does not take its shape keeps the generic path.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" + ) + head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + if device is None: + device = ( + head_weight.device + if isinstance(head_weight, torch.Tensor) + else torch.device("cuda", torch.cuda.current_device()) + ) + state = cls() + for site in sites: + spec = SITES[site] + weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) + for rows in (1, 16, 32, 64) if spec.wide else (1,): + if rows > spec.wide_rows: + continue + state._project(site, weight.new_zeros(rows, _width(spec)), weight, warm=True) + del weight + if "dense_gate_up" in sites: + gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) + if _ctm_op.supports_situ_mul(gu): + for linear_beta in (None, 1.0): + k3_situ_mul(gu, 1.0, linear_beta) + state._ran.add(("situ_mul", linear_beta is not None)) + if lm_head is not None and _head_takes_module(lm_head): + x = head_weight.new_zeros(1, head_weight.shape[1]) + if _head_op.supports(x, head_weight): + workspace = K3HeadGemvWorkspace.create( + head_weight.shape[0], head_weight.shape[1], head_weight.device + ) + k3_head_gemv(x, head_weight, workspace) + state.head_workspace = workspace + state._ran.add(("lm_head",)) + torch.cuda.synchronize(device) + return state + + def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: + """``x @ weight.T`` (``[..., N]``: bf16, the site's sigmoid columns through the sigmoid; fp32 at an + ``out_fp32`` site; on a SiLU-and-mul site, ``silu(gate) * up`` of ``x = [gate | up]`` first) for ``site``'s + weight on its decode kernel, or None where none takes the call: more rows than the site's kernels take, + another shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its + module.""" + return self._project(site, x, weight, warm=False) + + def _project( + self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool + ) -> Optional[torch.Tensor]: + spec = SITES[site] + width = _width(spec) + if ( + weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (spec.n, spec.k) + or not weight.is_contiguous() + or x.dtype != torch.bfloat16 + or x.dim() < 1 + or x.shape[-1] != width + ): + return None + rows = x.numel() // width + if 0 < rows <= MAX_ROWS: + kernel = spec.small + elif spec.wide and MAX_ROWS < rows <= spec.wide_rows: + kernel = "wide" + else: + return None + key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) + capturing = _capturing() + if capturing and not warm and key not in self._ran: + return None + y = _run(spec, kernel, _dense_rows(x.reshape(rows, width)), weight) + if y is None: + return None + if not capturing: + self._ran.add(key) + return y.view(*x.shape[:-1], spec.n) + + def lm_head_logits( + self, rows: torch.Tensor, lm_head: nn.Module, gather: bool = True + ) -> Optional[torch.Tensor]: + """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on + ``gemm/k3_head_gemv`` (``gather`` False: the shard alone, ``[M, vocab / TP]``); None where it does not take + the call (more than `MAX_ROWS` rows, another head, ...).""" + workspace = self.head_workspace + if workspace is None or not _head_takes_module(lm_head): + return None + weight = lm_head.weight + if ( + tuple(weight.shape) != (workspace.n_out, workspace.k_in) + or weight.device != workspace.partials.device + or rows.dim() != 2 + or rows.dtype != torch.bfloat16 + or not 0 < rows.shape[0] <= MAX_ROWS + or rows.shape[1] != workspace.k_in + ): + return None + if _capturing() and ("lm_head",) not in self._ran: + return None + group = lm_head.mapping.tp_group + if gather and len(group) > 1 and mpi_disabled(): + return None + x = _dense_rows(rows) + if not _head_op.supports(x, weight): + return None + local = k3_head_gemv(x, weight, workspace) + if not gather or len(group) == 1: + return local + gathered = allgather(local, None, group) + return torch.cat(torch.split(gathered, rows.shape[0], dim=0), dim=-1) + + def dense_mlp( + self, + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + beta: float, + linear_beta: Optional[float], + ) -> Optional[torch.Tensor]: + """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on + ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on + ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then + runs the module.""" + gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] + rows = x.shape[0] if x.dim() == 2 else 0 + if ( + not 0 < rows <= MAX_ROWS + or linear_beta == 0.0 + or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) + or tuple(down_weight.shape) != (down.n, down.k) + or gate_up.n != 2 * down.k + ): + return None + if ( + _capturing() + and not { + ("dense_gate_up", "long", 0), + ("dense_down", "long", 0), + ("situ_mul", linear_beta is not None), + } + <= self._ran + ): + return None + gu = self.project("dense_gate_up", x, gate_up_weight) + if gu is None or not _ctm_op.supports_situ_mul(gu): + return None + return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) + + def embed_norm( + self, + input_ids: torch.Tensor, + table: torch.Tensor, + norm: nn.Module, + bank: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one + ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not + apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose + kernel has not run eagerly).""" + if not _plain_rmsnorm(norm): + return None + ids = input_ids.reshape(-1) + if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): + return None + key = ("embed_norm", ids.numel(), ids.dtype) + capturing = _capturing() + if capturing and key not in self._ran: + return None + normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) + if not capturing: + self._ran.add(key) + return normed + + +class _K3Head: + """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" + + __slots__ = ("_gemvs", "_lm_head") + + def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: + self._gemvs = gemvs + self._lm_head = lm_head + + def __call__(self, rows: torch.Tensor) -> torch.Tensor: + logits = self._gemvs.lm_head_logits(rows, self._lm_head) + return self._lm_head(rows) if logits is None else logits + + +class K3LogitsProcessor(nn.Module): + """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. + + It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock + processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. + """ + + def __init__(self, stock: nn.Module) -> None: + super().__init__() + self.stock = stock + self.gemvs: Optional[K3DecodeGemvs] = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) + return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) + + def lm_head_shard(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: + """This rank's bf16 vocabulary shard of ``lm_head(rows)`` on ``gemm/k3_head_gemv``, without the gather, or + None where ``gemvs`` does not take the rows. A speculative worker that keeps its logits vocabulary-sharded + computes them here, so they hold the values of the logits this processor gathers.""" + if self.gemvs is None: + return None + return self.gemvs.lm_head_logits(rows, lm_head, gather=False) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_moe.py new file mode 100644 index 000000000000..43f643d49f40 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/decode_moe.py @@ -0,0 +1,408 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's latent MoE on the catalog's Kimi K3 MoE entries. + +**At most 8 tokens** (`MAX_TOKENS`), a MoE layer runs: + +* `moe/k3_moe_front`, one kernel: this rank's slice of the MoE head (its latent-down rows and its router rows) as one + GEMV, the slices' all-gather over the TP group's `K3MoeHeadWorkspace`, the top-16 routing, the MXFP8 latent, and + the shared experts' gate_up + SiTU; +* `moe/k3_moe`: this rank's routed partial over its experts (on a `K3MoeState`); +* the latent all-reduce. On a pushing step (`DecodeStep.latent_push`: a pure decode step captured into a CUDA graph + whose attention layers all run the decode kernels) the routed experts run as `moe/k3_moe`'s push form, which stores + this rank's partial into every rank's `K3LatentExchange`, and `comm/k3_latent_reduce` sums the partials in the MNNVL + one-shot's order, so with the same bits, while the experts' grid completes. On every other step it is the routed + experts' all-reduce (one-shot, as `decode_comm.use_decode_one_shot` sets); +* the tail. The latent norm's weight is folded into the latent up projection at load, so + `[RMSNorm(latent) slice | shared activation] @ [latent up columns | shared down]` is this rank's share of the MoE + output. Where the next layer's pre-attention step (or the final norm) reduces it, the layer hands it on unreduced + as a `PendingTail`, which the consumer runs with its all-reduce and residual update as one `comm/k3_sandwich_tail` + kernel (`decode_comm.py`). Elsewhere the replicated tail runs: the latent RMS applied to the fp32 output of one + GEMV with the folded latent up weight, plus the shared experts' down projection and its all-reduce. + +**A wide decode step** (9 to 64 tokens, `WIDE_MAX_TOKENS`) keeps the sharded head and the row-parallel tail on +M-general ops: the head GEMV, `comm/mnnvl_allgather_split`, then `moe/k3_route_quant` and `moe/k3_moe` (on a +`K3MoeWideState`) beside the shared gate_up + SiTU, the latent all-reduce, and one GEMV of +`[RMSNorm(latent) slice | padding | shared activation]` with the tail weight. That is this rank's unreduced share, +which the consumer reduces with a plain all-reduce. + +The GEMVs run on the decode GEMV sites of `decode_gemv.py` where they take the call, else on the stock GEMM ops. + +`K3DecodeMoe` holds what every MoE layer shares: the head workspace and the latent exchange (both collective over +the TP group), the two `k3_moe` builds' scratch, and the TP group's MNNVL workspace (`decode_comm.K3DecodeComm`'s). +`K3DecodeMoeLayer` holds one layer's decode weights and its `k3_moe` counters. The target builds both in +`post_load_weights`, before any CUDA-graph capture, and runs every kernel once there so none compiles under a capture. + +Every MoE layer's push and reduce go to the one exchange, in the stream's order: each push is followed by exactly one +reduce of the same token count before the next push, on every rank in the same order (the exchange's call-order +invariant, `comm/k3_latent_reduce`). Every kernel between a reduce and the next push waits for its predecessor (or +launches without programmatic dependent launch), which the decode kernels a pushing step runs do; pushes run only in +CUDA-graph replays, so no other step's kernels sit between two pushes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_latent_reduce import ( + K3LatentExchange, + k3_latent_reduce, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allgather_split import ( + mnnvl_allgather_split, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe import ( + K3MoeHeadWorkspace, + K3MoeLayer, + K3MoeState, + K3MoeWideState, + is_supported, + k3_moe, + k3_moe_push, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe_front import ( + front_weight, + k3_moe_front, + weight_supported, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_route_quant import k3_route_quant +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel + +from .decode_comm import K3DecodeComm, PendingTail, wide_all_reduce + +# The most tokens of the front and the small k3_moe build (one token tile), and of the wide build. +MAX_TOKENS = 8 +WIDE_MAX_TOKENS = 64 + +# The tail weight's latent columns are zero-padded to whole 128-column k-tiles of the tail kernels +# (comm/k3_sandwich_tail takes the TP16 tail weight [7168, 256 + 384]). +TAIL_K_TILE = 256 + + +@dataclass(eq=False) +class K3DecodeMoe: + """What every MoE layer's decode path shares on one device: the TP group's ``K3MoeHeadWorkspace`` (the front's + all-gather), the ``k3_moe`` builds for up to 8 and up to 64 tokens (their scratch; the layers run one at a time on + one stream), the TP group's ``MnnvlWorkspace`` for a wide step's head all-gather, and the TP group's + ``K3LatentExchange`` for a pushing step's latent all-reduce (None: every step keeps the routed experts' + all-reduce). Built by `create`.""" + + head: K3MoeHeadWorkspace + small: K3MoeState + wide: K3MoeWideState + mnnvl: MnnvlWorkspace + exchange: Optional[K3LatentExchange] = None + + @classmethod + def create( + cls, mapping, device, i_tp: int, num_local: int, mnnvl: MnnvlWorkspace, push: bool = True + ) -> "K3DecodeMoe": + """The state for ``mapping``'s TP group on ``device``, for experts of ``i_tp`` intermediate columns per rank, + ``num_local`` of them on this rank. Collective (the head workspace and the latent exchange): every rank of the + group calls it at the same point, eagerly, before any CUDA-graph capture. ``push``: build the latent exchange + (where it takes the TP size: 4, 8 or 16 ranks); without it every step keeps the routed experts' all-reduce.""" + head = K3MoeHeadWorkspace.create(mapping) + exchange = None + if push: + try: + exchange = K3LatentExchange.create(mapping) + # Raised before any collective step, on every rank of the group alike: a TP size the exchange does not take. + except ValueError: + exchange = None + return cls( + head, + K3MoeState(device, i_tp, num_local), + K3MoeWideState(device, i_tp, num_local), + mnnvl, + exchange, + ) + + +def _experts(moe: nn.Module) -> tuple: + """The routed experts' TRTLLM-Gen W4A8_MXFP4_MXFP8 buffers ``k3_moe`` reads in place.""" + backend = moe.routed_experts.backend + return ( + backend.w3_w1_weight, + backend.w3_w1_weight_scale, + backend.w2_weight, + backend.w2_weight_scale, + ) + + +def layout_gaps(moe: nn.Module, tp_size: int, max_snapshots: int) -> list: + """Why the decode path does not take MoE layer ``moe`` (empty when it takes it), read once the routed experts' + buffers exist: the checks in order, up to the first that fails. ``max_snapshots``: the model's snapshot bank + rows.""" + gate, backend = moe.gate, moe.routed_experts.backend + shared = moe.shared_experts + gate_up, down = shared.gate_up_proj.weight, shared.down_proj.weight + if not moe._reduce_routed_output: + return ["a routed output the model does not reduce"] + if (moe.num_experts, moe.top_k, moe.moe_hidden_size, moe.hidden_size) != (896, 16, 3584, 7168): + return [ + f"experts / top-k / latent / hidden {(moe.num_experts, moe.top_k, moe.moe_hidden_size)}" + ] + if (gate.num_expert_group, gate.topk_group) != (1, 1): + return ["grouped routing"] + if moe.moe_hidden_size % (8 * tp_size) or moe.num_experts % (4 * tp_size): + return [f"a head slice of TP {tp_size}"] + latent = (gate.weight, moe.routed_expert_down_proj.weight, moe.routed_expert_up_proj.weight) + if not isinstance(moe.routed_expert_up_proj, nn.Linear) or any( + w.dtype != torch.bfloat16 for w in latent + ): + return ["latent projections or a router other than bf16"] + if ( + gate_up.dtype != torch.bfloat16 + or down.dtype != torch.bfloat16 + or shared.gate_up_proj.bias is not None + or shared.down_proj.bias is not None + or down.shape[0] != moe.hidden_size + ): + return ["a shared expert other than bf16 and unbiased"] + names = ( + "w3_w1_weight", + "w3_w1_weight_scale", + "w2_weight", + "w2_weight_scale", + "expert_size_per_partition", + ) + if ( + not all(hasattr(backend, name) for name in names) + or not is_supported(*_experts(moe), backend.expert_size_per_partition)[0] + ): + return ["routed-expert buffers k3_moe does not read"] + shared_cols, width = gate_up.shape[0] // 2, moe.moe_hidden_size // tp_size + if not weight_supported(tp_size, shared_cols, gate_up.shape[1], gate_up.device): + return ["a MoE front of this TP size and shared width"] + tail_cols = width + (-width % TAIL_K_TILE) + shared_cols + probe = gate_up.new_empty + if not K3DecodeComm.takes_tail( + probe(1, moe.moe_hidden_size), + probe(1, shared_cols), + probe(moe.hidden_size, tail_cols), + max_snapshots, + ): + return ["a row-parallel tail the sandwich tail kernel does not take"] + return [] + + +def fold_latent_norm(moe: nn.Module) -> None: + """Fold the latent RMSNorm's weight into the latent up projection's columns; the norm keeps a weight of ones, so + every step computes the same function and the tails may normalize before slicing.""" + up, norm = moe.routed_expert_up_proj, moe.routed_expert_norm + with torch.no_grad(): + up.weight.mul_(norm.weight.to(up.weight.dtype)[None, :]) + norm.weight.fill_(1) + + +@dataclass(eq=False) +class K3DecodeMoeLayer: + """One MoE layer's decode path: the front weight (this rank's head slice padded to whole tiles, then the shared + gate_up re-ordered), the head slice (a view of it), the row-parallel tail weight + ``[latent up columns lo:lo+width | padding | shared down]``, the routing bias, and its ``k3_moe`` handles on the + shared builds. Built by `create` once the weights are final (and the latent norm folded).""" + + state: K3DecodeMoe + front_weight: torch.Tensor + head_weight: torch.Tensor + tail_weight: torch.Tensor + tail_pad: Optional[torch.Tensor] + # The gate's routing bias detached (a view of the parameter): an op given a tensor that requires grad while + # autograd records, as in post_load_weights, returns outputs that require grad, which k3_moe cannot take. + bias: torch.Tensor + lo: int + width: int + shared_cols: int + small: K3MoeLayer + wide: K3MoeLayer + + @classmethod + def create( + cls, moe: nn.Module, state: K3DecodeMoe, tp_rank: int, tp_size: int + ) -> "K3DecodeMoeLayer": + """Build MoE layer ``moe``'s decode weights and its handles on ``state``'s builds (``layout_gaps`` empty).""" + width = moe.moe_hidden_size // tp_size + experts = moe.num_experts // tp_size + gate_up = moe.shared_experts.gate_up_proj.weight + shared_down = moe.shared_experts.down_proj.weight + up = moe.routed_expert_up_proj.weight + lo = tp_rank * width + pad = -width % TAIL_K_TILE + with torch.no_grad(): + head = torch.cat( + [ + moe.routed_expert_down_proj.weight[lo : lo + width], + moe.gate.weight[tp_rank * experts : (tp_rank + 1) * experts], + ] + ) + front = front_weight(head, gate_up) + parts = [up[:, lo : lo + width]] + if pad: + parts.append(up.new_zeros(up.shape[0], pad)) + parts.append(shared_down) + tail = torch.cat(parts, dim=1).contiguous() + weights = _experts(moe) + return cls( + state=state, + front_weight=front, + head_weight=front[: head.shape[0]], + tail_weight=tail, + tail_pad=up.new_zeros(WIDE_MAX_TOKENS, pad) if pad else None, + bias=moe.gate.e_score_correction_bias.detach(), + lo=lo, + width=width, + shared_cols=gate_up.shape[0] // 2, + small=state.small.layer(*weights), + wide=state.wide.layer(*weights), + ) + + def warm_up(self, moe: nn.Module) -> None: + """One call of every kernel of the decode path on zero inputs (M = 1), so none compiles under a capture: the + front (collective: every rank makes the same call), both ``k3_moe`` builds, ``k3_route_quant`` and, with the + latent exchange, the push build and the reduce (one push + reduce pair, collective). Once per model: every + layer's calls compile the same builds.""" + device = self.front_weight.device + x = torch.zeros(1, moe.hidden_size, dtype=torch.bfloat16, device=device) + ids, weights, x_fp8, x_sf, _ = self._front(moe, x) + offset = moe.routed_experts.backend.slot_start + k3_moe(x_fp8, x_sf, ids, weights, offset, self.small) + exchange = self.state.exchange + if exchange is not None: + # Two calls on one layer must not run back to back (moe/k3_moe): the push build starts once the call + # above has ended. + torch.cuda.synchronize(device) + k3_moe_push(x_fp8, x_sf, ids, weights, offset, self.small, exchange) + k3_latent_reduce(1, exchange) + logits = torch.zeros(1, moe.num_experts, dtype=torch.float32, device=device) + latent = torch.zeros(1, moe.moe_hidden_size, dtype=torch.bfloat16, device=device) + ids, weights, x_fp8, x_sf = k3_route_quant( + logits, self.bias, latent, float(moe.gate.routed_scaling_factor), early_trigger=True + ) + k3_moe(x_fp8, x_sf, ids, weights, offset, self.wide) + torch.cuda.synchronize(device) + + def _front(self, moe: nn.Module, x: torch.Tensor): + return k3_moe_front( + x.contiguous(), + self.front_weight, + self.bias, + float(moe.gate.routed_scaling_factor), + self.shared_cols, + *moe._situ_betas, + self.state.head, + ) + + def takes(self, hidden_states: torch.Tensor, step, partial_tail: bool) -> bool: + """Whether this decode path runs the layer on ``step`` (bf16 rows ``hidden_states``): any step of at most + `MAX_TOKENS` tokens; with ``partial_tail``, also a wide decode step.""" + rows = hidden_states.shape[0] + if hidden_states.dtype != torch.bfloat16: + return False + return 0 < rows <= MAX_TOKENS or (partial_tail and step.wide and rows <= WIDE_MAX_TOKENS) + + def forward( + self, + moe: nn.Module, + hidden_states: torch.Tensor, + gemvs, + partial_tail: bool, + push: bool = False, + ) -> Union[torch.Tensor, PendingTail]: + """The MoE output of ``hidden_states`` (``takes`` holds): with ``partial_tail``, this rank's unreduced share + (a ``PendingTail`` at most `MAX_TOKENS` tokens); else the reduced output. ``gemvs``: the decode GEMVs' state, + or None. ``push`` (the step's ``DecodeStep.latent_push``): at most `MAX_TOKENS` tokens, the latent all-reduce + is the routed experts' push form plus ``comm/k3_latent_reduce`` on the state's exchange (the same bits as the + routed experts' all-reduce); a state without the exchange ignores it.""" + if hidden_states.shape[0] > MAX_TOKENS: + return self._wide(moe, hidden_states, gemvs) + ids, weights, x_fp8, x_sf, shared_act = self._front(moe, hidden_states) + offset = moe.routed_experts.backend.slot_start + exchange = self.state.exchange + if push and exchange is not None: + k3_moe_push(x_fp8, x_sf, ids, weights, offset, self.small, exchange) + latent = k3_latent_reduce(hidden_states.shape[0], exchange) + else: + routed = k3_moe(x_fp8, x_sf, ids, weights, offset, self.small) + latent = moe.routed_experts.all_reduce(routed) + if partial_tail: + return PendingTail( + latent.contiguous(), + shared_act.contiguous(), + self.tail_weight, + self.lo, + float(moe.routed_expert_norm.variance_epsilon), + ) + # The replicated tail: the latent RMS on the fp32 accumulator of the folded latent up projection, plus the + # shared experts' reduced output, rounded to bf16 once. + shared = moe.shared_experts + shared_out = shared.down_proj(shared_act, layer_idx=shared.layer_idx) + up = _gemv(gemvs, "moe_up", latent, moe.routed_expert_up_proj.weight, out_fp32=True) + scale = torch.rsqrt( + latent.float().pow(2).mean(-1, keepdim=True) + moe.routed_expert_norm.variance_epsilon + ) + return (up * scale + shared_out.float()).bfloat16() + + def _wide(self, moe: nn.Module, hidden_states: torch.Tensor, gemvs) -> torch.Tensor: + """A wide decode step's MoE: this rank's unreduced share of the output, ``[M, hidden]`` bf16.""" + x = hidden_states.contiguous() + head = _gemv(gemvs, "moe_head", x, self.head_weight, out_fp32=True) + routed_in, router_logits = mnnvl_allgather_split(head, self.width, self.state.mnnvl) + shared = moe.shared_experts + + def _routed_partial(): + ids, weights, x_fp8, x_sf = k3_route_quant( + router_logits, self.bias, routed_in.contiguous(), + float(moe.gate.routed_scaling_factor), early_trigger=True, + ) # fmt: skip + return k3_moe( + x_fp8, x_sf, ids, weights, moe.routed_experts.backend.slot_start, self.wide + ) + + def _shared_activation(): + return shared._apply_activation( + _gemv(gemvs, "moe_shared_gate_up", x, shared.gate_up_proj.weight) + ) + + routed, shared_act = maybe_execute_in_parallel( + _routed_partial, + _shared_activation, + moe.moe_main_event, + moe.moe_shared_event, + moe.shared_expert_stream, + disable_on_compile=True, + ) + # The latent norm's weight is folded into the tail weight: normalize the whole latent row, keep this rank's + # columns. + normed = moe.routed_expert_norm(wide_all_reduce(moe.routed_experts.all_reduce, routed)) + parts = [normed[:, self.lo : self.lo + self.width]] + if self.tail_pad is not None: + parts.append(self.tail_pad[: x.shape[0]]) + parts.append(shared_act) + return _gemv(gemvs, "moe_tail", torch.cat(parts, dim=1), self.tail_weight) + + +def _gemv( + gemvs, site: str, x: torch.Tensor, weight: torch.Tensor, out_fp32: bool = False +) -> torch.Tensor: + """``x @ weight.T`` on ``site``'s decode GEMV where it takes the call (``decode_gemv.K3DecodeGemvs.project``), + else on the stock GEMM: cuBLAS through ``trtllm::dsv3_router_gemm_op`` for an fp32 output, ``F.linear`` else.""" + y = None if gemvs is None else gemvs.project(site, x, weight) + if y is not None: + return y + if out_fp32: + # The op reads the weight with a leading dimension of K: a strided weight would give wrong values. + if not weight.is_contiguous(): + raise ValueError( + f"the fp32 GEMM needs a contiguous weight, got strides {weight.stride()}" + ) + return torch.ops.trtllm.dsv3_router_gemm_op( + x.contiguous(), weight.t(), bias=None, out_dtype=torch.float32 + ) + return torch.nn.functional.linear(x, weight) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index e2a3f736e8f6..082f8798ae75 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -1,21 +1,3595 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 MXFP4, SM100, attention TP16 and MoE TP4 x EP4.""" +"""ModelingV2 target: Kimi K3 (MXFP4) / sm_100 / tp16 attention, routed experts moe_tp 4 x moe_ep 4. +Kimi K3's language model: 93 layers, hidden 7168. Every fourth layer (3, 7, ..., 91) is MLA attention with 96 +query heads; the others are Kimi Delta Attention (KDA), a gated linear-attention recurrence with per-request state. +Layer 0's MLP is dense. Every other layer's MLP is a latent MoE: 896 routed experts, top-16, expert width 3584. +Attention residuals mix each sublayer's input from a bank of earlier outputs. The routed experts are MXFP4 (the +checkpoint's compressed-tensors default, read as W4A8_MXFP4_MXFP8); everything else is bf16, and so is the KV pool. + +`tp16_moetp4ep4` is `tensor_parallel_size: 16` with `moe_tensor_parallel_size: 4` and +`moe_expert_parallel_size: 4`, no attention data parallelism, all 16 GPUs in one NVLink domain. Attention is +head-split (6 MLA query heads per rank), and every rank holds a quarter of the width of a quarter of the experts. + +**Each step is classified once, on the host, from its shape** (`decode_step`, a `DecodeStep` or None), and the +classification decides which kernels each module runs: + +* **small**: at most 8 tokens (`DECODE_MAX_TOKENS`, one token tile), context requests included. The token-count + kernels take it: decode GEMVs, the MoE front and routed experts, the sandwiches, the embedding and residual + epilogues. +* **decode**: a pure decode step of R <= 8 generation requests with the same T <= 8 tokens each (one token without + speculation, 1 + 7 drafts with DSpark) and no context request, so at most 64 tokens. The request-aware kernels + take it: MLA attention and its KV store, the KDA verify, the drafter's attention. +* **wide**: a decode step of more than one token tile (DSpark verify of several requests). It keeps the decode + layout's MoE head and tail, on M-general ops. + +Every other step (prefill, mixed steps, decode steps above those bounds) runs the **generic path**: this target's +text model (`KimiLinearModel` below: decoder layers, attention residuals, the MLA / KDA / MoE runtimes), computed as +the built-in Kimi K3 text model computes it, on stock modules and ops that have no catalog entries yet. One weight +differs: where the MoE decode path takes a layer, its latent norm's weight is folded into the latent up projection +(the same function, rounded differently). And the routed experts' top-16 is computed outside the TRTLLM-Gen kernel +(`KimiK3MoeRoutingMethod`), with the kernel's arithmetic. `UNCERTIFIED_GENERIC_CALLS` names the stock code. + +The text model hands each step's classification to its attention modules, which run a **decode step** on the K3 +decode kernels' catalog entries: + +* KDA (`K3DecodeKDA`): one token per request, the fused input projection and the plain decode in one + `ssm/k3_kda_decode_attn` launch. Verify tokens (DFlash / DSpark at an even verify width up to 8): the cache manager + then keeps the last verify's draft records (per draft the row innovations, beta * k and the decay), and every + verify of the layer, on any step, runs the kernels that write them and replay the accepted drafts from them: + `ssm/k3_kda_attn` for one request of 8 tokens (the projection fused in), else `ssm/k3_kda_verify`. +* MLA (`K3DecodeMLA`): `attention/k3_mla_qkv` (the query path and the step's latent KV rows into the paged cache), + then `attention/k3_mla_attn_vb_out` (the attention, v_b and the output gate in one launch). + +The projections around them run on the decode GEMV sites of `decode_gemv.py` (the [W_a; W_g] and KDA verify-row +projections, `o_proj` on every classified step), as do the LM head, the embedding and layer 0's dense MLP. A +classified step's attention-residual epilogues (the selection and the RMSNorm after it) take the fused kernels up to +one token tile, 32 tokens on a wide decode step. On any step of at most 16 tokens but a wide one, the post-attention +all-reduce carries the residual update (`decode_comm.py`): at most 8 tokens on an attention decode branch, o_proj, the +all-reduce and the update are one `comm/k3_sandwich_oproj` kernel; otherwise the attention's unreduced o_proj output +goes through `comm/mnnvl_allreduce_attn_res`. The stock MNNVL all-reduces send one-shot up to 4 MiB, a wide decode +step's up to the stock 1 MiB. The state those kernels share (the KDA projection's Lamport buffers, the MLA attention +workspace, the decode GEMVs' state, the TP group's MNNVL and sandwich workspaces, the MoE path's) lives in typed +objects this target creates in `post_load_weights`, before any graph capture. + +A MoE layer on a step of at most 8 tokens runs `decode_moe.py`: `moe/k3_moe_front` (this rank's head slice, its +all-gather, the routing, the MXFP8 latent and the shared experts' gate_up + SiTU in one kernel), `moe/k3_moe`, the +latent all-reduce (on a pushing step, `DecodeStep.latent_push`, `moe/k3_moe`'s push form plus `comm/k3_latent_reduce`), +then the row-parallel tail, which the next layer's pre-attention step (the final norm's, after the +last layer) runs with its all-reduce and residual update as one `comm/k3_sandwich_tail` kernel. A wide decode step's +MoE keeps the sharded head and the row-parallel tail, on `moe/k3_route_quant`, `moe/k3_moe` and M-general ops. + +**What this target asserts rather than adapts**: SM 10.0; the topology above; the MXFP4 checkpoint's quantization +(W4A16_MXFP4 with no per-layer declarations, so the routed experts run the W4A8_MXFP4_MXFP8 default and the excluded +modules stay bf16); bf16 weights and a bf16 KV pool; +tokens_per_block 64 (the MLA generation kernels K3's 96 heads reach exist only at 64); the V2 hybrid KV / state +manager, which holds the KDA states, with block reuse off and fp32 recurrent states; an all-reduce strategy of AUTO or +MNNVL. The construction-time ones fail in `__init__`, the per-engine ones on the first forward, each naming the +setting. A layer the decode kernels do not take fails the weight load. + +**Text only.** The checkpoint is the vision-language wrapper. This target builds and loads no vision tower (its +weights are a predicted non-load, `weights.py`), and a step carrying multimodal input raises. + +**Speculative decoding** goes through the stock one-engine shell: DSpark or DFlash with an external drafter +checkpoint, and SA. The DSpark drafter is this target's `K3DSparkDrafter`, the stock GQA drafter with a decode step's +block on the drafter entries (`attention/k3_drafter_attn_qknorm` and the drafter's decode GEMV sites), which this +target puts in place of the shell's stock one. Where every attention all-reduce runs over MNNVL, the drafter also runs +on the TP group's collective state: its context projection's `fc` is split by input feature over the group, +`hidden_norm` applied in the all-reduce of the partial products (`comm/mnnvl_fusion_allreduce`), and a block's +residual adds and RMSNorms run in its all-reduces (`comm/k3_sandwich_plain` with the projection up to 8 rows, else +`comm/mnnvl_fusion_allreduce`). The DFlash / DSpark worker is this target's (`spec_worker.py`): the stock worker with +the decode kernels `trtllm::k3_spec_accept`, `k3_ctx_kv` and `k3_markov` behind its `k3_decode`. +""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( + k3_drafter_attn_qknorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_vb_out import ( + k3_mla_attn_vb_out, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_attn_workspace import ( + K3MlaAttnWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_mla_qkv import k3_mla_qkv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_attn import k3_kda_attn +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_buffers import K3KdaBuffers +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_decode_attn import ( + k3_kda_decode_attn, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify +from tensorrt_llm._torch.attention.backends import AttentionMetadata + +# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. +from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.decode_view import k3_mla_decode_view +from tensorrt_llm._torch.distributed import AllReduce from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.models.modeling_dspark import ( + GQADSparkForCausalLM, + draft_is_embedded_in_target, +) +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.moe.fused_moe import ( + ConfigurableMoE, + SiTuActivation, + TRTLLMGenFusedMoE, + create_moe, +) +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.config_utils import is_mla +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from . import decode_comm as _decode_comm +from . import decode_gemv as _decode_gemv +from . import decode_moe as _decode_moe +from . import spec_worker as _spec_worker +from . import weights as _weights + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, +# and the certification is per arch. +_SM = (10, 0) + +#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in +#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router +#: and fused-A GEMMs), then the decode kernels'. +REQUIRED_TRTLLM_OPS = ( + "attn_res_fwd", + "attn_res_rmsnorm_fwd", + "attn_res_add_rmsnorm_fwd", + "attn_res_add_rmsnorm_persistent_fwd", + "kda_prefill", + "kda_decode", + "kda_mtp_decode", + "dsv3_router_gemm_op", + "dsv3_fused_a_gemm_op", + "k3_kda_decode_attn", + "k3_kda_attn", + "k3_kda_verify", + "k3_mla_qkv", + "k3_mla_attn_vb_out", + # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). + "k3_decode_gemv", + "k3_ctm_gemv_wide", + "k3_ctm_gemv_long", + "k3_situ_mul", + "k3_head_gemv", + "k3_embed_norm", + "allgather", + # The decode path's collectives (decode_comm.py) and MoE (decode_moe.py). + "mnnvl_allreduce_attn_res", + "k3_sandwich_oproj", + "k3_sandwich_tail", + "k3_moe_front", + "k3_moe", + "k3_latent_reduce", + "k3_route_quant", + "mnnvl_allgather_split", + # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces + # with the residual add and RMSNorm (the split context projection's with hidden_norm). + "k3_ctm_gemv", + "k3_ctm_gemv_swiglu", + "k3_drafter_attn_qknorm", + "k3_sandwich_plain", + "mnnvl_fusion_allreduce", +) + +#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. +#: A renamed field upstream fails here, loudly, instead of reading as a default. +REQUIRED_ENGINE_FIELDS = { + "attn_metadata": ( + "num_contexts", + "num_generations", + "num_seqs", + "num_tokens", + "seq_lens", + "tokens_per_block", + "kv_cache_manager", + "mamba_metadata", + ), + "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), +} + +#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every +#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry +#: leaves the list when a catalog entry replaces it. +UNCERTIFIED_GENERIC_CALLS = ( + # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. + "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", + "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", + "tensorrt_llm._torch.models.modeling_utils.DecoderModel", + # The text model's stock modules. + "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", + "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", + "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", + "tensorrt_llm._torch.moe.fused_moe.create_moe", + "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", + "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", + "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", + "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", + "tensorrt_llm._torch.modules.situ.SituAndMul", + "tensorrt_llm._torch.modules.rms_norm.RMSNorm", + "tensorrt_llm._torch.distributed.AllReduce", + "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", + # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries + # do not take a block, its context projection without the TP group's collective state, its context k / v, its + # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. + "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", + "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", + "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", + "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", +) + +# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware +# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). +DECODE_MAX_TOKENS = 8 +MAX_REQUESTS = 8 +MAX_TOKENS_PER_REQUEST = 8 + +# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own +# get_model_defaults sets it for the same reason). +_TOKENS_PER_BLOCK = 64 + +_LANG_PREFIX = "language_model." + + +# ---------------------------------------------------------------------------------------------------------------------- +# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and +# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, +# `K3DecodeMLA` below). +# ---------------------------------------------------------------------------------------------------------------------- + +# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections +# instead of the min-latency fused GEMM op (read once at import). +_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( + os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" +) + + +# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table +# holds cos=1/sin=0), but the chunked-context path indexes the table by +# absolute position, so it must cover max_position_embeddings (~512MB per +# backend for the 1M-position checkpoint); a smaller table is read out of +# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context +# deployments. +_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" + + +class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): + """DeepSeek-V3 routing, computed outside the TRTLLM-Gen MoE kernel. + + The kernel's own top-16 of 896 experts is the slow part of a generic MoE call of a few tokens. Outside it, the + scheduler routes with the backend's fused route + MXFP8 quantize (``trtllm::k3_route_quant``) up to 64 tokens and + with ``noaux_tc_op`` above, and the kernel takes the expert ids and weights. Both compute the kernel's arithmetic + and break ties the same way; a bf16 weight can differ by one ulp where it lies within fp32 rounding of a bf16 + rounding boundary. From about a thousand tokens on, ``noaux_tc_op`` takes longer than the kernel's routing. + """ + + @property + def requires_separated_routing(self) -> bool: + return True + + +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + @property + def routing_method(self) -> KimiK3MoeRoutingMethod: + """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return KimiK3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + +def _get_text_config(pretrained_config: "PretrainedConfig"): + """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" + if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( + not hasattr(pretrained_config, "linear_attn_config") + and hasattr(pretrained_config, "text_config") + ): + return pretrained_config.text_config + return pretrained_config + + +def _is_kda_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] + + +def _is_mla_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" + + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" + + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" + + +KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" + + +_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") + + +# Turning the port on should not require knowing which of the three topologies +# is the right one: "1" selects the measured policy (``split``), so enabling the +# feature and choosing the policy are one action through one variable. +_ATTN_RES_TOPOLOGY_ON = "1" + + +def _read_attn_res_topology() -> str: + """Default ``per_token``: the persistent kernel is opt-in. + + ``1`` is the only accepted on-value and resolves to ``split``. The named + topologies stay available for measurement: ``persistent`` uses the + persistent kernel at every shape it implements, ``per_token`` at none. + """ + raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") + if raw == _ATTN_RES_TOPOLOGY_ON: + return "split" + if raw not in _ATTN_RES_TOPOLOGIES: + # Loudly, for the same reason as the token ceiling below: a mistyped A/B + # arm that silently fell back to the default would measure one side + # twice and report no difference. + raise ValueError( + f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " + f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " + f"(which means 'split'), got {raw!r}" + ) + return raw + + +def _read_fused_attn_res_max_tokens() -> int: + """Resolved after the topology, because its default follows it. + + With the persistent port off -- the default -- the ceiling is 1, the + pre-existing gate: the fused epilogue is taken at the single-token decode + shape and nowhere else. Enabling the port raises it to 32, the top of the + measured range, so that "off" keeps meaning "unchanged". + """ + default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" + raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) + try: + value = int(raw) + except ValueError: + # Failing loudly matters more than usual here: a mistyped A/B arm that + # silently fell back to the default would measure the candidate twice + # and report no difference. + raise ValueError( + f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value < 1: + raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") + return value + + +_ATTN_RES_TOPOLOGY = _read_attn_res_topology() + + +_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() + + +def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: + """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. + + No token ceiling: the persistent grid is sized by the SM count, not by the + token count, so prefill is the case it exists for. + """ + del M # deliberately unused; see above + return H == 7168 and 2 <= N <= 9 + + +def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: + """Pick between the two fused kernels for this call site. + + ``persistent`` takes the persistent kernel at every shape it implements; + ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in + for the prefill/decode boundary. Shapes the persistent kernel does not + implement fall through to the caller's existing gate and land on the + unfused path. + """ + if not _persistent_attn_res_applicable(M, H, N): + return False + if _ATTN_RES_TOPOLOGY == "persistent": + return True + return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> Optional[torch.Tensor]: + """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. + + Returns ``None`` when the call falls outside the fused kernel's + contract (dtype/shape/arch) so the caller can use the exact fp32 reference + instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` + layout. Candidate order matches the reference: snapshots first, the + running prefix sum last. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + try: + attn_res_op = torch.ops.trtllm.attn_res_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output, _rsigma, _probs, _logits = attn_res_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + ) + return output.reshape(M, H) + + +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that + the shape gate let the call through, and a rejected call looks exactly like + a disabled one in the logs. Emitted at debug level: it is once per distinct + shape, not once per process, so it is a diagnostic rather than a summary. + """ + logger.debug_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, + default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, + # which is the measured window; larger token counts have not been measured + # and fall back to the unfused add + attn_res_fwd + RMSNorm path. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + _, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + None, + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) + return output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in + ``_apply_attn_res_rmsnorm_fused``). + + The production residual add produces a BF16 tensor that remains live across + the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual output, + while avoiding a separate add launch and a re-read of the intermediate by + attention-residual selection. + """ + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # Same measured window as _apply_attn_res_rmsnorm_fused above. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + updated_prefix_sum, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + addend.reshape(M, 1, H).contiguous(), + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = attn_res_add_rmsnorm_op( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + +def _apply_attn_res( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> torch.Tensor: + """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). + + prefix_sum: ``[num_tokens, hidden_size]`` + block_residual: ``[num_snapshots, num_tokens, hidden_size]`` + + Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's + contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. + Only the fallback boundary restores the HF ``[M, K, H]`` layout. + """ + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + block_residual_hf = block_residual.transpose(0, 1) + v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.eps) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the + fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_rmsnorm_fused( + prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in + ``_apply_attn_res_and_rmsnorm``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + + +# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where +# they are faster than the add -> attn_res -> RMSNorm chain. +_WIDE_ATTN_RES_MAX_TOKENS = 32 + + +def _attn_res_max_tokens(step: DecodeStep) -> int: + """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step + ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" + return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS + + +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + + +# The routed experts' module path prefixes. ``exclude_modules`` matches with +# wildcards and walks ancestor prefixes, so an empty prefix would widen what +# matches instead of just missing. +_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") + + +# Routed-expert quantization used when the checkpoint declares nothing per +# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors +# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that +# checkpoint is what this default has always served. +_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 + + +def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_packed_mxfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), + w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), + w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), + w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), + w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), + w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( + kinds=("weight_packed", "weight_scale"), + loader=_load_packed_mxfp4_expert, + loaded_slots_attr="_packed_mxfp4_loaded_slots", + needs_layer_finalize=False, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + self.num_experts = cfg.num_experts + self.top_k = cfg.num_experts_per_token + self.moe_hidden_size = cfg.routed_expert_hidden_size + # ValueError (not assert): these guard unsupported checkpoint + # configurations and must stay active under ``python -O``. + if self.moe_hidden_size is None: + raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") + if not getattr(cfg, "latent_moe_use_norm", False): + raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") + + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) + self._situ_betas = (situ_beta, situ_linear_beta) + dtype = torch.bfloat16 + + # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 + # accumulate/output (checkpoint stores the gate weight in bf16; saves a + # per-layer input cast + fp32 splitK pair on the bs1 decode path). + # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True + self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) + + routed_moe_model_config = self._routed_moe_model_config(model_config) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) + routed_moe_kwargs = dict( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_intermediate_size, + dtype=dtype, + # Kimi owns the latent reduction so it can order that collective + # after the shared expert's auxiliary-stream reduction. + reduce_results=False, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + # Let CommunicationFactory select the best available strategy. + communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. As measured once: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver + # instead, which reports the rejection trail. + # + # CUTLASS is absent on purpose: it is the fallback target, so + # "degraded to CUTLASS" is not a thing that can happen to it. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + ) + self._check_trtllm_situ_quant( + routed_moe_model_config.moe_backend, routed_quant_config.quant_algo + ) + + self.routed_experts = create_moe(**routed_moe_kwargs) + if not isinstance(self.routed_experts, ConfigurableMoE): + raise RuntimeError( + "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." + ) + if self.routed_experts.layer_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "dynamic EPLB or replicated expert slots." + ) + local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) + if local_expert_ids != list( + range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) + ): + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming currently requires a " + "contiguous static expert partition." + ) + self.local_expert_ids = tuple(local_expert_ids) + self.experts_per_rank = len(local_expert_ids) + self.expert_lo = local_expert_ids[0] + self.expert_hi = self.expert_lo + self.experts_per_rank + + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = model_config.mapping.tp_size > 1 + self._reduce_routed_output = ( + use_shared_tp + and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + ) + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self.shared_experts = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=shared_model_config, + reduce_output=use_shared_tp, + layer_idx=layer_idx, + is_shared_expert=True, + ) + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.moe_main_event = torch.cuda.Event() + self.moe_shared_event = torch.cuda.Event() + self.routed_expert_down_proj = nn.Linear( + cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype + ) + # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer + # fallback is the same fp32-variance eager math as KimiK3RMSNorm). + self.routed_expert_norm = RMSNorm( + hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. + self.decode_moe: Optional[_decode_moe.K3DecodeMoeLayer] = None + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, projection.weight.t(), None, None + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """The routed-expert ``(moe_tp, moe_ep)`` split: the user config's explicit + ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` (4 x 4, asserted at + construction).""" + return mapping.moe_tp_size, mapping.moe_ep_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Routed-expert quantization for ``layer_idx``: the MXFP4 checkpoint declares nothing per layer (asserted at + construction), so the experts keep the ``W4A8_MXFP4_MXFP8`` default. + + An exclusion outranks the default: ``create_weights`` treats an override + as authoritative over anything ``__post_init__`` wrote, so this return + value stands in for both quantization passes and exclusion is the one + that runs second. It is matched as a pattern, so it is asked only about + real module names. + """ + quant_config = model_config.quant_config + if quant_config is not None and any( + quant_config.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + logger.debug( + "Kimi K3 layer %d routed experts: excluded from quantization, " + "keeping them unquantized", + layer_idx, + ) + return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) + + logger.debug( + "Kimi K3 layer %d routed experts: no per-layer quant config in the " + "checkpoint, defaulting to %s", + layer_idx, + _K3_DEFAULT_ROUTED_QUANT_ALGO, + ) + return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) + + @staticmethod + def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: + """Reject a routed-expert format trtllm-gen has no fused SiTu cubin for. + + trtllm-gen has fused SiTu FC1 cubins for two input formats and no + standalone SiTu activation kernel, so anything else has to die here + rather than in a cubin lookup deep inside the runner. Checked against + the resolved backend, not the K3 architecture branch, because the + generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also + land on TRTLLM. + + The admitted set is read off the backend rather than restated here, + because restating it is what broke. This guard was written in #17865 + when MXFP4 was the only fused SiTu drop; #17940 then added the NVFP4 + (group-16 ``Bmm_E2m1_E2m1E2m1_..._siTuGlu_*``) cubins and updated + ``TRTLLMGenFusedMoE``'s set without touching this copy. For the week + in between, an NVFP4 K3 checkpoint could not start at all -- and not + only when TRTLLM was asked for by name, because + ``ModelConfig.resolve_moe_backend`` sends every K3 architecture to + TRTLLM, so the default AUTO configuration hit this raise too. The unit + tests did not catch it: they call ``create_moe`` directly and never + reach this guard, so the kernel path stayed green while the model path + was closed. + + A staticmethod, not an inline block, so that the invariant is + reachable from a test without constructing the whole runtime. + """ + situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() + if moe_backend != "TRTLLM" or quant_algo in situ_supported: + return + supported = ", ".join(sorted(algo.name for algo in situ_supported)) + raise ValueError( + f"Kimi K3 routed experts are quantized as {quant_algo}, which the " + "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " + f"exist only for {supported}. Set moe_config.backend to CUTLASS " + "or MEGAMOE_CUTEDSL." + ) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config, with the split of ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. + supported_backends = { + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + f"got {model_config.moe_backend!r}." + ) + if model_config.moe_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "EPLB or replicated expert slots." + ) + mapping = model_config.mapping + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + logger.info_once( + f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " + f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", + key="kimi_k3_moe_tp_ep_split", + ) + + mapping_dict = mapping.to_dict() + mapping_dict["moe_cluster_size"] = 1 + mapping_dict["moe_tp_size"] = moe_tp + mapping_dict["moe_ep_size"] = moe_ep + routed_mapping = Mapping.from_dict(mapping_dict) + + routed_model_config = copy.copy(model_config) + routed_model_config._frozen = False + routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) + routed_model_config.mapping = routed_mapping + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend in { + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + }: + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) + if configured_moe_max_num_tokens < default_moe_max_num_tokens: + logger.info_once( + "Kimi K3 MegaMoE raises moe_max_num_tokens from " + f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " + "because the global DP SymmBuffer requires capacity for " + "max_num_tokens * dp_size.", + key=( + "kimi_k3_megamoe_capacity_override_" + f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" + ), + ) + routed_model_config.moe_max_num_tokens = max( + configured_moe_max_num_tokens, + default_moe_max_num_tokens, + ) + routed_model_config._frozen = True + return routed_model_config + + def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: + """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel + share (``partial_tail``): the decode path takes the step with that tail.""" + return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + partial_tail: bool = False, + *, + step: DecodeStep, + ) -> Union[torch.Tensor, _decode_comm.PendingTail]: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification + (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where + ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at + most 8 tokens, a tensor on a wide decode step.""" + decode = self.decode_moe + if decode is not None and decode.takes(hidden_states, step, partial_tail): + return decode.forward( + self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push + ) + if partial_tail: + raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") + return self.forward_generic(hidden_states, all_rank_num_tokens) + + def forward_generic( + self, hidden_states: torch.Tensor, all_rank_num_tokens=None + ) -> torch.Tensor: + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None -from ..modeling import K3TargetSpec, ModelingV2KimiK3 -from ..modeling import _check_construction as _check_target_construction + def _routed_output(): + # Latent down/up projections via the min-latency fused GEMM op: + # at <=16 tokens (decode graphs) it runs a single pipelined + # bf16 kernel per projection instead of cuBLAS's split-K GEMV + + # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); + # for larger token counts the op falls back to cuBLAS internally. + # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear + # (A/B escape hatch). When the FP8 weight-read conversion has + # replaced the projection module, call it directly: its weight is + # an e4m3 buffer the bf16 dsv3 op must not read, and its forward + # is already a single fused GEMM (fp8_swap_ab_gemm). + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + if self._reduce_routed_output: + return y + # Communication-backed paths return a complete routed result. + y = self.routed_expert_norm(y) + return self._routed_projection(y, self.routed_expert_up_proj) -_TARGET = K3TargetSpec(sm=(10, 0), quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8, moe_tp=4, moe_ep=4) + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. The shared GatedMLP includes its output all-reduce on + # the auxiliary stream. The join below must precede the routed + # all-reduce: concurrent collectives on different streams can corrupt + # SYMM_MEM all-reduce state. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + disable_on_compile=True, + ) + if self._reduce_routed_output: + routed_latent = moe_all_reduce(routed_out) + routed_latent = self.routed_expert_norm(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + return routed_out + shared_out + + +def resolve_attention_quant_config( + config: ModelConfig | None, layer_idx: int, projection: str +) -> QuantConfig: + """Resolve a checkpoint projection, including mixed-precision exclusions.""" + if config is None: + return QuantConfig() + global_config = config.quant_config or QuantConfig() + names = [ + f"{prefix}layers.{layer_idx}.self_attn.{projection}" + for prefix in ("language_model.model.", "model.", "") + ] + if any(global_config.is_module_excluded_from_quantization(name) for name in names): + return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) + declarations = config.quant_config_dict or {} + matches = [declarations[name] for name in names if name in declarations] + if matches: + selected = matches[0] + if any(match.quant_algo != selected.quant_algo for match in matches[1:]): + raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") + elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: + selected = QuantConfig() + else: + selected = global_config + if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " + f"quantization {selected.quant_algo}" + ) + if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): + raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") + result = copy.copy(selected) + result.kv_cache_quant_algo = global_config.kv_cache_quant_algo + return result + + +class KimiMLARuntime(nn.Module): + """Wraps K3 MLA and applies its external TP output reduction.""" + + def __init__( + self, + cfg: "PretrainedConfig", + layer_idx: int, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the + # final output reduction in this wrapper so the output gate remains + # between attention and the row-parallel o_proj. + mapping = model_config.mapping + reduce_output = mapping.tp_size > 1 + self._o_allreduce = ( + AllReduce( + mapping=mapping, + strategy=model_config.allreduce_strategy, + dtype=torch.bfloat16, + ) + if reduce_output + else None + ) + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in ( + "q_a_proj", + "kv_a_proj_with_mqa", + "q_b_proj", + "kv_b_proj", + "g_proj", + "o_proj", + ) + } + attention_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + attention_config._frozen = model_config._frozen + self.mixer = K3DecodeMLA( + hidden_size=cfg.hidden_size, + num_heads=cfg.num_attention_heads, + q_lora_rank=cfg.q_lora_rank, + kv_lora_rank=cfg.kv_lora_rank, + qk_nope_head_dim=cfg.qk_nope_head_dim, + qk_rope_head_dim=cfg.qk_rope_head_dim, + v_head_dim=cfg.v_head_dim, + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + use_output_gate=cfg.mla_use_output_gate, + max_position_embeddings=max_positions, + model_config=attention_config, + aux_stream_dict=aux_stream_dict, + ) + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" + return self.mixer.will_run_decode_branch(attn_metadata, step) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the + mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer( + None, hidden_states, attn_metadata, step=step, project_output=project_output + ) + return self._reduce(out) if project_output and reduce_output else out + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + out = self.mixer.forward_generic(None, hidden_states, attn_metadata) + return self._reduce(out) if reduce_output else out + + def _reduce(self, out: torch.Tensor) -> torch.Tensor: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + return self._o_allreduce(out) if self._o_allreduce is not None else out + + +class KimiLinearDecoderLayer(nn.Module): + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + dtype = torch.bfloat16 + + self.is_kda = _is_kda_layer(cfg, layer_idx) + is_mla = _is_mla_layer(cfg, layer_idx) + if self.is_kda == is_mla: + raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") + + if self.is_kda: + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in projection_names + } + attention_config._frozen = model_config._frozen + self.linear_attn = K3DecodeKDA( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + ) + + self.is_moe = ( + cfg.num_experts is not None + and layer_idx >= cfg.first_k_dense_replace + and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 + ) + if self.is_moe: + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's + # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take + # (decode_gemv.SITES); otherwise it stays within one node. + spans_nodes = self._mnnvl_allreduce() is not None + if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. + self.mlp = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, + ) + self._situ = (situ_beta, situ_linear_beta) + # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics + # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied + # after downcast, use_gemma=False). + self.input_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3). The res norms stay + # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by + # _apply_attn_res and the fused attn_res op, never called as + # modules. + self.attn_res_block_size = cfg.attn_res_block_size + assert self.attn_res_block_size is not None, ( + "Kimi K3 runtime expects attn_res_block_size to be set" + ) + self.self_attention_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set + # by the target's post_load_weights. + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + self.sandwich_oproj = False + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + step: DecodeStep, + capture: Optional[Tuple[Any, int]] = None, + prenormed: bool = False, + pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, + defer_moe_tail: bool = False, + ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + ``block_residual`` is a preallocated snapshot bank in kernel-native + ``[K_max, M, H]`` layout. Returns the running prefix sum and the + number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. + + ``step`` is the step's classification (``decode_step``), handed to the + attention module. + + ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is + this layer's input norm, and the layer's input, the step's embedding, + already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). + + The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens + runs the attention's all-reduce and the residual update in one + collective once the target has built its ``decode_comm``: with o_proj + (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, + the step and the attention's decode branch, else on the attention's + unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). + + ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the + previous layer's MoE output, unreduced (``defer_moe_tail``); then + ``hidden_states`` is the prefix sum without it, and this layer's + pre-attention step reduces and adds it: a ``PendingTail`` in one + ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its + all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer + the ``sandwich_tail`` call also stores the tap into the capture slot + where the speculative metadata exposes it (``capture_view``); other + tapped layers take the tap after the step. + + ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` + instead, ``partial`` this layer's MoE output unreduced, for the next + consumer to reduce. + """ + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + attn_res_max_tokens = _attn_res_max_tokens(step) + tail = ( + pending_moe_partial + if isinstance(pending_moe_partial, _decode_comm.PendingTail) + else None + ) + # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum + # straight into the bank row it snapshots. + snapshot_row = None + if tail is not None and self.layer_idx % self.attn_res_block_size == 0: + snapshot_row = block_residual[num_snapshots] + # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the + # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). + tap_view = None + if tail is not None and capture is not None: + tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) + + if prenormed: + assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 + assert pending_moe_partial is None and capture is None + elif tail is not None: + hidden_states, prefix_sum = self.decode_comm.sandwich_tail( + tail, + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + updated_out=snapshot_row, + tap=tap_view, + tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, + ) + elif pending_moe_partial is not None: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + elif capture is not None: + # The mixture tap needs the PRE-norm value, which the fused + # attn-res + RMSNorm kernel does not expose. Keep the two steps + # split on captured layers only and fuse everywhere else. + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if capture is not None and pending_moe_partial is not None and tap_view is None: + # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output + # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. + tapped = ( + _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + if _AUX_ATTN_RES_STREAM_ENABLED + else prefix_sum + ) + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + + if self.layer_idx % self.attn_res_block_size == 0: + if not prenormed and snapshot_row is None: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, step): + if ( + self.sandwich_oproj + and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS + and attention.will_run_decode_branch(attn_metadata, step) + ): + core = attention(hidden_states, attn_metadata, step=step, project_output=False) + hidden_states, prefix_sum = comm.sandwich_oproj( + core, + self._o_proj().weight, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + if comm is not None and step.wide: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) + else: + attention_output = attention(hidden_states, attn_metadata, step=step) + if prefix_sum is None: + prefix_sum = attention_output + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention_output, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + if self.is_moe and defer_moe_tail: + partial = self.block_sparse_moe( + hidden_states, all_rank_num_tokens, partial_tail=True, step=step + ) + return prefix_sum, num_snapshots, partial + if self.is_moe: + hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) + else: + hidden_states = self._dense_mlp(hidden_states, step) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def forward_generic( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + ) -> Tuple[torch.Tensor, int]: + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + if capture is not None: + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, None): + partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + elif prefix_sum is None: + prefix_sum = attention.forward_generic(hidden_states, attn_metadata) + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention.forward_generic(hidden_states, attn_metadata), + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + if self.is_moe: + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + hidden_states = self.block_sparse_moe.forward_generic( + hidden_states, all_rank_num_tokens + ) + else: + hidden_states = self.mlp(hidden_states) + return prefix_sum + hidden_states, num_snapshots + + def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: + """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode + kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" + gemvs = self.decode_gemvs + if gemvs is not None and step.small: + out = gemvs.dense_mlp( + hidden_states, + self.mlp.gate_up_proj.weight, + self.mlp.down_proj.weight, + *self._situ, + ) + if out is not None: + return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out + return self.mlp(hidden_states) + + def _mnnvl_allreduce(self): + """The MNNVL all-reduce of this layer's attention output, or None.""" + attention = self.linear_attn if self.is_kda else self.self_attn + return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) + + def _o_proj(self) -> nn.Module: + """This layer's attention output projection (row parallel).""" + return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj + + def _o_allreduce(self) -> AllReduce: + """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" + return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce + + def accepts_moe_partial(self, num_snapshots: int) -> bool: + """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target + built its collective state, and the snapshot bank is not empty.""" + return num_snapshots > 0 and self.decode_comm is not None + + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + + +class KimiLinearModel(DecoderModel): + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = _get_text_config(model_config.pretrained_config) + self._text_cfg = cfg + dtype = torch.bfloat16 + + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream() for _ in range(4)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + } + + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + + # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) + # by _apply_attn_res and the fused attn_res op. + self.output_attn_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.num_attn_res_snapshots = ( + cfg.num_hidden_layers + cfg.attn_res_block_size - 1 + ) // cfg.attn_res_block_size + # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the + # weights are final. None: every step embeds on the generic path. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + + @property + def kda_draft_records(self) -> bool: + """Whether the hybrid cache manager keeps the last verify's draft records (``kda_state_tok``), from which + ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn`` replay the accepted drafts. The engine reads it once the weights + are loaded, to build the manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer + taking the K3 kernels. Otherwise the KDA verify replays the accepted drafts from the replay caches (the + built-in verify).""" + spec_config = getattr(self.model_config, "spec_config", None) + return bool( + spec_config is not None + and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) + and spec_config.tokens_per_gen_step in (2, 4, 6, 8) + and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) + ) + + def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, + from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from + load-time state, so every rank and every replay of the graph decides alike.""" + if not torch.cuda.is_current_stream_capturing(): + return False + kda = [layer.linear_attn for layer in self.layers if layer.is_kda] + mla = [layer.self_attn for layer in self.layers if not layer.is_kda] + return latent_push( + step, + capturing=True, + breakable=is_in_breakable_cuda_graph(), + kda_draft_records=self.kda_draft_records, + kda_decode_kernels=all(m.takes_k3_kernels and m.k3_buffers is not None for m in kda), + mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), + ) + + def _defer_moe_tail( + self, + i: int, + hidden_states: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + step: DecodeStep, + ) -> bool: + """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE + takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after + the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps + its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated + tail.""" + layer = self.layers[i] + if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): + return False + if spec_metadata is not None and ( + capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) + ): + return False + return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] + step = decode_step(attn_metadata, num_tokens) + if step is None: + return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) + if self._latent_push(attn_metadata, step): + step = replace(step, latent_push=True) + # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. + prenormed = None + if ( + inputs_embeds is None + and self.decode_gemvs is not None + and len(self.layers) > 0 + and self.num_attn_res_snapshots > 0 + ): + table = self.embed_tokens.weight + block_residual = table.new_empty( + self.num_attn_res_snapshots, num_tokens, table.shape[1] + ) + prenormed = self.decode_gemvs.embed_norm( + input_ids, table, self.layers[0].input_layernorm, block_residual + ) + if prenormed is not None: + hidden_states = prenormed + else: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) + # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. + pending_moe_partial = None + for i, layer in enumerate(self.layers): + defer_moe_tail = self._defer_moe_tail( + i, hidden_states, num_snapshots, spec_metadata, capture_set, step + ) + outputs = layer( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + step=step, + capture=self._capture(spec_metadata, capture_set, i), + prenormed=i == 0 and prenormed is not None, + pending_moe_partial=pending_moe_partial, + defer_moe_tail=defer_moe_tail, + ) + if defer_moe_tail: + hidden_states, num_snapshots, pending_moe_partial = outputs + else: + (hidden_states, num_snapshots), pending_moe_partial = outputs, None + + if isinstance(pending_moe_partial, _decode_comm.PendingTail): + normed, _ = self.layers[-1].decode_comm.sandwich_tail( + pending_moe_partial, + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + return normed + if pending_moe_partial is not None: + _, normed = _apply_attn_res_add_and_rmsnorm( + hidden_states, + _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + _attn_res_max_tokens(step), + ) + return normed + return self._final_norm( + hidden_states, + block_residual, + num_snapshots, + spec_metadata, + capture_set, + _attn_res_max_tokens(step), + ) + + def _generic_forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor], + inputs_embeds: Optional[torch.FloatTensor], + spec_metadata, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + for i, layer in enumerate(self.layers): + hidden_states, num_snapshots = layer.forward_generic( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + capture=self._capture(spec_metadata, capture_set, i), + ) + return self._final_norm( + hidden_states, block_residual, num_snapshots, spec_metadata, capture_set + ) + + def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + return (spec_metadata, self.layers[i - 1].layer_idx) + return None + + def _final_norm( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + max_fused_tokens: Optional[int] = None, + ) -> torch.Tensor: + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) + + return _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + max_fused_tokens, + ) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. +# ---------------------------------------------------------------------------------------------------------------------- + + +class K3DecodeKDA(KimiKDALinearAttention): + """Kimi K3's KDA attention: the built-in module, with the decode kernels on the steps they take. + + * A decode step of one token per request runs the fused input projection and the plain decode in one + ``ssm/k3_kda_decode_attn`` launch. + * With the cache manager's per-token states (``KimiLinearModel.kda_draft_records``), every verify of the layer, on + any step, runs the kernels that keep them: ``ssm/k3_kda_attn`` for one request of 8 tokens (the projection fused + in), else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it takes + them). The built-in verify replays drafts from caches these kernels do not fill, so the two never run on one + manager. + * On every step ``decode_step`` classifies, ``o_proj`` runs on the ``o_proj`` decode GEMV site where it takes the + rows, then the module's all-reduce. + + Every other step runs the built-in module. The kernels read one ``[q | k | v | g | f_a | b]`` weight, built at the + checkpoint load from the module's own, and the device's ``K3KdaBuffers`` and decode GEMVs' state, which the + target sets in ``post_load_weights``. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # The fused [q | k | v | g | f_a | b | pad] projection weight; the six projections' weights and the built-in + # fused [q | k | v | g] and [f_a | b | pad] ones are views of it. + self.k3_proj_weight: Optional[torch.Tensor] = None + # The fused projection's Lamport buffers: one set per device, shared by every KDA layer. + self.k3_buffers: Optional[K3KdaBuffers] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @property + def takes_k3_kernels(self) -> bool: + """Whether the decode kernels can run this layer: its fused projection weight is built.""" + return self.k3_proj_weight is not None + + def finalize_decode_weights(self) -> None: + """The built-in fused weights, then one ``[q | k | v | g | f_a | b | pad]`` weight of both.""" + super().finalize_decode_weights() + assert ( + self.use_full_rank_gate + and self.gate_lower_bound is not None + and self._qkvg_proj_weight is not None + and self._bfa_proj_weight is not None + and self._qkvg_proj_weight.dtype == self._bfa_proj_weight.dtype == torch.bfloat16 + ), ( + f"Kimi K3 KDA layer {self.layer_idx}: the decode kernels read the bf16 fused projections the built-in " + "module builds on CUDA at head dim 128, with a full-rank output gate and a gate lower bound" + ) + rows = self._qkvg_proj_weight.shape[0] + with torch.no_grad(): + fused = self._merge_projection_weights( + (self.q_proj, self.k_proj, self.v_proj, self.g_proj, self.f_a_proj, self.b_proj), + pad_rows_to=8, + ) + self.k3_proj_weight = fused + self._qkvg_proj_weight, self._bfa_proj_weight = fused[:rows], fused[rows:] + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode branch: outside a breakable CUDA graph.""" + return not is_in_breakable_cuda_graph() + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic`` where ``will_run_decode_branch`` does not hold. On the decode branch: the plain decode + on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the built-in dispatch; then + ``o_proj`` on its decode GEMV site and the TP all-reduce. + + ``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce), and, on the decode branch only, + ``project_output=False`` the post-o_norm core ``[N, H * 128]`` (no ``o_proj``).""" + if not self.will_run_decode_branch(attn_metadata, step): + return self.forward_generic(hidden_states, attn_metadata, reduce_output) + if ( + step.decode + and step.tokens_per_request == 1 + and self.takes_k3_kernels + and self.k3_buffers is not None + ): + core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) + else: + core = self._forward_impl(hidden_states, attn_metadata) + if not project_output: + return core.reshape(-1, self.proj_size) + return self._k3_project_output(core, reduce_output) + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + if reduce_output: + return super().forward(hidden_states, attn_metadata) + return self.o_proj( + self._builtin_core(hidden_states, attn_metadata).reshape(-1, self.proj_size) + ) + + def _builtin_core( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """The built-in forward's core ``[N, H, 128]``: inside a breakable CUDA graph from the eager + ``maybe_bcg_kda_core_inplace``, else from ``_forward_impl``.""" + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + return core + return self._forward_impl(hidden_states, attn_metadata) + + def _k3_project_output(self, core: torch.Tensor, reduce_output: bool = True) -> torch.Tensor: + """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then, with + ``reduce_output``, the TP all-reduce.""" + core2d = core.reshape(-1, self.proj_size) + out = None + if self.decode_gemvs is not None: + out = self.decode_gemvs.project("o_proj", core2d, self.o_proj.weight) + if out is None: + if reduce_output: + return self._project_output(core) + out = self.o_proj(core2d) + if reduce_output and self._o_allreduce is not None: + out = self._o_allreduce(out) + return out + + def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: + """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; + each slot's conv window and state advance in place.""" + mamba_metadata = attn_metadata.mamba_metadata + slots = getattr(mamba_metadata, "generation_state_indices", None) + if slots is None: + slots = mamba_metadata.state_indices[: x.shape[0]] + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + w_q, w_k, w_v = self._get_mtp_conv_weights() + core = k3_kda_decode_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.conv, + layer_cache.temporal, + slots, + self.k3_buffers, + float(self.gate_lower_bound), + self.head_k_dim**-0.5, + float(self.o_norm.eps), + ) + # Speculative decoding's replay caches keep their committed conv window in step with the pool's. + self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) + return core + + def forward_verify( + self, + x2d, + num_steps, + layer_cache, + conv_pool, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The built-in verify, unless the cache manager keeps the per-token states; then the decode kernels'.""" + if not (layer_cache.has_kda_replay_caches and layer_cache.kda_state_tok is not None): + return super().forward_verify( + x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices, output=output + ) + assert self.takes_k3_kernels and self.k3_buffers is not None, ( + f"Kimi K3 KDA layer {self.layer_idx}: the cache manager keeps the per-token verify states, which only " + "the decode kernels write, and this layer has no fused projection weight or buffers" + ) + core = self._k3_verify(x2d, num_steps, layer_cache, ssm_pool, slot_indices) + return self._store_core(core, output) + + def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch.Tensor: + """The core output ``[N T, H, 128]`` of N requests of T verify tokens: ``ssm/k3_kda_attn`` for one request of + 8 tokens, else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it + takes them). Each slot's state after the golden token, its drafts' states and its conv window are written in + place.""" + num_requests = x.shape[0] // num_steps + slots = slot_indices[:num_requests] + w_q, w_k, w_v = self._get_mtp_conv_weights() + constants = (float(self.gate_lower_bound), self.head_k_dim**-0.5, float(self.o_norm.eps)) + if num_requests == 1 and num_steps == 8: + out = k3_kda_attn( + x.contiguous(), + self.k3_proj_weight, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + self.k3_buffers, + num_steps - 1, + *constants, + ) + else: + rows = None + if self.decode_gemvs is not None: + rows = self.decode_gemvs.project("kda_proj", x, self.k3_proj_weight) + if rows is None: + rows = torch.nn.functional.linear(x, self.k3_proj_weight) + out = k3_kda_verify( + rows, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + num_steps - 1, + *constants, + ) + return out.view(-1, self.num_heads, self.head_dim) + + +class K3DecodeMLA(KimiK3MLAAttention): + """Kimi K3's MLA attention: the built-in module, with a decode step's attention on the decode kernels. + + A decode step runs ``x [W_a; W_g]^T`` with the gate columns through a sigmoid (the ``mla_ag`` decode GEMV site + where it takes the rows, else one GEMM), then ``attention/k3_mla_qkv`` (the q_a / kv_a RMSNorms, q_b and the k_b + absorption into the fused query, the step's latent rows stored into the paged cache) and + ``attention/k3_mla_attn_vb_out`` (the attention over the paged cache, v_b and the output gate in one launch), then + ``o_proj`` (the ``o_proj`` site where it takes the rows, else the module). Every other step, and a decode step whose + cache the kernels do not read (``k3_mla_decode_view`` says why), runs the built-in module. + + ``[W_a; W_g]`` is built at load from the module's weights, which become views of it. The attention workspace is + the device's ``K3MlaAttnWorkspace``; it and the decode GEMVs' state are set by the target in + ``post_load_weights``. + """ + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + # [W_a; W_g]: the fused q_a / kv_a projection's rows, then the output gate's. + self.k3_ag_weight: Optional[torch.Tensor] = None + # The decode attention's workspace: one per device, shared by every MLA layer. + self.k3_workspace: Optional[K3MlaAttnWorkspace] = None + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + def post_load_weights(self) -> None: + """The built-in post-load, then ``[W_a; W_g]`` (once: CUDA graphs captured since read it).""" + super().post_load_weights() + if self.k3_ag_weight is not None: + return + gaps = self._k3_layout_gaps() + assert not gaps, ( + f"Kimi K3 MLA layer {self.layer_idx}: the decode kernels do not take {'; '.join(gaps)}" + ) + qkv_a, gate = self.kv_a_proj_with_mqa, self.g_proj + rows = qkv_a.weight.shape[0] + with torch.no_grad(): + fused = torch.cat([qkv_a.weight, gate.weight]) + qkv_a.weight = nn.Parameter(fused[:rows], requires_grad=False) + gate.weight = nn.Parameter(fused[rows:], requires_grad=False) + self.k3_ag_weight = fused + + def _k3_layout_gaps(self) -> list: + """What of this layer the decode kernels do not take (empty when they take all of it).""" + if not (self.use_output_gate and self.fuse_qkv_a_proj and not self.is_lite): + return ["a layer without the output gate or the fused q_a / kv_a projection"] + linears = (self.kv_a_proj_with_mqa, self.g_proj, self.q_b_proj, self.o_proj) + checks = ( + (not self.mapping.has_cp_helix(), "helix context parallelism"), + (not self.apply_rotary_emb and not self.llama_4_scaling, "RoPE or llama-4 scaling"), + (self.sparse_attn_hooks is None, "sparse attention"), + ( + self.kv_cache_dtype != "fp8_ds_mla" + and not getattr(self.mqa, "has_fp8_kv_cache", False) + and not getattr(self.mqa, "has_fp4_kv_cache", False), + "a quantized KV cache", + ), + ( + all(m.weight.dtype == torch.bfloat16 and m.bias is None for m in linears) + and self.k_b_proj_trans.dtype == self.v_b_proj.dtype == torch.bfloat16, + "projections other than bf16 and unbiased", + ), + ( + not getattr(self.q_a_layernorm, "is_nvfp4", False) + and not getattr(self.kv_a_layernorm, "use_gemma", False), + "an NVFP4 q_a norm or a Gemma kv_a norm", + ), + ( + self.num_heads_tp % 6 == 0 + and self.kv_lora_rank == 512 + and self.qk_rope_head_dim == 64 + and self.q_lora_rank == 1536, + f"{self.num_heads_tp} heads, latent {self.kv_lora_rank}, rope {self.qk_rope_head_dim}, " + f"q_lora {self.q_lora_rank}", + ), + ) + return [why for ok, why in checks if not ok] + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode kernels (see ``_k3_step_view``). Only there does it return + the gated attention output before ``o_proj``; the KV writes forbid running the attention twice, so a caller + that needs it asks first.""" + return self._k3_step_view(attn_metadata, step) is not None + + def _k3_step_view(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> Optional[dict]: + """The paged-cache view the decode kernels read on ``step``: a decode step, this layer's ``[W_a; W_g]`` and + workspace built, outside a breakable CUDA graph, and a cache ``k3_mla_decode_view`` takes. Else None.""" + if not ( + step.decode + and self.k3_ag_weight is not None + and self.k3_workspace is not None + and not is_in_breakable_cuda_graph() + ): + return None + return self._k3_decode_view(attn_metadata, step.num_tokens) + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic``, except on a decode step whose cache the decode kernels read + (``will_run_decode_branch``). There only, ``project_output=False`` returns the gated attention output + ``[M, H * 128]``, the input of ``o_proj``.""" + view = self._k3_step_view(attn_metadata, step) + if view is None: + if not project_output: + raise ValueError("project_output=False needs a step will_run_decode_branch takes") + return self.forward_generic(position_ids, hidden_states, attn_metadata) + x = hidden_states[: step.num_tokens].contiguous() + rows = self.kv_a_proj_with_mqa.weight.shape[0] + gemvs = self.decode_gemvs + ag = None if gemvs is None else gemvs.project("mla_ag", x, self.k3_ag_weight) + if ag is None: + ag = torch.nn.functional.linear(x, self.k3_ag_weight) + ag[:, rows:].sigmoid_() + fused_q = k3_mla_qkv( + ag, + self.q_a_layernorm.weight, + float(self.q_a_layernorm.variance_epsilon), + self.q_b_proj.weight, + self.k_b_proj_trans, + self.kv_a_layernorm.weight, + float(self.kv_a_layernorm.variance_epsilon), + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + ) + attn_output = self.create_output(x, 0) + k3_mla_attn_vb_out( + fused_q, + view["pool"], + view["row_stride"], + view["page_table"], + view["page_offset"], + view["seq_len"], + view["softmax_scale"], + self.v_b_proj, + attn_output, + self.k3_workspace, + gate=ag, + gate_col0=rows, + ) + if not project_output: + return attn_output + out = None if gemvs is None else gemvs.project("o_proj", attn_output, self.o_proj.weight) + if out is None: + out = self._project_output([attn_output], position_ids, attn_metadata, None) + return out + + def forward_generic( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + return super().forward(position_ids, hidden_states, attn_metadata) + + def _k3_decode_view(self, attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[dict]: + """The paged latent cache as the decode kernels read it this step, or None when they do not read it (the + reason is logged once).""" + view = k3_mla_decode_view(self.mqa, attn_metadata, num_tokens) + if isinstance(view, str): + logger.info_once( + f"Kimi K3 MLA: the built-in path for a decode step the decode kernels do not read ({view})", + key=f"k3_mla_decode_view_{view}", + ) + return None + return view + + +# ---------------------------------------------------------------------------------------------------------------------- +# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. +# ---------------------------------------------------------------------------------------------------------------------- + +#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at +#: max_draft_len 7. +DRAFTER_ATTN_SPLITS = frozenset( + {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} +) +# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k +# RMSNorm's epsilon and the NeoX RoPE's base. +_DRAFTER_HEADS = (6, 1) +_DRAFTER_HEAD_DIM = 64 +_DRAFTER_PAGE = 64 +_DRAFTER_EPS = 1e-5 +_DRAFTER_ROPE_BASE = 10000.0 + + +def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: + """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input + feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" + if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: + return None + width = in_features // tp_size + return tp_rank * width, (tp_rank + 1) * width + + +class K3FcSlice(nn.Module): + """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is + ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's + ranks is ``fc``'s output.""" + + def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: + super().__init__() + self.weight = nn.Parameter(weight, requires_grad=False) + self.start = start + self.end = end + + @classmethod + def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: + """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns + (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" + columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) + if columns is None: + return None + start, end = columns + return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the + GEMM reads the column block in place, through the rows' stride.""" + return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) + + +class K3DSparkDrafter(GQADSparkForCausalLM): + """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. + + A block the entries take runs, per layer: + + * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; + * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to + its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the + worker writes the accepted rows' context k / v before a block reads them; + * the output projection on ``drafter_o``, then the module's all-reduce; + * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the + module's all-reduce. + + The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state + (below); a projection whose site does not take its rows runs its module. Every other block runs the stock + ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or + normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under + CUDA-graph capture. + + Where the target hands the drafter the TP group's collective state (``use_decode_comm``): + + * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and + ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; + * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the + post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the + block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the + projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection + after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is + followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers + whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. + + Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the + stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); + the context k / v and the Markov head stay upstream code. + """ + + def __init__( + self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" + ) -> None: + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. + self._k3_take: Optional[bool] = None + # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes + # only these. + self._k3_attn_ran: set = set() + # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + # The zero residual of the split context projection's all-reduce, up to a decode step's rows. + self._k3_zero_rows: Optional[torch.Tensor] = None + # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms + # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. + self._k3_norms_fuse = False + self._k3_sandwich_forms: frozenset = frozenset() + + def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: + """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the + target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are + loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. + + * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they + split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` + sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. + * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and + RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose + shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not + compile is not used.""" + self.decode_comm = comm + self._k3_split_fc() + self._k3_norms_fuse = self._k3_norms_take_comm() + forms = set() + if self._k3_norms_fuse: + layers = self.model.layers + for form, swiglu, weights in ( + ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), + ("down", True, [layer.mlp.down_proj.weight for layer in layers]), + ): + if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( + weights[0], swiglu=swiglu + ): + forms.add(form) + self._k3_sandwich_forms = frozenset(forms) + logger.info( + "Kimi K3 DSpark drafter: residual adds and RMSNorms " + + ( + f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " + "mnnvl_fusion_allreduce)" + if self._k3_norms_fuse + else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" + ) + ) + + def _k3_norms_take_comm(self) -> bool: + """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and + post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no + quantized output), and every output and down projection is row parallel with its all-reduce.""" + hidden = self.config.hidden_size + + def plain(norm: nn.Module) -> bool: + weight = getattr(norm, "weight", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (hidden,) + and hasattr(norm, "variance_epsilon") + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "return_hp_output", True) + ) + + def reduces(linear: nn.Module) -> bool: + return ( + getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" + and getattr(linear, "reduce_output", False) + and getattr(linear, "all_reduce", None) is not None + ) + + return plain(self.model.norm) and all( + plain(layer.input_layernorm) + and plain(layer.post_attention_layernorm) + and reduces(layer.self_attn.o_proj) + and reduces(layer.mlp.down_proj) + for layer in self.model.layers + ) + + def _k3_tp_all_reduce(self) -> Optional[nn.Module]: + """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" + o_proj = self.model.layers[0].self_attn.o_proj + if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": + return None + return getattr(o_proj, "all_reduce", None) + + def _k3_split_fc(self) -> None: + """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero + residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is + a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP + all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" + fc = getattr(self, "fc", None) + if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): + return + mapping = self.model_config.mapping + weight = getattr(fc, "weight", None) + sliced = None + if ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and getattr(fc, "bias", None) is None + and self._k3_tp_all_reduce() is not None + ): + sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) + if sliced is None: + logger.info( + "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " + f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" + ) + return + self.fc = sliced + self._k3_zero_rows = weight.new_zeros( + MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] + ) + logger.info( + f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " + f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" + ) + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" + result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + self._k3_split_fc() + return result + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. + + With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group + with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode + step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then + ``hidden_norm``. Otherwise the stock projection.""" + fc = self.fc + if not isinstance(fc, K3FcSlice): + return super().project_target_hidden(hidden_states) + partial = fc(hidden_states.to(fc.weight.dtype)) + rows = partial.shape[0] + zeros = self._k3_zero_rows + if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): + normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) + return normed + if rows == 0: + return self.hidden_norm(partial) + return self.hidden_norm(self._k3_tp_all_reduce()(partial)) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class + docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows + ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows + as a view instead of gathering them.""" + keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) + if keys is None: + return super().dflash_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ) + out = self._k3_block_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ctx_rows_start, + ) + if not torch.cuda.is_current_stream_capturing(): + self._k3_attn_ran |= keys + return out + + def _k3_layers_take(self) -> bool: + """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head + layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, + non-causal unwindowed block attention, and no block convolution or post-attention gate.""" + if self._k3_take is None: + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + take = ( + self._use_fused_qk_norm_rope + and not self.has_block_conv + and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate + ) + for layer_idx, layer in enumerate(self.model.layers): + attn, mlp = layer.self_attn, layer.mlp + pos = getattr(attn, "pos_embd_params", None) + rope = getattr(pos, "rope", None) + norms = (attn.q_norm, attn.k_norm) + take = take and ( + (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS + and attn.head_dim == _DRAFTER_HEAD_DIM + and getattr(attn, "is_qk_norm", False) + and not getattr(attn, "use_gemma_rms_norm", True) + and rope is not None + and pos.is_neox + and getattr(pos, "mrope_section", None) is None + and rope.theta == _DRAFTER_ROPE_BASE + and all( + norm.variance_epsilon == _DRAFTER_EPS + and norm.weight.dtype == torch.bfloat16 + and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) + for norm in norms + ) + and attn.qkv_proj.bias is None + and attn.o_proj.bias is None + and mlp.activation is torch.nn.functional.silu + and mlp.swiglu_limit is None + and mlp.swiglu_alpha in (None, 1.0) + and mlp.swiglu_beta in (None, 0.0) + and mlp.gate_up_proj.bias is None + and mlp.down_proj.bias is None + and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) + ) + self._k3_take = bool(take) + logger.info( + f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " + + ( + "on k3_drafter_attn_qknorm and the drafter GEMV sites" + if self._k3_take + else "on the stock block forward (a layer the drafter entries do not take)" + ) + ) + return self._k3_take + + def _k3_block_keys( + self, + noise_embedding: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor], + ctx_page_table: Optional[torch.Tensor], + ) -> Optional[frozenset]: + """The attention's compile keys for this block when the drafter entries take it, else None.""" + if ( + self.dflash_attention_backend != "TRTLLM" + or ctx_kv_cache is None + or ctx_page_table is None + or noise_embedding.dtype != torch.bfloat16 + or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS + or is_in_breakable_cuda_graph() + or not self._k3_layers_take() + ): + return None + num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads + page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) + dense = ( + num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + ) + strides = set() + for layer_idx in range(len(self.model.layers)): + cache = ctx_kv_cache[layer_idx] + if ( + cache.dtype != torch.bfloat16 + or cache.dim() != 5 + or tuple(cache.shape[1:]) != page + or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) + or cache.stride(0) % 8 + ): + return None + strides.add(cache.stride(0)) + keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) + if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: + return None + return keys + + def _k3_block_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: torch.Tensor, + ctx_page_table: torch.Tensor, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + batch, block = noise_embedding.shape[:2] + rows = batch * block + ctx_len = num_ctx_per_req[:batch].to(torch.int32) + if ctx_rows_start is None: + page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) + else: + # The batch's page-table rows are one contiguous run: a view, no gather. + page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] + positions = query_positions.reshape(-1).contiguous() + hidden = noise_embedding.reshape(rows, -1) + layers = self.model.layers + comm = self._k3_fused_comm(hidden) + residual = None + if comm is not None: + # The fused all-reduces read the residual and return the updated one: the block's input serves as the + # first residual, uncopied. + residual = hidden + normed = layers[0].input_layernorm(hidden) + for layer_idx, layer in enumerate(layers): + attn = layer.self_attn + if comm is None: + if residual is None: + residual = hidden.clone() + normed = layer.input_layernorm(hidden) + else: + normed, residual = layer.input_layernorm(hidden, residual) + qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) + out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) + k3_drafter_attn_qknorm( + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + positions, + attn.q_norm.variance_epsilon, + attn.pos_embd_params.rope.theta, + ctx_kv_cache[layer_idx], + page_table, + ctx_len, + attn.num_heads, + attn.num_key_value_heads, + out, + ) + if comm is not None: + # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after + # the last layer the final norm. + next_norm = ( + layers[layer_idx + 1].input_layernorm + if layer_idx + 1 < len(layers) + else self.model.norm + ) + normed, residual = self._k3_project_norm( + comm, out, attn.o_proj, residual, layer.post_attention_layernorm + ) + normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) + continue + hidden = self._k3_project("drafter_o", out, attn.o_proj) + hidden, residual = layer.post_attention_layernorm(hidden, residual) + hidden = self._k3_mlp(layer.mlp, hidden) + if comm is not None: + return normed + out, _ = self.model.norm(hidden, residual) + return out + + def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: + """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its + all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace + holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" + comm = self.decode_comm + if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): + return None + return comm + + def _k3_project_norm( + self, + comm: _decode_comm.K3DecodeComm, + x: torch.Tensor, + linear: nn.Module, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, + the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the + projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` + where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" + if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( + x, linear.weight, residual, norm + ): + return comm.sandwich_plain(x, linear.weight, residual, norm) + gemvs = self.decode_gemvs + partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) + if partial is None: + partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_mlp_norm( + self, + comm: _decode_comm.K3DecodeComm, + mlp: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down + projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul + form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down + projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; + then ``comm/mnnvl_fusion_allreduce``.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + partial = None + if gate_up is not None: + if "down" in self._k3_sandwich_forms and comm.takes_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ): + return comm.sandwich_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ) + partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + if partial is None: + partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: + """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel + projection's all-reduce.""" + y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) + if y is None: + return linear(x) + return _row_parallel_reduce(linear, y) + + def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on + ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + down = ( + None + if gate_up is None + else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + ) + if down is None: + return mlp(x) + return _row_parallel_reduce(mlp.down_proj, down) + + +def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: + """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" + all_reduce = getattr(linear, "all_reduce", None) + if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: + return all_reduce(partial) + return partial + + +# ---------------------------------------------------------------------------------------------------------------------- +# The target: step classification, the construction checks and the class route returns. +# ---------------------------------------------------------------------------------------------------------------------- + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. + + The checkpoint names its language-model modules `language_model.`; the text model's are ``, with + `layers.*` under `model.`. + """ + config = model_config.pretrained_config + text = copy.copy(model_config) + text._frozen = False + text.pretrained_config = config.text_config + excluded = text.quant_config.exclude_modules + if excluded: + text.quant_config = copy.copy(text.quant_config) + renamed = [] + for name in excluded: + if name.startswith(_LANG_PREFIX): + name = name[len(_LANG_PREFIX) :] + if name.startswith("layers."): + name = "model." + name + renamed.append(name) + text.quant_config.exclude_modules = renamed + text.skip_create_weights_in_init = True + text._frozen = True + return text + + +@dataclass(frozen=True) +class DecodeStep: + """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` + generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" + + num_tokens: int + num_requests: Optional[int] = None + tokens_per_request: Optional[int] = None + #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` + #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. + latent_push: bool = False + + @property + def small(self) -> bool: + """Whether the step fits one token tile of the token-count kernels.""" + return self.num_tokens <= DECODE_MAX_TOKENS + + @property + def decode(self) -> bool: + """Whether the step is a pure decode step the request-aware kernels take.""" + return self.num_requests is not None + + @property + def wide(self) -> bool: + """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode + layout's MoE head and tail, on M-general ops.""" + return self.decode and not self.small + + +def latent_push( + step: DecodeStep, + *, + capturing: bool, + breakable: bool, + kda_draft_records: bool, + kda_decode_kernels: bool, + mla_decode_branch: bool, +) -> bool: + """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at + most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a + breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels + (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA + layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). + Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel + between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a + graph launch orders every pushing replay behind whatever ran before it.""" + if not (step.decode and step.small) or not capturing or breakable: + return False + if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): + return False + return mla_decode_branch + + +def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: + """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). + + ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers + and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input + here is fixed by that shape, so a captured step and its replays are classified alike. + """ + if num_tokens <= 0: + return None + requests = _decode_requests(attn_metadata, num_tokens) + if requests is not None: + return DecodeStep(num_tokens, requests, num_tokens // requests) + if num_tokens <= DECODE_MAX_TOKENS: + return DecodeStep(num_tokens) + return None + + +def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: + """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" + if attn_metadata.num_contexts != 0: + return None + requests = attn_metadata.num_generations + if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: + return None + tokens = num_tokens // requests + if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: + return None + seq_lens = getattr(attn_metadata, "seq_lens", None) + if seq_lens is not None and seq_lens.device.type == "cpu": + lens = seq_lens[:requests] + if lens.numel() != requests or bool((lens != tokens).any()): + return None + return requests def _check_construction(model_config: ModelConfig) -> None: - _check_target_construction(model_config, _TARGET) + """The settings this target is built for that are fixed before the first step.""" + capability = torch.cuda.get_device_capability() + assert capability == _SM, ( + f"this target is certified on sm_{_SM[0]}{_SM[1]}, running on sm_{capability[0]}{capability[1]}" + ) + mapping = model_config.mapping + topology = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.enable_attention_dp, + ) + assert topology == (16, 16, 1, 4, 4, False), ( + "the tp16_moetp4ep4 target needs world_size 16, tensor_parallel_size 16, pipeline_parallel_size 1, " + "moe_tensor_parallel_size 4, moe_expert_parallel_size 4 and enable_attention_dp false; the engine built " + f"(world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" + ) + assert model_config.torch_dtype == torch.bfloat16, ( + f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" + ) + kv_algo = model_config.quant_config.kv_cache_quant_algo + assert kv_algo is None, ( + f"this target's MLA kernels read a bf16 KV pool; kv_cache_config.dtype resolved to {kv_algo}" + ) + quant_algo = model_config.quant_config.quant_algo + assert quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( + "this target loads the MXFP4 checkpoint, whose compressed-tensors config the model config reads as " + "W4A16_MXFP4 with no per-layer declarations (its routed experts run the W4A8_MXFP4_MXFP8 default); the " + f"engine read {quant_algo} with {len(model_config.quant_config_dict or {})} per-layer declarations" + ) + strategy = model_config.allreduce_strategy + assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( + f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" + ) + + +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(KimiLinearForCausalLM): + """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. + + It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through + weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model + keeps the built-in one's. + """ + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The V2 hybrid manager holds the KDA states; the step contract requires it.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + config = model_config.pretrained_config + assert getattr(config, "text_config", None) is not None, ( + "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " + "text_config" + ) + _check_construction(model_config) + text = _text_model_config(model_config) + spec_config = getattr(text, "spec_config", None) + assert ( + spec_config is None + or spec_config.spec_dec_mode.is_sa() + or spec_config.spec_dec_mode.is_dflash() + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" + # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 + # weight-read conversion of the shared / latent MLPs. + self._fp8_weight_read_moe_mlp = False + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + cfg = text.pretrained_config + SpecDecOneEngineForCausalLM.__init__( + self, + KimiLinearModel(text), + text, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self._own_spec_branch(text) + self._step_checked = False + # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the + # processor the speculative worker's target logits and a parallel drafter's logits call too. + stock_logits_processor = self.logits_processor + self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) + draft_model = getattr(self, "draft_model", None) + if getattr(draft_model, "logits_processor", None) is stock_logits_processor: + draft_model.logits_processor = self.logits_processor + # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which + # must therefore be the text config, as the built-in wrapper leaves it. + model_config._frozen = False + model_config.pretrained_config = self.config + model_config._frozen = True + + def _own_spec_branch(self, model_config: ModelConfig) -> None: + """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a + standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is + `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" + spec_config = model_config.spec_config + if spec_config is None: + return + mode = spec_config.spec_dec_mode + standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) + if ( + standalone_dspark + and self.draft_config is not None + and not is_mla(self.draft_config.pretrained_config) + ): + drafter = K3DSparkDrafter( + self.draft_config, dflash_attention_backend=spec_config.attention_backend + ) + if mode.is_parallel_draft(): + drafter.logits_processor = self.logits_processor + self.epilogue[self.epilogue.index(self.draft_model)] = drafter + self.draft_model = drafter + if mode.is_dflash(): + worker_cls = _spec_worker.KimiK3DFlashWorker + elif standalone_dspark: + worker_cls = _spec_worker.KimiK3DSparkWorker + else: + return + worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) + worker.set_draft_model(self.draft_model) + self.epilogue[self.epilogue.index(self.spec_worker)] = worker + self.spec_worker = worker + + def load_weights(self, weights, *args, **kwargs): + _weights.load(self, weights) + + def cache_derived_state(self) -> None: + """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of + every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: + a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites + are among them, and the drafter gets the state too.""" + super().cache_derived_state() + if self.model.decode_gemvs is not None: + return + gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head) + self.model.decode_gemvs = gemvs + self.logits_processor.gemvs = gemvs + for layer in self.model.layers: + if not layer.is_moe: + layer.decode_gemvs = gemvs + if isinstance(self.draft_model, K3DSparkDrafter): + self.draft_model.decode_gemvs = gemvs + + def post_load_weights(self) -> None: + """The state the decode kernels share, built once per device before any CUDA-graph capture and handed to + every layer of its kind: the KDA projection's Lamport buffers and the MLA decode attention's workspace; the + decode GEMVs' state (built by ``cache_derived_state``) handed to every attention module; and, where every + attention all-reduce runs over MNNVL, the TP group's collective state (``K3DecodeComm``, collective: every + rank builds it here) handed to every layer, with whether its sandwich takes the layer's o_proj, the decode + path's one-shot ceiling on every stock MNNVL all-reduce (``use_decode_one_shot``), and the MoE decode path + (``_build_decode_moe``); then the speculative worker's decode kernels (``_gate_spec_worker_kernels``) and the + DSpark drafter's collectives (``_gate_drafter_comm``).""" + super().post_load_weights() + kda = [layer.linear_attn for layer in self.model.layers if layer.is_kda] + mla = [layer.self_attn.mixer for layer in self.model.layers if not layer.is_kda] + if kda[0].k3_buffers is not None: + return # built by an earlier call; CUDA graphs captured since hold it + device = self.model.embed_tokens.weight.device + buffers = K3KdaBuffers.create(device) + workspace = K3MlaAttnWorkspace.create(device, mla[0].num_heads_tp // 6) + for module in kda: + module.k3_buffers = buffers + for module in mla: + module.k3_workspace = workspace + for module in kda + mla: + module.decode_gemvs = self.model.decode_gemvs + layers = self.model.layers + comm = None + if all(layer._mnnvl_allreduce() is not None for layer in layers): + for layer in layers: + layer.sandwich_oproj = _decode_comm.K3DecodeComm.takes_oproj( + layer._o_proj(), self.model.num_attn_res_snapshots + ) + oproj = next((layer._o_proj().weight for layer in layers if layer.sandwich_oproj), None) + comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, oproj) + for layer in layers: + layer.decode_comm = comm + _decode_comm.use_decode_one_shot(self) + moe_layers = 0 if comm is None else self._build_decode_moe(comm) + logger.info( + "Kimi K3 decode kernels: KDA on k3_kda_decode_attn, k3_kda_attn and k3_kda_verify " + f"({sum(m.takes_k3_kernels for m in kda)} / {len(kda)} layers take them), MLA on k3_mla_qkv and " + f"k3_mla_attn_vb_out ({len(mla)} layers); the post-attention all-reduce of at most " + f"{_decode_comm.AR_ATTN_RES_MAX_TOKENS} tokens " + + ( + "with the residual update: k3_sandwich_oproj with o_proj at most " + f"{_decode_comm.SANDWICH_MAX_TOKENS} decode tokens " + f"({sum(layer.sandwich_oproj for layer in layers)} / {len(layers)} layers), else " + "mnnvl_allreduce_attn_res" + if comm is not None + else "unfused (an attention all-reduce does not run over MNNVL)" + ) + + f"; MoE on k3_moe_front, k3_moe and the row-parallel tail ({moe_layers} layers)" + ) + self._gate_spec_worker_kernels(comm) + self._gate_drafter_comm(comm) + + def _build_decode_moe(self, comm: _decode_comm.K3DecodeComm) -> int: + """The MoE decode path (``decode_moe.py``) on every MoE layer it takes: the shared state (collective: the + front's head workspace), each layer's decode weights with its latent norm folded into the latent up + projection, the decode GEMVs' state; then one call of every kernel of the path (collective: the front and + the sandwich tail). Every rank builds it here. Returns the number of MoE layers it takes.""" + mapping = self.model_config.mapping + moes = [layer.block_sparse_moe for layer in self.model.layers if layer.is_moe] + gaps = { + id(moe): _decode_moe.layout_gaps( + moe, mapping.tp_size, self.model.num_attn_res_snapshots + ) + for moe in moes + } + takes = [moe for moe in moes if not gaps[id(moe)]] + for reason in sorted({gap for moe in moes for gap in gaps[id(moe)]}): + logger.info_once( + f"Kimi K3 MoE decode path off on some layers: {reason}", + key=f"k3_decode_moe_off_{reason}", + ) + if not takes: + return 0 + backend = takes[0].routed_experts.backend + state = _decode_moe.K3DecodeMoe.create( + mapping, + backend.w3_w1_weight.device, + backend.w3_w1_weight.shape[1] // 2, + backend.expert_size_per_partition, + comm.mnnvl, + ) + for moe in takes: + _decode_moe.fold_latent_norm(moe) + moe.decode_moe = _decode_moe.K3DecodeMoeLayer.create( + moe, state, mapping.tp_rank, mapping.tp_size + ) + moe.decode_gemvs = self.model.decode_gemvs + first = takes[0].decode_moe + first.warm_up(takes[0]) + logger.info( + "Kimi K3 MoE decode path: the latent all-reduce at <= 8 tokens on " + + ( + "the routed experts' push form and k3_latent_reduce on pushing steps (a pure decode step captured " + "into a CUDA graph on the decode kernels), else the routed experts' all-reduce" + if state.exchange is not None + else f"the routed experts' all-reduce (no latent exchange at TP {mapping.tp_size})" + ) + ) + comm.compile_tail(takes[0].moe_hidden_size, first.shared_cols, first.tail_weight) + return len(takes) + + def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, + ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input + that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's + ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the + vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks + each step's own conditions. Returns the setting (False without such a worker).""" + worker = getattr(self, "spec_worker", None) + if not hasattr(worker, "k3_decode"): + return False + gemvs = self.model.decode_gemvs + worker.k3_decode = ( + comm is not None and gemvs is not None and gemvs.head_workspace is not None + ) + logger.info( + f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " + + ( + "on" + if worker.k3_decode + else "off (no MNNVL decode state or no k3_head_gemv LM head)" + ) + ) + return worker.k3_decode + def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built + it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context + projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms + in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on + the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate + requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` + workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the + drafter took the state (False without such a drafter).""" + drafter = getattr(self, "draft_model", None) + if comm is None or not isinstance(drafter, K3DSparkDrafter): + return False + drafter.use_decode_comm(comm) + return True -class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(ModelingV2KimiK3): - """Concrete deployment of the shared Kimi K3 text model.""" + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """First-forward checks of the engine surface and the per-engine settings.""" + objects = { + "attn_metadata": attn_metadata, + "kv_cache_manager": attn_metadata.kv_cache_manager, + } + missing = [ + f"{owner}.{name}" + for owner, names in REQUIRED_ENGINE_FIELDS.items() + for name in names + if not hasattr(objects[owner], name) + ] + assert not missing, f"engine fields this target reads are missing: {missing}" + assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( + f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " + f"{attn_metadata.tokens_per_block}" + ) + manager = attn_metadata.kv_cache_manager + assert isinstance(manager, MambaHybridCacheManagerV2), ( + "this target needs the V2 hybrid KV / state cache manager " + "(kv_cache_config.use_kv_cache_manager_v2); the engine built " + f"{type(manager).__name__}" + ) + assert not manager.enable_block_reuse, ( + "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" + ) + kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) + state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype + assert state_dtype == torch.float32, ( + "this target's KDA kernels keep fp32 recurrent states (kv_cache_config.mamba_ssm_cache_dtype float32 or " + f"auto); the engine built a {state_dtype} state pool" + ) + self._step_checked = True - target_spec = _TARGET + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + if kwargs.pop("multimodal_params", None): + raise ValueError( + "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" + ) + if not self._step_checked: + self._check_step_contract(attn_metadata) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py new file mode 100644 index 000000000000..abc1a1388213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py @@ -0,0 +1,1343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import functools +import inspect +from typing import Optional + +import torch +from torch import nn + +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.speculative.dflash import ( + DFlashSpecMetadata, + DFlashWorker, + dflash_allocated_ctx_limit, + dflash_noise_block_embedding, +) +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase +from tensorrt_llm.logger import logger + + +@functools.lru_cache(maxsize=None) +def _takes_ctx_rows_start_cls(cls) -> bool: + return "ctx_rows_start" in inspect.signature(cls.dflash_forward).parameters + + +def _takes_ctx_rows_start(draft_model) -> bool: + """Whether the draft model's ``dflash_forward`` accepts ``ctx_rows_start``.""" + return _takes_ctx_rows_start_cls(type(draft_model)) + + +def dflash_ctx_rows_kwargs(draft_model, ctx_rows_start: Optional[int]) -> dict: + """The ``ctx_rows_start`` keyword (from ``prepare_1st_drafter_inputs``) for a draft model whose ``dflash_forward`` + takes it; no keyword for another draft model or without a start.""" + if ctx_rows_start is None or not _takes_ctx_rows_start(draft_model): + return {} + return {"ctx_rows_start": ctx_rows_start} + + +def capture_view( + spec_metadata: Optional[SpecMetadata], layer_id: int, num_tokens: int +) -> Optional[torch.Tensor]: + """``layer_id``'s slot of the DFlash capture buffer for ``num_tokens`` rows (a strided view a kernel can write the + tap into), or None when ``spec_metadata`` is not DFlash's or the layer is not captured.""" + if ( + not isinstance(spec_metadata, DFlashSpecMetadata) + or spec_metadata.captured_hidden_states is None + ): + return None + i = spec_metadata._layer_to_idx.get(layer_id) + if i is None: + return None + return spec_metadata.captured_hidden_states[ + :num_tokens, i * spec_metadata.hidden_size : (i + 1) * spec_metadata.hidden_size + ] + + +class KimiK3DFlashWorker(DFlashWorker): + # Set by a Kimi K3 target running its decode kernels: decode steps then run the acceptance and the drafter's + # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context + # K/V with trtllm::k3_ctx_kv. + k3_decode = False + + def _mask_token_id(self, draft_model) -> int: + """The drafter's mask token id, resolved once.""" + if self._resolved_mask_token_id is None: + if ( + hasattr(self.spec_config, "mask_token_id") + and self.spec_config.mask_token_id is not None + ): + self._resolved_mask_token_id = self.spec_config.mask_token_id + elif hasattr(draft_model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.mask_token_id + elif hasattr(draft_model.model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.model.mask_token_id + else: + raise ValueError( + "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " + "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." + ) + return self._resolved_mask_token_id + + @staticmethod + def _trained_mask_embedding(draft_model, mask_token_id: int) -> Optional[torch.Tensor]: + """The drafter's own trained mask row, if it kept one for the mask id in use.""" + if mask_token_id != getattr(draft_model, "mask_token_id", None): + return None + return getattr(draft_model, "mask_token_embedding", None) + + def _k3_accept_applies( + self, logits, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Whether this step's acceptance and the drafter's inputs run as ``trtllm::k3_spec_accept``. + + That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no + rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target + logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block + table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. + """ + if not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return False + K = spec_metadata.runtime_draft_len + if ( + logits.dim() != 2 + or logits.dtype != torch.float32 + or logits.shape[0] != num_gens * (K + 1) + ): + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + embed = draft_model.draft_model_full.model.embed_tokens + return accept_op.supports( + logits.shape[1], num_gens, self._compute_block_size, K, embed.weight.shape[1] + ) + + def _k3_accept_step_applies( + self, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Everything ``_k3_accept_applies`` checks but the target logits.""" + if not self.k3_decode: + return False + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + + if num_contexts != 0 or not 0 < num_gens <= 8 or self.guided_decoder is not None: + return False + if ( + not spec_metadata.is_all_greedy_sample + or self._can_use_rejection_sampling(spec_metadata) + or getattr(spec_metadata, "enable_penalty", False) + or type(self)._reshape_draft_tokens_for_accept + is not SpecWorkerBase._reshape_draft_tokens_for_accept + or type(self)._reshape_logits_for_accept + is not SpecWorkerBase._reshape_logits_for_accept + ): + return False + K = spec_metadata.runtime_draft_len + draft_tokens = spec_metadata.draft_tokens + if ( + K <= 0 + or draft_tokens is None + or draft_tokens.dtype != torch.int32 + or draft_tokens.numel() != num_gens * K + ): + return False + mgr = attn_metadata.kv_cache_manager + mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) + if ( + not isinstance(mgr, MambaHybridCacheManagerV2) + or type(mgr).update_mamba_states is not MambaHybridCacheManagerV2.update_mamba_states + or not mgr.use_kda_replay_update + or getattr(mgr, "prev_num_accepted_tokens", None) is None + or getattr(mgr, "_dummy_request_mask", None) is None + or mamba_metadata is None + or mamba_metadata.state_indices.dtype != torch.int32 + ): + return False + if ( + self._ctx_block_tables is None + or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None + or getattr(attn_metadata, "kv_lens_cuda", None) is None + ): + return False + embed = draft_model.draft_model_full.model.embed_tokens + return getattr(embed, "tp_size", 1) == 1 and embed.weight.dtype == torch.bfloat16 + + def target_logits( + self, hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) -> torch.Tensor: + """This rank's bf16 vocabulary shard of the target logits, [rows, vocab / TP] (``_k3_head_shard``: the head + GEMM without its all-gather and fp32 cast), for a step whose acceptance runs as the sharded + ``trtllm::k3_spec_accept`` (see ``_k3_logits_shard``); the gathered fp32 logits otherwise. At such a step the + acceptance is their only reader: the one-model spec sampler takes the worker's tokens and rejects requests for + logits or log probabilities, so the shard is also what the step returns as ``"logits"``.""" + self._k3_step_shard = None + shard = self._k3_logits_shard( + hidden_states, lm_head, attn_metadata, spec_metadata, draft_model + ) + if shard is None: + return super().target_logits( + hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) + logits = self._k3_head_shard(logits_processor, lm_head, hidden_states) + self._k3_step_shard = (logits, shard) + return logits + + @staticmethod + def _k3_head_shard(logits_processor, lm_head, rows: torch.Tensor) -> torch.Tensor: + """This rank's bf16 vocabulary shard of ``lm_head(rows)``, without the head's all-gather: from the logits + processor's own head kernel where it has one that takes the rows (``lm_head_shard``; the Kimi K3 target's + ``gemm/k3_head_gemv``), so the shard holds the values of the processor's gathered logits; else from the head's + own GEMM.""" + head_shard = getattr(logits_processor, "lm_head_shard", None) + logits = None if head_shard is None else head_shard(rows, lm_head) + if logits is None: + logits = lm_head.apply_linear(rows, lm_head.bias) + return logits + + def _k3_logits_shard(self, hidden_states, lm_head, attn_metadata, spec_metadata, draft_model): + """``(workspace, first column)`` when this step's target logits stay vocabulary-sharded, else None. + + That needs a step ``_k3_accept_step_applies`` takes, plain TP whose all-reduces own an MNNVL workspace for + this mapping, and an unquantized, bias-free, unpadded, evenly split column-parallel bf16 lm_head with shards + the kernel splits. The exchange is collective, so every rank must decide alike: every condition is the + batch's or the configuration's. The workspace is allocated (collectively) and the kernel compiled outside + CUDA-graph capture only, so a graph whose warmup allocated no workspace is captured on the gathered path. + """ + if not self.k3_decode: + return None + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + mapping = self.mapping + if getattr(self, "_k3_shard_head", None) is not lm_head: + reason = self._k3_logits_shard_reason(lm_head) + if reason is None and MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None: + # Not kept: re-checked on later steps, since the model's MNNVL workspace may not exist yet. + logger.info_once( + "DFlash: the target logits are all-gathered (the TP all-reduces are not MNNVL)", + key="dflash_k3_logits_no_mnnvl", + ) + return None + if mapping is not None and mapping.tp_size > 1: + if reason is None: + logger.info( + f"DFlash: decode-step target logits stay vocabulary-sharded ({lm_head.weight.shape[0]} " + f"columns per rank); trtllm::k3_spec_accept exchanges the row maxima, and a request's " + f"logits post-processors get them gathered" + ) + else: + logger.info_once( + f"DFlash: the target logits are all-gathered ({reason})", key=reason + ) + self._k3_shard_head, self._k3_shard_ok = lm_head, reason is None + if not self._k3_shard_ok: + return None + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + K = spec_metadata.runtime_draft_len + if hidden_states.shape[0] != num_gens * (K + 1) or not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return None + embed = draft_model.draft_model_full.model.embed_tokens + columns = lm_head.weight.shape[0] + if not accept_op.supports( + columns, num_gens, self._compute_block_size, K, embed.weight.shape[1], mapping.tp_size + ): + return None + if torch.cuda.is_current_stream_capturing(): + ws = accept_op.existing_workspace(mapping) + if ws is None: + return None + else: + ws = accept_op.workspace(mapping) + return ws, mapping.tp_rank * columns + + def _k3_logits_shard_reason(self, lm_head): + """Why the target logits cannot stay vocabulary-sharded under this configuration and head (None if they can, + given an MNNVL workspace for the mapping, which ``_k3_logits_shard`` checks).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + if mapping is None or mapping.tp_size <= 1: + return "no TP" + if mapping.enable_attention_dp or mapping.has_cp() or mapping.has_pp(): + return "attention DP, CP or PP" + if ( + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or getattr(lm_head, "gather_output_sizes", None) is not None + ): + return "the lm_head is not an evenly split, unpadded column-parallel head" + if ( + lm_head.bias is not None + or lm_head.has_any_quant + or lm_head.weight.dtype != torch.bfloat16 + ): + return "the lm_head has a bias, is quantized or is not bf16" + if not accept_op.supports_columns(lm_head.weight.shape[0], mapping.tp_size): + columns = lm_head.weight.shape[0] + return f"trtllm::k3_spec_accept cannot split {columns} columns over {mapping.tp_size} ranks" + return None + + def _k3_static_arange(self, lo: int, hi: int) -> torch.Tensor: + """``arange(lo, hi)`` (int64), built once outside CUDA-graph capture.""" + cache = self.__dict__.setdefault("_k3_aranges", {}) + t = cache.get((lo, hi)) + if t is None: + t = torch.arange(lo, hi, dtype=torch.long, device="cuda") + if not torch.cuda.is_current_stream_capturing(): + cache[(lo, hi)] = t + return t + + def _k3_mask_row(self, draft_model) -> torch.Tensor: + """The noise block's mask embedding, bf16 [hidden], cached: the drafter's own trained mask row if it kept one + for the mask id in use, else the embedding's row (``dflash_noise_block_embedding`` reads the same row).""" + row = getattr(self, "_k3_mask_embedding", None) + if row is None: + embed = draft_model.draft_model_full.model.embed_tokens + mask_token_id = self._mask_token_id(draft_model) + trained = self._trained_mask_embedding(draft_model, mask_token_id) + if trained is not None: + row = trained.to(embed.weight.dtype).reshape(-1).contiguous() + else: + row = embed.weight[mask_token_id].contiguous() + if torch.cuda.is_current_stream_capturing(): + return row + self._k3_mask_embedding = row + return row + + def _k3_ctx_pool_key(self): + """The identity of the pool the ctx cache is bound to now (KV cache estimation rebinds it to a new one).""" + return ( + tuple(t.data_ptr() for t in self._ctx_kv_buf) if self._ctx_kv_buf is not None else None + ) + + def _k3_ctx_kv_applies(self, draft_model, projected: torch.Tensor) -> bool: + """Whether ``trtllm::k3_ctx_kv`` writes this step's context K/V: the manager-bound paged pool with K and V, + bf16, the fused K/V weight without bias or context input norm, k_norm, NeoX RoPE from flashinfer's fp32 + cache, a context the op can split (up to 64 tokens: B <= 8 requests of K + 1 <= 8, ``ctx_op.pick_split``), one + allocation for every layer's pool. Checked once per token count and + bound pool; the pool view is kept for the bound pool only, so a rebind never leaves the kernel writing the + replaced pool.""" + pool_key = self._k3_ctx_pool_key() + if getattr(self, "_k3_ctx_pool_bound", None) != pool_key: + self._k3_ctx_pool_bound = pool_key + self._k3_ctx_pool = None + self._k3_ctx_kv_ok = {} + cache = self._k3_ctx_kv_ok + n = projected.shape[0] + ok = cache.get(n) + if ok is not None: + return ok + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctx_kv import op as ctx_op + from tensorrt_llm._torch.models import modeling_dflash + + reason = None + if not (self._ctx_paged and self._ctx_block_tables is not None): + reason = "the context pool is not the manager-bound paged pool" + elif projected.dtype != torch.bfloat16 or self._ctx_kv_buf[0].dtype != torch.bfloat16: + reason = "not bf16" + else: + if draft_model._fused_kv_weight is None: + draft_model._build_fused_kv_buffers() + view = ctx_op.pool_view(list(self._ctx_kv_buf)) + if ( + draft_model._fused_kv_bias is not None + or getattr(draft_model, "_input_ln_eps", None) is not None + ): + reason = "K/V bias or context input norm" + elif draft_model._k_norm_stacked is None or not draft_model._is_neox: + reason = "no k_norm or not NeoX RoPE" + elif ( + modeling_dflash._flashinfer_rope is None + or draft_model._get_cos_sin_cache().shape[1] != 64 + ): + reason = "RoPE is not flashinfer's full-rotary fp32 cache" + elif view is None or self._ctx_kv_buf[0].size(1) != 2: + reason = "the layers' pools are not K/V views of one allocation" + elif not ctx_op.pick_split( + draft_model._fused_kv_weight.shape[0], + projected.shape[1], + draft_model._num_kv_heads, + self._draft_tokens_per_req, + n, + projected.device, + ): + reason = f"{n} context tokens (at most 64) or the K/V weight shape" + else: + self._k3_ctx_pool = view + ok = reason is None + if not ok: + logger.info(f"DFlash: context K/V on the Python path ({reason})") + if not torch.cuda.is_current_stream_capturing(): + cache[n] = ok + return ok + + def _k3_ctx_kv( + self, draft_model, projected, ctx_positions, num_accepted, slots, rows + ) -> torch.Tensor: + """``trtllm::k3_ctx_kv``: every drafter layer's context K/V of this step into the paged pool and + ``_ctx_len += num_accepted`` (clamped); returns the context length each request may advertise.""" + if self._k3_ctx_pool is None or self._k3_ctx_pool_bound != self._k3_ctx_pool_key(): + raise RuntimeError("k3_ctx_kv: the context pool was rebound after the path was chosen") + flat, layer_off, page_stride, kv_stride, head_stride = self._k3_ctx_pool + return torch.ops.trtllm.k3_ctx_kv( + projected.contiguous(), + draft_model._fused_kv_weight, + draft_model._k_norm_stacked, + draft_model._get_cos_sin_cache(), + ctx_positions, + num_accepted, + self._ctx_len, + slots, + rows, + self._ctx_block_tables, + self._ctx_block_counts, + flat, + layer_off, + page_stride, + kv_stride, + head_stride, + draft_model._k_norm_eps, + self._max_ctx, + self._ctx_page_size, + self._compute_block_size, + draft_model._num_kv_heads, + ) + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Called with this step's acceptance when ``k3_spec_accept`` produced it (a hook for drafter families).""" + + def _k3_accept(self, logits, attn_metadata, spec_metadata, draft_model, shard=None): + """``trtllm::k3_spec_accept`` for a step ``_k3_accept_applies`` accepted: the acceptance, the block table, + the KDA replay record, kv_lens + 1 and the drafter's inputs in one launch. ``shard``: the exchange + workspace and first column when ``logits`` are this rank's vocabulary shard (``_k3_logits_shard``). + Returns (accepted tokens, num accepted, {rewind, bonus, qpos, cpos, noise}).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + num_gens = attn_metadata.num_seqs + K = spec_metadata.runtime_draft_len + mgr = attn_metadata.kv_cache_manager + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + force = 0.0 if is_warmup else float(self.force_num_accepted_tokens) + mode, _, _ = accept_op.force_mode(force, K) + if mode == accept_op._kernel_module().FORCE_FRAC: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + else: + dummy_rng = getattr(self, "_k3_dummy_rng", None) + if dummy_rng is None: + dummy_rng = self._k3_dummy_rng = ( + torch.zeros(4, dtype=torch.float32, device=logits.device), + torch.zeros(2, dtype=torch.int64, device=logits.device), + ) + pool, counter = dummy_rng + embed = draft_model.draft_model_full.model.embed_tokens + accepted, num_acc, rewind, bonus, qpos, cpos, noise = torch.ops.trtllm.k3_spec_accept( + logits, + spec_metadata.draft_tokens.reshape(num_gens, K), + attn_metadata.draft_kv_cache_block_offsets, + self._ctx_pool_idx, + self._ctx_block_divisor, + self._ctx_block_counts, + self._ctx_block_tables, + mgr.prev_num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices, + mgr._dummy_request_mask, + attn_metadata.kv_lens_cuda, + self._batch_to_slot, + self._ctx_len, + self._max_ctx, + embed.weight, + self._k3_mask_row(draft_model), + pool, + counter, + force, + self._compute_block_size, + *self._k3_shard_args(shard), + ) + self._on_acceptance(accepted, num_acc, attn_metadata, spec_metadata) + return ( + accepted, + num_acc, + dict(rewind=rewind, bonus=bonus, qpos=qpos, cpos=cpos, noise=noise), + ) + + @staticmethod + def _k3_shard_args(shard) -> tuple: + """``trtllm::k3_spec_accept``'s exchange arguments for a vocabulary-sharded step (none otherwise).""" + if shard is None: + return () + ws, first_column = shard + return ( + ws["uc"], + ws["mc"], + ws["flags"], + ws["rank"], + ws["slots"], + ws["push_copies"], + first_column, + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + # Set when target_logits kept this step's logits vocabulary-sharded (only k3_spec_accept reads them). + step_shard, self._k3_step_shard = getattr(self, "_k3_step_shard", None), None + shard = None + if step_shard is not None: + if step_shard[0] is not logits: + raise RuntimeError( + "DFlash: the target logits are not the vocabulary shard target_logits returned" + ) + shard = step_shard[1] + + raw_logits = logits + K = spec_metadata.runtime_draft_len + + if K == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + # Lazy init buffers and attach worker reference for prepare() + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + # A decode step of the supported mode runs the acceptance and the drafter's inputs as one kernel + # (trtllm::k3_spec_accept), which also decodes the block table below. + if shard is None: + fused_accept = self._k3_accept_applies( + logits, attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ) + elif self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + fused_accept = True + else: + raise RuntimeError( + "DFlash: the target logits are vocabulary-sharded, but the fused acceptance no longer applies" + ) + # Before any store: prefill and decode both address pages through it. + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + if not fused_accept: + self._refresh_ctx_block_tables(attn_metadata, batch_size) + + # Save context lengths so both warmup and a failed forward can roll + # back the in-place _ctx_len updates made during drafting. + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # Never allocate the snapshot while capturing a CUDA graph: the + # clone would live in the graph memory pool, and its replay-time + # writes could alias blocks reused by later captures. Rollback is + # only meaningful for eager/warmup forwards anyway; a failure + # during capture aborts the graph itself, and captured ops do not + # mutate _ctx_len until replay. + self._saved_ctx_len = self._ctx_len.clone() + self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_req_ctx_pos = dict(self._req_ctx_pos) + self._ctx_len_restore_pending = True + + self._execute_guided_decoder_if_present(logits) + + prep = None + if fused_accept: + # Saved before the kernel updates kv_lens (warmup also snapshots kv_lens_cuda here). + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + accepted_tokens, num_accepted_tokens, prep = self._k3_accept( + logits, attn_metadata, spec_metadata, draft_model, shard + ) + else: + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + # Opt-in acceptance recording (env-gated; eager-mode measurement + # runs only). Skipped for CUDA-graph batches (capture/replay/warmup + # use synthetic requests and forbid the host sync). + if ( + self._accept_stats is not None + and num_gens > 0 + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + ): + self._accept_stats.on_accept( + spec_metadata.request_ids[num_contexts:batch_size], + num_accepted_tokens[num_contexts:batch_size].tolist(), + ) + + if fused_accept: + # The kernel recorded the KDA replay acceptance and advanced kv_lens_cuda. + self._kv_rewind_amount = prep["rewind"] + self._kv_rewind_nc = num_contexts + self._kv_rewind_bs = batch_size + self._kv_rewind_pending = True + attn_metadata.update_for_spec_dec() + else: + # Update GDN/Mamba recurrent states to the accepted token's state. + if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) + + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] + + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + ) + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + self._rebuild_batch_to_slot(spec_metadata.request_ids) + + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + prep=prep, + ) + + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + ctx_kv_cache=inputs["ctx_kv_cache"], + ctx_page_table=inputs["ctx_page_table"], + **dflash_ctx_rows_kwargs(draft_model, inputs["ctx_rows_start"]), + ) + + # Gather K logits per gen request from the block outputs. + # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. + block_size = self._compute_block_size + gen_hidden_states = self._draft_block_hidden_states( + draft_model, hidden_states_out, num_gens, block_size, K + ) + gen_logits = self._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) + + # DFlash 2: replace the independent per-position picks with one + # coherent path through the block (absent for plain DFlash). + if getattr(draft_model, "has_candidate_selector", False): + gen_logits = self._apply_dflash2_selector( + draft_model, + gen_logits, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + spec_metadata, + ) + vocab_size = gen_logits.shape[-1] + + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + ) + + # Opt-in confidence-calibration recording (env-gated). The + # confidence values come from an optional provider on the + # recorder (None here -> no calibration rows collected; the + # DSpark confidence-scheduled verification MR supplies the + # real provider, see accept_stats.py). Guarded off for + # CUDA-graph batches and d2t vocab-mapped drafters. + if ( + self._accept_stats is not None + and self._accept_stats.confidence_provider is not None + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + and self._d2t is None + ): + self._accept_stats.record_draft_confidence( + spec_metadata.request_ids[num_contexts:batch_size], + draft_model, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + gen_draft_tokens, + ) + + else: + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + gen_vocab = vocab_size if num_gens > 0 else None + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + self._rollback_guided_decoder_after_verify(num_accepted_tokens) + + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + # Restore context lengths after warmup; real runs keep the updates. + if is_warmup: + self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_host() + self._ctx_len_restore_pending = False + + outputs = { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + if shard is not None: + # "logits" is this rank's vocabulary shard: the engine gathers it for any logits post-processor. + outputs["logits_vocab_shard"] = True + return outputs + + def _draft_block_hidden_states( + self, + draft_model, + hidden_states_out: torch.Tensor, + num_gens: int, + block_size: int, + num_draft_tokens: int, + ) -> torch.Tensor: + """The block-output rows that produce the K draft logits per gen request (``_draft_slot_ids``). + + The slot ids depend only on the shapes, so they are built once per shape, outside CUDA-graph + capture, and cached. When they are one contiguous run of rows (a single request, or the + shift_label convention with a block of K slots, whose requests' slots follow each other) the + rows are returned as a view, without a gather; otherwise they are gathered. + """ + rows = hidden_states_out.shape[0] + key = (num_gens, block_size, num_draft_tokens, rows) + cache = self.__dict__.setdefault("_draft_block_rows", {}) + entry = cache.get(key) + if entry is None: + ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. + ids = ids.clamp(max=rows - 1) + if torch.cuda.is_current_stream_capturing(): + return hidden_states_out[ids] + host = ids.tolist() + if host and host == list(range(host[0], host[0] + len(host))): + entry = host[0] + else: + entry = ids + cache[key] = entry + if isinstance(entry, int): + return hidden_states_out[entry : entry + num_gens * num_draft_tokens] + return hidden_states_out[entry] + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata, + ) -> torch.Tensor: + """fp32 logits of the block positions over the full vocabulary (the TP lm_head all-gathers its + shards). A drafter family whose refinement and draft sampler reduce across ranks themselves may + return this rank's vocabulary shard instead.""" + return draft_model.logits_processor( + gen_hidden_states, draft_model.lm_head, attn_metadata, True + ) + + def _ctx_rows_start(self, num_contexts: int, num_gens: int) -> Optional[int]: + """The first row of the gen requests' page-table rows where those rows are one contiguous run, else None. + + The manager's block table is keyed by batch position, so the gen requests' rows are [num_contexts, + num_contexts + num_gens): a drafter can view them instead of gathering ``ctx_cache_batch_idx``'s rows. The + private arena's table is keyed by slot, which need not be contiguous. + """ + if num_gens > 0 and self._ctx_block_tables is not None: + return num_contexts + return None + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DFlashSpecMetadata, + draft_model: nn.Module, + total_target_tokens: int = 0, + prep: Optional[dict] = None, + ): + """Prepare inputs for DFlash's draft forward. + + ``prep``: ``trtllm::k3_spec_accept``'s bonus tokens, positions and noise embedding for a decode step it + took (see ``_k3_accept``); computed here otherwise. + + For gen requests, builds: + - noise_embedding: token embeddings for [accepted + mask] tokens + - query_positions: position IDs for the query tokens + - num_ctx_per_req: per-request context length in the pool + - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed + views of the persistent per-layer K/V pool. + - ctx_rows_start: where ctx_cache_batch_idx is the contiguous rows + [ctx_rows_start, ctx_rows_start + num_gens) of ctx_page_table, their + start (``_ctx_rows_start``), else None. + """ + num_contexts = attn_metadata.num_contexts + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + + mask_token_id = self._mask_token_id(draft_model) + + # Get the embed_tokens layer from the draft model + embed_tokens = draft_model.draft_model_full.model.embed_tokens + hidden_dim = ( + spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] + ) + + if num_gens > 0: + gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] + gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] + + total_tokens_per_req = self._draft_tokens_per_req # K+1 + K = spec_metadata.runtime_draft_len + + # Get captured multi-layer hidden states from spec_metadata + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + has_target_features = ( + captured_hs is not None + and hasattr(draft_model, "fc") + and hasattr(draft_model, "hidden_norm") + ) + + # Use cached block_size (resolved once on first call) + block_size = self._compute_block_size + query_tokens_per_req = block_size + + # Get slots for gen requests from pre-computed mapping + slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] + K_plus_1 = K + 1 + + if prep is not None: + gen_rows_out = self._k3_static_arange(num_contexts, num_contexts + num_gens) + offsets_kp1 = self._k3_static_arange(0, K_plus_1) + bonus = prep["bonus"] + query_position_ids = prep["qpos"] + ctx_position_ids = prep["cpos"] + noise_embed_2d = prep["noise"] + else: + gen_rows_out = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + + bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) + bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() + + ctx_len_gen = self._ctx_len[slots] + j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") + offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") + + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, + bonus, + mask_token_id, + query_tokens_per_req, + self._trained_mask_embedding(draft_model, mask_token_id), + ) + + # Accumulate new accepted features into context buffers + fused_ctx_kv = False + if has_target_features: + gen_start = attn_metadata.num_ctx_tokens + # Target now processes exactly K+1 tokens per gen req, so the + # captured slice is already the full set we need to project. + gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] + gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) + projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) + fused_ctx_kv = prep is not None and self._k3_ctx_kv_applies( + draft_model, projected_to_store + ) + if fused_ctx_kv: + num_ctx_fused = self._k3_ctx_kv( + draft_model, + projected_to_store, + ctx_position_ids, + gen_num_accepted, + slots, + gen_rows_out, + ) + # k3_ctx_kv adds the accepted tokens in the kernel; empty the dummy slot after + # it, as _advance_ctx_len does after the add on the torch path. + self._ctx_len[self._dummy_slot].zero_() + if has_target_features and not fused_ctx_kv: + gen_num_accepted_long = gen_num_accepted.long() + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) + if self._ctx_block_tables is not None: + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) + else: + col_idx = col_idx.clamp(max=self._max_ctx - 1) + + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) + col_flat = col_idx.reshape(-1) + proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] + pos_flat = ctx_position_ids.long().reshape(-1) + mask_1d = write_mask.reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + k_new, v_new = draft_model.precompute_context_kv( + proj_flat.to(cache_dtype), pos_flat + ) + mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + slot_long = slot_flat.long() + col_long = col_flat.long() + if self._ctx_paged: + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + self._advance_ctx_len(slots, gen_num_accepted_long) + + if fused_ctx_kv: + num_ctx_per_req_t = num_ctx_fused + else: + num_ctx_per_req_t = self._ctx_len[slots] + if not fused_ctx_kv and self._ctx_block_tables is not None: + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + noise_embedding = noise_embed_2d + query_positions = query_position_ids.long() + + # Update seq_lens for gen requests to K+1 (the number of tokens + # the target forward actually processed). + attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( + total_tokens_per_req + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req + else: + noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) + query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") + num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") + slots = torch.empty(0, dtype=torch.long, device="cuda") + gen_rows_out = slots + bonus = torch.empty(0, dtype=torch.long, device="cuda") + + return { + "noise_embedding": noise_embedding, + "query_positions": query_positions, + "num_ctx_per_req": num_ctx_per_req_t, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": bonus, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), + "ctx_rows_start": self._ctx_rows_start(num_contexts, num_gens), + } + + +class KimiK3DSparkWorker(KimiK3DFlashWorker, DSparkWorker): + # Whether this step's block logits are the Kimi K3 decode path's vocabulary shard (see ``_draft_block_logits``). + _k3_sharded_block_logits = False + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Block logits the Kimi K3 decode path kept vocab-sharded (see ``_draft_block_logits``) run the + whole chain as ``trtllm::k3_markov``, which also returns the greedy draft tokens and + next_new_tokens. Other TP-sharded logits slice markov_w2's rows to this rank's contiguous + shard and chain through the TP-aware global argmax. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # Unlike the d2t guard this one cannot move to set_draft_model: it + # keys on the runtime logits width, and reproducing that at init would + # duplicate the draft head's sharding rules. A standalone drafter + # borrows the target lm_head, whose gather_output defaults to True, so + # the logits normally arrive full-vocab and this branch is skipped. + self._k3_markov = None + k3_sharded, self._k3_sharded_block_logits = self._k3_sharded_block_logits, False + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + if k3_sharded: + return self._k3_markov_chain( + draft_model, gen_logits, first_prev_tokens, vocab_slice + ) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) + + def _keep_draft_logits_sharded(self, draft_model, spec_metadata, num_gens: int) -> bool: + """Keep the draft logits vocab-sharded: ``trtllm::k3_markov`` reduces across the ranks itself. + + The chain's global argmaxes are the greedy draft tokens, so the lm_head's all-gather of the + full logits and the draft sampler's gather are both skipped. Needs a Kimi K3 target running its + decode kernels (``k3_decode``), plain TP, greedy drafting (rejection sampling reads + full-vocabulary probabilities), an unquantized bias-free column-parallel bf16 lm_head whose + shards tile the Markov vocabulary, bf16 Markov weights of the kernel's rank, MNNVL (the model's + all-reduces own an MNNVL workspace for this mapping), and a block, shard and batch of + ``num_gens`` requests the kernel can split. + """ + if not self.k3_decode: + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + lm_head = getattr(draft_model, "lm_head", None) + if ( + lm_head is None + or not getattr(draft_model, "has_markov_head", False) + or mapping is None + or mapping.tp_size <= 1 + or mapping.enable_attention_dp + or spec_metadata.wants_advanced_draft_sampling + or getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "bias", None) is not None + or lm_head.weight.dtype != torch.bfloat16 + or draft_model.markov_w1.dtype != torch.bfloat16 + or draft_model.markov_w2.dtype != torch.bfloat16 + or lm_head.weight.shape[0] * mapping.tp_size != draft_model.markov_w2.shape[0] + or MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None + ): + return False + shard = lm_head.weight.shape[0] + block = spec_metadata.runtime_draft_len + rank = k3_markov_op._kernel_module().MARKOV_RANK + if ( + not 0 < block <= k3_markov_op.WORKSPACE_MAX_BLOCK + or draft_model.markov_w1.dim() != 2 + or draft_model.markov_w1.shape[1] != rank + or tuple(draft_model.markov_w2.shape[1:]) != (rank,) + ): + return False + if k3_markov_op.pick_grid(shard, block, num_gens) == 0: + logger.warning_once( + f"DSpark Markov head: trtllm::k3_markov cannot split a {shard}-row vocab shard for " + f"{num_gens} requests; the draft logits are all-gathered and the chain runs unfused.", + key=f"dspark_k3_markov_shard_{num_gens}", + ) + return False + return True + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + ) -> torch.Tensor: + """This rank's bf16 shard of the block logits when ``k3_markov`` takes them (it converts them to fp32 + exactly, as ``.float()`` would), from the drafter's logits processor's head kernel where it takes the rows, + else the head's own GEMM (``_k3_head_shard``); otherwise the base class' fp32 logits.""" + num_gens = gen_hidden_states.shape[0] // max(spec_metadata.runtime_draft_len, 1) + self._k3_sharded_block_logits = self._keep_draft_logits_sharded( + draft_model, spec_metadata, num_gens + ) + if self._k3_sharded_block_logits: + return self._k3_head_shard( + draft_model.logits_processor, draft_model.lm_head, gen_hidden_states + ) + return super()._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): + """The base acceptance; under ``k3_decode`` its outputs are kept for ``k3_markov``'s next_new_tokens.""" + accepted_tokens, num_accepted_tokens = super().sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + if self.k3_decode: + self._on_acceptance(accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata) + return accepted_tokens, num_accepted_tokens + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Keeps this step's acceptance for ``k3_markov``'s next_new_tokens (and the metadata whose KV lengths it + rewinds).""" + self._k3_acceptance = ( + accepted_tokens, num_accepted_tokens, attn_metadata.num_contexts, spec_metadata, attn_metadata + ) # fmt: skip + + def _k3_markov_chain( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + vocab_slice: slice, + ) -> torch.Tensor: + """The Markov chain as ``trtllm::k3_markov`` on this rank's shard of the block logits. + + Returns the corrected logits (fp32); keeps the kernel's greedy tokens for + ``sample_draft_tokens`` and its next_new_tokens (built from this step's acceptance) for + ``_prepare_next_new_tokens``. + """ + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + + acceptance = getattr(self, "_k3_acceptance", None) + if acceptance is None: + raise RuntimeError("DSpark Markov head: k3_markov needs this step's acceptance first") + accepted_tokens, num_accepted_tokens, num_contexts, spec_metadata, attn_metadata = ( + acceptance + ) + num_gens = gen_logits.shape[0] + # The draft forward's pending KV-length rewind (see _apply_kv_rewind_after_draft) runs in the kernel: every + # reader of kv_lens_cuda in the draft forward precedes it. In warmup the restore of kv_lens_cuda that + # follows overwrites it, as it would the Python rewind. + rewind = getattr(self, "_kv_rewind_amount", None) + fold = ( + getattr(self, "_kv_rewind_pending", False) + and rewind is not None + and getattr(attn_metadata, "kv_lens_cuda", None) is not None + and rewind.dtype == torch.int32 + and self._kv_rewind_bs - self._kv_rewind_nc == rewind.numel() <= num_gens + ) + corrected, tokens, next_new = k3_markov_op.markov_chain( + self.mapping, + gen_logits.contiguous(), + first_prev_tokens.long(), + draft_model.markov_w1, + draft_model.markov_w2[vocab_slice], + vocab_slice.start, + accepted_tokens, + num_accepted_tokens[num_contexts : num_contexts + num_gens], + spec_metadata.batch_indices_cuda[num_contexts : num_contexts + num_gens], + kv_lens=attn_metadata.kv_lens_cuda if fold else None, + rewind=rewind if fold else None, + rewind_first=self._kv_rewind_nc if fold else 0, + ) + if fold: + self._kv_rewind_amount = None + self._kv_rewind_pending = False + self._k3_markov = (corrected, tokens, next_new, accepted_tokens, num_accepted_tokens) + return corrected + + def sample_draft_tokens( + self, + logits, + spec_metadata, + batch_size, + *, + num_contexts=0, + draft_step=None, + mapping_lm_head_tp=None, + ): + """Greedy block drafts straight from ``k3_markov``. + + When ``logits`` are the corrected logits the kernel just returned, its per-position global + argmax (first maximum, lowest vocabulary index among equal values) is the token the + TP-gathered greedy sampler would pick. Anything else goes to the base sampler. + """ + chain = getattr(self, "_k3_markov", None) + if ( + chain is not None + and chain[0] is logits + and mapping_lm_head_tp is None + and not spec_metadata.wants_advanced_draft_sampling + ): + self._k3_markov_next = (chain[1], chain[2], chain[3], chain[4]) + return chain[1] + self._k3_markov_next = None + return super().sample_draft_tokens( + logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + draft_step=draft_step, + mapping_lm_head_tp=mapping_lm_head_tp, + ) + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """``k3_markov``'s next_new_tokens when the drafts are its tokens for the whole batch (no context + requests); otherwise the base assembly.""" + chain = getattr(self, "_k3_markov_next", None) + # The step ends here: release its acceptance (which holds the step's attention and spec metadata) and the + # kernel's outputs, so they do not keep a captured graph's pool or metadata alive after the step. + self._k3_markov_next = None + self._k3_acceptance = None + self._k3_markov = None + if ( + chain is not None + and chain[0] is next_draft_tokens + and chain[2] is accepted_tokens + and chain[3] is num_accepted_tokens + and next_draft_tokens.shape[0] == batch_size + ): + return chain[1] + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py index 24233d194709..b4177f698d8b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/weights.py @@ -1,6 +1,42 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The shared Kimi K3 text checkpoint loader.""" +"""Weight loading: Kimi K3 (MXFP4) / sm_100 / tp16_moetp4ep4. -from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD -from ..weights import load as load +The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and +`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: + +* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the + routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under + moe_tp 4 x moe_ep 4), and checks its own key coverage. +* The vision tower and the projector are a predicted non-load: listed here, never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_utils import filter_weights + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") + + +def load(model, weights) -> None: + """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) + assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = _LANG_PREFIX + KimiLinearForCausalLM.load_weights(model, lm_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/__init__.py new file mode 100644 index 000000000000..d4c5a174ed15 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 Kimi K3 targets with FP8 fused attention on SM103 and SM107.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py new file mode 100644 index 000000000000..627d026db775 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py @@ -0,0 +1,433 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FP8-preserving fused decode kernels for the NVFP4 Kimi K3 targets.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _kda_decode( + QKVG, + BFA, + FB, + WQ, + WK, + WV, + A_LOG, + DT, + NORM, + CONV, + STATE, + SLOTS, + OUT, + H: tl.constexpr, + D: tl.constexpr, + PROJ_STRIDE: tl.constexpr, + BFA_STRIDE: tl.constexpr, + CONV_SLOT: tl.constexpr, + CONV_CHANNEL: tl.constexpr, + CONV_TIME: tl.constexpr, + STATE_SLOT: tl.constexpr, + LOWER: tl.constexpr, + SCALE: tl.constexpr, + EPS: tl.constexpr, +): + row, head = tl.program_id(0), tl.program_id(1) + slot = tl.load(SLOTS + row).to(tl.int64) + k = tl.arange(0, 128) + channel = head * 128 + k + fa = tl.load(BFA + row * BFA_STRIDE + k) + fb = tl.load(FB + channel[None, :] * 128 + k[:, None]) + gate_tile = tl.dot(tl.broadcast_to(fa[None, :], (16, 128)), fb) + gate = tl.sum(tl.where(tl.arange(0, 16)[:, None] == 0, gate_tile, 0.0), 0) + gate = gate.to(tl.bfloat16).to(tl.float32) + beta = tl.sigmoid(tl.load(BFA + row * BFA_STRIDE + 128 + head).to(tl.float32)) + q = tl.full((128,), 0, tl.float32) + key = tl.full((128,), 0, tl.float32) + value = tl.full((128,), 0, tl.float32) + conv_base = slot * CONV_SLOT + channel * CONV_CHANNEL + for tap in tl.static_range(3): + cq = tl.load(CONV + conv_base + tap * CONV_TIME).to(tl.float32) + ck = tl.load(CONV + conv_base + D * CONV_CHANNEL + tap * CONV_TIME).to(tl.float32) + cv = tl.load(CONV + conv_base + 2 * D * CONV_CHANNEL + tap * CONV_TIME).to(tl.float32) + q += cq * tl.load(WQ + channel * 4 + tap) + key += ck * tl.load(WK + channel * 4 + tap) + value += cv * tl.load(WV + channel * 4 + tap) + xq = tl.load(QKVG + row * PROJ_STRIDE + channel) + xk = tl.load(QKVG + row * PROJ_STRIDE + D + channel) + xv = tl.load(QKVG + row * PROJ_STRIDE + 2 * D + channel) + q += xq.to(tl.float32) * tl.load(WQ + channel * 4 + 3) + key += xk.to(tl.float32) * tl.load(WK + channel * 4 + 3) + value += xv.to(tl.float32) * tl.load(WV + channel * 4 + 3) + q *= tl.sigmoid(q) + key *= tl.sigmoid(key) + value *= tl.sigmoid(value) + q *= tl.rsqrt(tl.sum(q * q, 0) + 1.0e-6) * SCALE + key *= tl.rsqrt(tl.sum(key * key, 0) + 1.0e-6) + decay = tl.exp( + LOWER * tl.sigmoid(tl.exp(tl.load(A_LOG + head)) * (gate + tl.load(DT + channel))) + ) + offsets = slot * STATE_SLOT + head * 16384 + k[:, None] * 128 + k[None, :] + state = tl.load(STATE + offsets).to(tl.float32) * decay[None, :] + innovation = (value - tl.sum(state * key[None, :], 1)) * beta + state += innovation[:, None] * key[None, :] + tl.store(STATE + offsets, state) + output = tl.sum(state * q[None, :], 1) + output *= tl.rsqrt(tl.sum(output * output, 0) / 128 + EPS) + output *= tl.load(NORM + k) + output *= tl.sigmoid(tl.load(QKVG + row * PROJ_STRIDE + 3 * D + channel).to(tl.float32)) + tl.store(OUT + row * D + channel, output) + for section in tl.static_range(3): + base = conv_base + section * D * CONV_CHANNEL + old1 = tl.load(CONV + base + CONV_TIME) + old2 = tl.load(CONV + base + 2 * CONV_TIME) + new = tl.load(QKVG + row * PROJ_STRIDE + section * D + channel) + tl.store(CONV + base, old1) + tl.store(CONV + base + CONV_TIME, old2) + tl.store(CONV + base + 2 * CONV_TIME, new) + + +@torch.library.custom_op("trtllm::k3_fp8_kda_decode", mutates_args=("conv", "state")) +def kda_decode( + qkvg: torch.Tensor, + bfa: torch.Tensor, + fb: torch.Tensor, + wq: torch.Tensor, + wk: torch.Tensor, + wv: torch.Tensor, + a_log: torch.Tensor, + dt: torch.Tensor, + norm: torch.Tensor, + conv: torch.Tensor, + state: torch.Tensor, + slots: torch.Tensor, + lower: float, + scale: float, + eps: float, +) -> torch.Tensor: + """Fuse f_b, convolution, normalized delta recurrence and gated RMSNorm.""" + h = state.shape[1] + if state.shape[2:] != (128, 128) or conv.shape[1:] != (3 * h * 128, 3): + raise ValueError("Fused FP8 KDA requires 128-wide heads and width-four convolution") + out = qkvg.new_empty((qkvg.shape[0], h, 128)) + _kda_decode[(qkvg.shape[0], h)]( + qkvg, + bfa, + fb, + wq, + wk, + wv, + a_log, + dt, + norm, + conv, + state, + slots, + out, + h, + h * 128, + qkvg.stride(0), + bfa.stride(0), + *conv.stride(), + state.stride(0), + lower, + scale, + eps, + num_warps=8, + enable_fp_fusion=False, + ) + return out + + +@kda_decode.register_fake +def _kda_fake(qkvg, bfa, fb, wq, wk, wv, a_log, dt, norm, conv, state, slots, lower, scale, eps): + return qkvg.new_empty((qkvg.shape[0], state.shape[1], 128)) + + +@triton.jit +def _mla_qkv( + AG, + QA_NORM, + QB, + QB_SCALE, + KB, + KB_SCALE, + KV_NORM, + POOL, + PAGES, + LENGTHS, + OUT, + M: tl.constexpr, + H: tl.constexpr, + AG_STRIDE: tl.constexpr, + QB_SCALE_STRIDE: tl.constexpr, + KB_S0: tl.constexpr, + KB_S1: tl.constexpr, + KB_S2: tl.constexpr, + KS_S0: tl.constexpr, + KS_S1: tl.constexpr, + KS_S2: tl.constexpr, + ROW_STRIDE: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + PAGE_OFFSET: tl.constexpr, + T: tl.constexpr, + QA_EPS: tl.constexpr, + KV_EPS: tl.constexpr, +): + head, chunk = tl.program_id(0), tl.program_id(1) + rows = chunk * 16 + tl.arange(0, 16) + qa_cols = tl.arange(0, 2048) + qa = tl.load( + AG + rows[:, None] * AG_STRIDE + qa_cols[None, :], + (rows[:, None] < M) & (qa_cols[None, :] < 1536), + 0, + ).to(tl.float32) + rms = tl.rsqrt(tl.sum(qa * qa, 1) / 1536 + QA_EPS) + n = tl.arange(0, 256) + k = tl.arange(0, 128) + qb_rows = head * 192 + n + q = tl.full((16, 256), 0, tl.float32) + for block in range(12): + cols = block * 128 + k + a = tl.load(AG + rows[:, None] * AG_STRIDE + cols[None, :], rows[:, None] < M, 0).to( + tl.float32 + ) + a *= rms[:, None] * tl.load(QA_NORM + cols)[None, :].to(tl.float32) + a = a.to(tl.bfloat16).to(tl.float32) + a_scale = tl.maximum(tl.max(tl.abs(a), 1) / 448.0, 1.0e-12) + a8 = (a / a_scale[:, None]).to(tl.float8e4nv) + b = tl.load(QB + qb_rows[None, :] * 1536 + cols[:, None], n[None, :] < 192, 0.0) + bs = tl.load(QB_SCALE + (qb_rows // 128) * QB_SCALE_STRIDE + block, n < 192, 0) + q += tl.dot(a8, b) * a_scale[:, None] * bs[None, :] + q = q.to(tl.bfloat16).to(tl.float32) + # Only the NoPE portion is absorbed; the PE portion is copied unchanged. + qn = tl.gather(q, tl.broadcast_to(k[None, :], (16, 128)), 1) + q_scale = tl.maximum(tl.max(tl.abs(qn), 1) / 448.0, 1.0e-12) + q8 = (qn / q_scale[:, None]).to(tl.float8e4nv) + latent = tl.arange(0, 512) + kb = tl.load(KB + head * KB_S0 + latent[None, :] * KB_S1 + k[:, None] * KB_S2) + ks = tl.load(KB_SCALE + head * KS_S0 + (latent // 128) * KS_S1) + absorbed = tl.dot(q8, kb) * q_scale[:, None] * ks[None, :] + base = rows[:, None] * (H * 576) + head * 576 + tl.store(OUT + base + latent[None, :], absorbed, rows[:, None] < M) + pe = tl.arange(0, 64) + qpe = tl.gather(q, tl.broadcast_to((128 + pe)[None, :], (16, 64)), 1) + tl.store(OUT + base + 512 + pe[None, :], qpe, rows[:, None] < M) + if head == 0: + seq = rows // T + length = tl.load(LENGTHS + seq, rows < M, 0) + pos = length - T + rows % T + page = tl.load(PAGES + seq * PAGE_STRIDE + pos // 64, (rows < M) & (pos >= 0), 0) + cache_row = ((page.to(tl.int64) + PAGE_OFFSET) * 64 + pos % 64) * ROW_STRIDE + ckv = tl.load( + AG + rows[:, None] * AG_STRIDE + 1536 + latent[None, :], rows[:, None] < M, 0 + ).to(tl.float32) + kv_rms = tl.rsqrt(tl.sum(ckv * ckv, 1) / 512 + KV_EPS) + ckv *= kv_rms[:, None] * tl.load(KV_NORM + latent)[None, :].to(tl.float32) + tl.store( + POOL + cache_row[:, None] + latent[None, :], + ckv, + (rows[:, None] < M) & (pos[:, None] >= 0), + ) + kpe = tl.load(AG + rows[:, None] * AG_STRIDE + 2048 + pe[None, :], rows[:, None] < M, 0) + tl.store( + POOL + cache_row[:, None] + 512 + pe[None, :], + kpe, + (rows[:, None] < M) & (pos[:, None] >= 0), + ) + + +@torch.library.custom_op("trtllm::k3_fp8_mla_qkv", mutates_args=("pool",)) +def mla_qkv( + ag: torch.Tensor, + qa_norm: torch.Tensor, + qb: torch.Tensor, + qb_scale: torch.Tensor, + kb: torch.Tensor, + kb_scale: torch.Tensor, + kv_norm: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + lengths: torch.Tensor, + row_stride: int, + page_offset: int, + tokens_per_request: int, + qa_eps: float, + kv_eps: float, +) -> torch.Tensor: + """Fuse Q-A norm, FP8 Q-B/absorption and the paged latent-cache update.""" + h = kb.shape[0] + out = ag.new_empty((ag.shape[0], h * 576)) + _mla_qkv[(h, triton.cdiv(ag.shape[0], 16))]( + ag, + qa_norm, + qb, + qb_scale, + kb, + kb_scale, + kv_norm, + pool, + pages, + lengths, + out, + ag.shape[0], + h, + ag.stride(0), + qb_scale.stride(0), + *kb.stride(), + *kb_scale.stride(), + row_stride, + pages.stride(0), + page_offset, + tokens_per_request, + qa_eps, + kv_eps, + num_warps=8, + ) + return out + + +@mla_qkv.register_fake +def _qkv_fake( + ag, + qa_norm, + qb, + qb_scale, + kb, + kb_scale, + kv_norm, + pool, + pages, + lengths, + row_stride, + page_offset, + tokens_per_request, + qa_eps, + kv_eps, +): + return ag.new_empty((ag.shape[0], kb.shape[0] * 576)) + + +@triton.jit +def _mla_attn( + Q, + POOL, + PAGES, + LENGTHS, + VB, + VS, + GATE, + OUT, + H: tl.constexpr, + ROW_STRIDE: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + PAGE_OFFSET: tl.constexpr, + T: tl.constexpr, + SCALE: tl.constexpr, + GATE_STRIDE: tl.constexpr, + GATE_COL: tl.constexpr, +): + row, head = tl.program_id(0), tl.program_id(1) + seq = row // T + length = tl.load(LENGTHS + seq) + end = length - T + row % T + 1 + d = tl.arange(0, 1024) + v = tl.arange(0, 512) + q = tl.load(Q + row * H * 576 + head * 576 + d, d < 576, 0).to(tl.float32) + maximum = tl.full((), -float("inf"), tl.float32) + denominator = tl.full((), 0.0, tl.float32) + output = tl.full((512,), 0.0, tl.float32) + for start in range(0, tl.cdiv(tl.maximum(end, 0), 32)): + pos = start * 32 + tl.arange(0, 32) + page = tl.load(PAGES + seq * PAGE_STRIDE + pos // 64, pos < end, 0) + offset = ((page.to(tl.int64) + PAGE_OFFSET) * 64 + pos % 64) * ROW_STRIDE + kv = tl.load( + POOL + offset[:, None] + d[None, :], (pos[:, None] < end) & (d[None, :] < 576), 0 + ).to(tl.float32) + scores = tl.sum(kv * q[None, :], 1) * SCALE + scores = tl.where(pos < end, scores, -float("inf")) + new_max = tl.maximum(maximum, tl.max(scores, 0)) + alpha = tl.exp(maximum - new_max) + probability = tl.exp(scores - new_max) + value = tl.gather(kv, tl.broadcast_to(v[None, :], (32, 512)), 1) + output = output * alpha + tl.sum(probability[:, None] * value, 0) + denominator = denominator * alpha + tl.sum(probability, 0) + maximum = new_max + output = (output / tl.maximum(denominator, 1.0e-20)).to(tl.bfloat16).to(tl.float32) + # Quantize each 128-wide latent group before the FP8 V-B projection. + groups = tl.reshape(output, (4, 128)) + scales = tl.maximum(tl.max(tl.abs(groups), 1) / 448.0, 1.0e-12) + quantized = (groups / scales[:, None]).to(tl.float8e4nv).to(tl.float32) + n = tl.arange(0, 128) + projected_tile = tl.full((16, 128), 0.0, tl.float32) + for group in tl.static_range(4): + latent = tl.sum(tl.where(tl.arange(0, 4)[:, None] == group, quantized, 0.0), 0) + activation_scale = tl.sum(tl.where(tl.arange(0, 4) == group, scales, 0.0), 0) + activation = tl.broadcast_to(latent[None, :], (16, 128)).to(tl.float8e4nv) + weights = tl.load(VB + head * 128 * 512 + n[None, :] * 512 + group * 128 + n[:, None]) + weight_scale = tl.load(VS + head * 4 + group) + projected_tile += tl.dot(activation, weights) * activation_scale * weight_scale + projected = tl.sum(tl.where(tl.arange(0, 16)[:, None] == 0, projected_tile, 0.0), 0) + projected = projected.to(tl.bfloat16).to(tl.float32) + gate = tl.sigmoid(tl.load(GATE + row * GATE_STRIDE + GATE_COL + head * 128 + n).to(tl.float32)) + tl.store(OUT + row * H * 128 + head * 128 + n, projected * gate) + + +@torch.library.custom_op("trtllm::k3_fp8_mla_attn", mutates_args=()) +def mla_attn( + q: torch.Tensor, + pool: torch.Tensor, + pages: torch.Tensor, + lengths: torch.Tensor, + vb: torch.Tensor, + vb_scale: torch.Tensor, + gate: torch.Tensor, + row_stride: int, + page_offset: int, + tokens_per_request: int, + softmax_scale: float, + gate_col: int, +) -> torch.Tensor: + """Fuse paged causal attention, FP8 V-B projection and output gating.""" + h = vb.shape[0] + out = q.new_empty((q.shape[0], h * 128)) + _mla_attn[(q.shape[0], h)]( + q, + pool, + pages, + lengths, + vb, + vb_scale, + gate, + out, + h, + row_stride, + pages.stride(0), + page_offset, + tokens_per_request, + softmax_scale, + gate.stride(0), + gate_col, + num_warps=8, + ) + return out + + +@mla_attn.register_fake +def _attn_fake( + q, + pool, + pages, + lengths, + vb, + vb_scale, + gate, + row_stride, + page_offset, + tokens_per_request, + softmax_scale, + gate_col, +): + return q.new_empty((q.shape[0], vb.shape[0] * 128)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_comm.py new file mode 100644 index 000000000000..f5327daf9474 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_comm.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's residual-update collectives on the catalog's Kimi K3 MNNVL and sandwich entries. + +A layer's post-attention step adds the attention output to the running prefix sum, selects the attention residual +and applies the post-attention RMSNorm. Under TP the attention output is a sum over the TP group, so the step is that +all-reduce followed by the residual update. On a step of at most `AR_ATTN_RES_MAX_TOKENS` tokens (wide decode steps +aside) one collective runs both: + +* `K3DecodeComm.sandwich_oproj`, `comm/k3_sandwich_oproj`: o_proj, its all-reduce and the residual update in one + kernel, at most `SANDWICH_MAX_TOKENS` tokens of an o_proj of the TP16 per-rank shape [7168, 768]. The attention + hands over its gated o_proj input. +* `K3DecodeComm.allreduce_attn_res`, `comm/mnnvl_allreduce_attn_res`, everywhere else: the one-shot MNNVL all-reduce + of the attention's unreduced o_proj partial, with the residual update as its epilogue. + +The pre-attention step after a MoE layer whose row-parallel tail was handed on (a `PendingTail`, `decode_moe.py`) is +the same kind of collective: `K3DecodeComm.sandwich_tail`, `comm/k3_sandwich_tail`, runs the tail GEMV, its +all-reduce and the next layer's residual update (the final norm's, after the last layer) in one kernel. + +The DSpark drafter (`K3DSparkDrafter`) runs its collectives with a plain residual add and RMSNorm on the same state: + +* `K3DecodeComm.sandwich_plain`, `comm/k3_sandwich_plain`: a row-parallel projection of at most `SANDWICH_MAX_TOKENS` + tokens (a layer's attention output projection, or its MLP's SiLU-and-mul and down projection), its all-reduce, the + residual add and the RMSNorm in one kernel. +* `K3DecodeComm.allreduce_norm`, `comm/mnnvl_fusion_allreduce`: the all-reduce of an unreduced projection output of + any token count the MNNVL workspace holds, with the residual add and the RMSNorm in its epilogue (the context + projection's with a zero residual). + +The state is one `MnnvlWorkspace` and one `K3SandwichWorkspace` of the TP group (`K3DecodeComm.create`): collective +over the group and eager, built by the target in `post_load_weights` before any CUDA-graph capture. Every rank must +make the same calls on each in the same order. Which call a step takes is decided from its token count and kind and +from the load-time layout alone, which every rank of the group shares. + +The plain MNNVL all-reduces the decode path keeps (the stock modules') send one-shot up to +`DECODE_AR_ONE_SHOT_MAX_BYTES` (`use_decode_one_shot`), except a wide decode step's (`wide_all_reduce`). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import NamedTuple, Optional, Tuple + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_oproj import ( + K3SandwichWorkspace, + k3_sandwich_oproj, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_plain import ( + k3_sandwich_plain, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_sandwich_tail import ( + k3_sandwich_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allreduce_attn_res import ( + mnnvl_allreduce_attn_res, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_fusion_allreduce import ( + mnnvl_fusion_allreduce, + required_buffer_bytes, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) + +# The kernel's support predicate: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_sandwich import op as _sandwich_op +from tensorrt_llm._torch.distributed import AllReduceParams + +# The most tokens of a step whose post-attention all-reduce runs one-shot with the residual update in its epilogue. +AR_ATTN_RES_MAX_TOKENS = 16 + +# The most tokens the sandwich kernel takes (one token tile). +SANDWICH_MAX_TOKENS = _sandwich_op.MAX_TOKENS + +# The MNNVL workspace's buffer size: a one-shot call pushes T x 7168 bf16 from each of 16 ranks, 3.5 MiB at T = 16. +MNNVL_BUFFER_BYTES = 4 << 20 + +# The one-shot ceiling of the stock MNNVL all-reduces on the decode path: 8 tokens x 7168 x 16 ranks x 2 B is +# 1.75 MiB, which the stock 1 MiB ceiling would send two-shot. +DECODE_AR_ONE_SHOT_MAX_BYTES = 4 << 20 + +# A wide decode step's ceiling: the stock 1 MiB. At 16 ranks two-shot is faster for every wide step's rows. +WIDE_AR_ONE_SHOT_MAX_BYTES = 1 << 20 + + +def use_decode_one_shot(model: nn.Module) -> None: + """Every stock MNNVL all-reduce of ``model`` sends one-shot up to `DECODE_AR_ONE_SHOT_MAX_BYTES`. Each grows its + workspace on its first eager call of a larger size, before the capture of that size.""" + for module in model.modules(): + mnnvl = getattr(module, "mnnvl_allreduce", None) + if mnnvl is not None: + mnnvl.one_shot_max_bytes = DECODE_AR_ONE_SHOT_MAX_BYTES + mnnvl.early_trigger = True + + +def skip_all_reduce() -> AllReduceParams: + """All-reduce parameters under which a row-parallel module returns this rank's unreduced output (its projection + without its all-reduce): a ``Linear``'s ``all_reduce_params``, a ``GatedMLP``'s ``final_all_reduce_params``.""" + return AllReduceParams(enable_allreduce=False) + + +def wide_all_reduce(all_reduce: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``all_reduce(x)`` of a wide decode step (a stock ``AllReduce`` module, no fusion): its MNNVL all-reduce with + the `WIDE_AR_ONE_SHOT_MAX_BYTES` ceiling, else the module itself.""" + mnnvl = getattr(all_reduce, "mnnvl_allreduce", None) + if mnnvl is not None: + out = mnnvl( + x.contiguous(), AllReduceParams(), one_shot_max_bytes=WIDE_AR_ONE_SHOT_MAX_BYTES + ) + if out is not None: + return out + return all_reduce(x) + + +class PendingTail(NamedTuple): + """A MoE layer's row-parallel tail left to its consumer's fused pre-attention step (``K3DecodeComm.sandwich_tail``): + the reduced latent ``[T, 3584]``, the shared experts' activation, the tail weight ``[latent up columns | padding | + shared down]``, this rank's first latent column and the latent norm's epsilon.""" + + latent: torch.Tensor + act: torch.Tensor + weight: torch.Tensor + lo: int + lat_eps: float + + +def _eps(norm: nn.Module) -> float: + """The epsilon of a KimiK3RMSNorm or ``torch.nn.RMSNorm`` (``eps``; for torch's None, its default: the machine + epsilon of the weight's dtype) or a stock RMSNorm (``variance_epsilon``).""" + eps = norm.eps if hasattr(norm, "eps") else norm.variance_epsilon + return float(torch.finfo(norm.weight.dtype).eps if eps is None else eps) + + +def _res_args(res_proj: nn.Module, res_norm: nn.Module, out_norm: nn.Module) -> tuple: + """The residual update's weights and epsilons in the entries' order.""" + return ( + res_proj.weight.reshape(-1), + res_norm.weight, + out_norm.weight, + _eps(res_norm), + _eps(out_norm), + ) + + +@dataclass(eq=False) +class K3DecodeComm: + """The decode path's collective state for the TP group: the `MnnvlWorkspace` and the `K3SandwichWorkspace` the + post-attention steps run on. Built by `create`; owned by the target and shared by every layer.""" + + mnnvl: MnnvlWorkspace + sandwich: K3SandwichWorkspace + + @classmethod + def create(cls, mapping, oproj_weight: Optional[torch.Tensor] = None) -> "K3DecodeComm": + """The state for ``mapping``'s TP group. Collective: every rank of the group calls it at the same point, + eagerly, before any CUDA-graph capture; each workspace fails on every rank or on none. + + ``oproj_weight``: an o_proj weight ``takes_oproj`` holds for. The sandwich kernel compiles here for its + shape, with one call on a zero row of a zero weight, so no capture compiles it; the call advances the + sandwich workspace on every rank alike.""" + state = cls( + MnnvlWorkspace.create(mapping, MNNVL_BUFFER_BYTES), + K3SandwichWorkspace.create(mapping), + ) + if oproj_weight is not None: + weight = torch.zeros_like(oproj_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_oproj( + weight.new_zeros(1, weight.shape[1]), + weight, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + state.sandwich, + ) + torch.cuda.synchronize(weight.device) + return state + + def compile_tail(self, latent_size: int, act_size: int, tail_weight: torch.Tensor) -> None: + """Compile the sandwich tail kernel for a MoE tail of ``latent_size`` latent and ``act_size`` shared columns + and ``tail_weight``'s shape, with one call on a zero row of a zero weight, before any capture. Collective: every + rank of the group makes the call; it advances the sandwich workspace on every rank alike.""" + weight = torch.zeros_like(tail_weight) + hidden = weight.shape[0] + ones = weight.new_ones(hidden) + k3_sandwich_tail( + weight.new_zeros(1, latent_size), + weight.new_zeros(1, act_size), + weight, + 0, + 1e-6, + None, + weight.new_zeros(0, 1, hidden), + weight.new_zeros(hidden), + ones, + ones, + 1e-6, + 1e-6, + self.sandwich, + ) + torch.cuda.synchronize(weight.device) + + def takes_post_attention(self, hidden_states: torch.Tensor, step) -> bool: + """Whether one of this state's collectives runs the post-attention step of a layer whose attention input is + ``hidden_states``: at most `AR_ATTN_RES_MAX_TOKENS` bf16 rows of a hidden size the MNNVL entry takes, on any + step but a wide decode step (``step.wide``), whose post-attention update stays the all-reduce and the fused + add + attn_res + RMSNorm.""" + if step is not None and step.wide: + return False + if hidden_states.dim() != 2 or hidden_states.dtype != torch.bfloat16: + return False + rows, hidden = hidden_states.shape + return ( + hidden % 1024 == 0 + and hidden <= 8192 + and 0 < rows <= min(AR_ATTN_RES_MAX_TOKENS, self.mnnvl.max_one_shot_tokens(hidden)) + ) + + @staticmethod + def takes_oproj(o_proj: nn.Module, max_snapshots: int) -> bool: + """Whether ``sandwich_oproj`` takes the layer of output projection ``o_proj`` on a step of at most + `SANDWICH_MAX_TOKENS` tokens, decided once the weights are final: a bias-free, contiguous bf16 weight of the + kernel's shape (the TP16 per-rank [7168, 768]) and a snapshot bank of at most ``max_snapshots`` rows the + kernel's candidate count holds.""" + weight = getattr(o_proj, "weight", None) + return ( + getattr(o_proj, "bias", None) is None + and isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.is_cuda + and _sandwich_op.supports(weight.new_empty((1, weight.shape[1])), weight, max_snapshots) + ) + + @staticmethod + def takes_tail( + latent: torch.Tensor, act: torch.Tensor, tail_weight: torch.Tensor, max_snapshots: int + ) -> bool: + """Whether ``sandwich_tail`` takes a MoE tail of these tensors' shapes (rows of ``latent`` and ``act``, at most + `SANDWICH_MAX_TOKENS`; ``tail_weight``) with a snapshot bank of at most ``max_snapshots`` rows: the TP16 + per-rank shapes.""" + return _sandwich_op.supports_tail(latent, act, tail_weight, max_snapshots) + + def allreduce_attn_res( + self, + partial: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_allreduce_attn_res`` call: ``updated = prefix_sum + + allreduce(partial)`` (the sum alone without ``prefix_sum``), ``normed = out_norm(attn_res(block_residual..., + updated))``, the attention residual selected with ``res_proj`` and ``res_norm``. ``block_residual`` holds the + valid snapshots, ``[S, T, H]``.""" + return mnnvl_allreduce_attn_res( + partial.contiguous(), + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.mnnvl, + ) + + def sandwich_tail( + self, + pending: PendingTail, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + updated_out: Optional[torch.Tensor] = None, + tap: Optional[torch.Tensor] = None, + tap_updated: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of a MoE layer's row-parallel tail ``pending`` + (``[RMSNorm(latent)[:, lo:lo + 224] | act] @ weight.T``), in one ``comm/k3_sandwich_tail`` call. + ``updated_out``: a bf16 ``[T, H]`` tensor the call stores ``updated`` into (the consumer's snapshot bank row), + returned as ``updated``. ``tap``: a bf16 ``[T, H]`` view (a speculative capture slot) the call also stores + the pre-norm attention-residual mixture into, or ``updated`` with ``tap_updated``.""" + return k3_sandwich_tail( + pending.latent, + pending.act, + pending.weight, + pending.lo, + pending.lat_eps, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + tap=tap, + tap_updated=tap_updated, + updated_out=updated_out, + ) + + def sandwich_oproj( + self, + core: torch.Tensor, + o_weight: torch.Tensor, + prefix_sum: Optional[torch.Tensor], + block_residual: torch.Tensor, + res_proj: nn.Module, + res_norm: nn.Module, + out_norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` as ``allreduce_attn_res`` of ``core @ o_weight.T``, in one ``comm/k3_sandwich_oproj`` + call: ``core`` is this rank's gated o_proj input ``[T, 768]``, ``o_weight`` its o_proj slice. Bit for bit + o_proj followed by ``allreduce_attn_res``.""" + return k3_sandwich_oproj( + core.contiguous(), + o_weight, + prefix_sum, + block_residual, + *_res_args(res_proj, res_norm, out_norm), + self.sandwich, + ) + + def compile_plain(self, weight: torch.Tensor, swiglu: bool = False) -> bool: + """Compile the plain sandwich (``sandwich_plain``) for a projection of ``weight``'s shape (with ``swiglu``, a + down projection after the SiLU-and-mul) with one call on a zero row of a zero weight, before any capture; + False, with nothing launched, where the kernel does not take that shape. Collective: every rank of the group + makes the call; it advances the sandwich workspace on every rank alike.""" + zero = torch.zeros_like(weight) + hidden, k_in = zero.shape + x = zero.new_zeros(1, 2 * k_in if swiglu else k_in) + residual = zero.new_zeros(1, hidden) + ones = zero.new_ones(hidden) + if not _sandwich_op.supports_plain(x, zero, residual, ones, swiglu): + return False + k3_sandwich_plain(x, zero, residual, ones, 1e-6, self.sandwich, swiglu=swiglu) + torch.cuda.synchronize(zero.device) + return True + + @staticmethod + def takes_plain( + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> bool: + """Whether ``sandwich_plain`` takes the call: at most `SANDWICH_MAX_TOKENS` contiguous bf16 rows ``x`` of a + row-parallel slice ``weight`` [7168, K] (K a multiple of 128 up to 896; with ``swiglu`` K 896 and ``x`` the + ``[gate | up]`` rows, 2 K wide), ``residual`` [rows, 7168], and ``norm``'s bf16 [7168] weight.""" + return _sandwich_op.supports_plain(x, weight, residual, norm.weight, swiglu) + + def sandwich_plain( + self, + x: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + swiglu: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/k3_sandwich_plain`` call: ``updated = residual + allreduce(x @ + weight.T)`` (with ``swiglu``, of ``silu_and_mul(x) @ weight.T``), ``normed = norm(updated)`` for a plain + RMSNorm ``norm``. Bit for bit, by the kernel's statement, the projection on ``k3_ctm_gemv`` at split 1 (with + ``swiglu``, ``k3_ctm_gemv_swiglu`` at split 2) followed by ``allreduce_norm``'s one-shot call.""" + return k3_sandwich_plain( + x.contiguous(), + weight, + residual.contiguous(), + norm.weight, + _eps(norm), + self.sandwich, + swiglu=swiglu, + ) + + def takes_allreduce_norm(self, rows: int, hidden: int) -> bool: + """Whether ``allreduce_norm`` takes ``rows`` bf16 rows of ``hidden`` columns: the MNNVL workspace holds the + call at the decode path's one-shot ceiling (`DECODE_AR_ONE_SHOT_MAX_BYTES`, two-shot above it).""" + if rows <= 0 or hidden <= 0 or hidden % 8: + return False + world, buffer_bytes = self.mnnvl.world_size, self.mnnvl.buffer_bytes + need = required_buffer_bytes( + rows, hidden, world, torch.bfloat16, DECODE_AR_ONE_SHOT_MAX_BYTES + ) + two_shot = rows * hidden * world * 2 > DECODE_AR_ONE_SHOT_MAX_BYTES + return need <= buffer_bytes and not (two_shot and buffer_bytes % 32) + + def allreduce_norm( + self, partial: torch.Tensor, residual: torch.Tensor, norm: nn.Module + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(normed, updated)`` in one ``comm/mnnvl_fusion_allreduce`` call: ``updated = residual + + allreduce(partial)``, ``normed = norm(updated)`` for a plain RMSNorm ``norm``, sent one-shot up to + `DECODE_AR_ONE_SHOT_MAX_BYTES`. ``partial`` is this rank's unreduced ``[rows, hidden]`` bf16 output of a + row-parallel projection, of a shape ``takes_allreduce_norm`` holds for.""" + return mnnvl_fusion_allreduce( + partial.contiguous(), + self.mnnvl, + DECODE_AR_ONE_SHOT_MAX_BYTES, + residual.contiguous(), + norm.weight, + _eps(norm), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_gemv.py new file mode 100644 index 000000000000..301b709643b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_gemv.py @@ -0,0 +1,465 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's GEMVs, LM head and embedding on the catalog's single-GPU Kimi K3 entries. + +* **Per-site GEMVs** (`K3DecodeGemvs.project`): a projection of a decode step runs on the kernel measured fastest + at its call site's weight shape (`SITES`): at most `MAX_ROWS` rows on `gemm/k3_decode_gemv`, + `gemm/k3_ctm_gemv_wide`, `gemm/k3_ctm_gemv_long`, `gemm/k3_ctm_gemv` or `gemm/k3_ctm_gemv_swiglu`, and, where the + site lists it, more rows (up to `WIDE_ROWS`) on `gemm/k3_ctm_gemv_wide`. The sites are the decode kernels' fused + projections (MLA's [W_a; W_g] with the gate rows through a sigmoid, KDA's [q | k | v | g | f_a | b]), the attention + output projection, the built-in MLA path's q_a / kv_a, q_b and gate projections, the MoE decode path's projections + (`decode_moe.py`), two of them with fp32 outputs, and the DSpark drafter's q / k / v, output, gate / up and down + projections (`K3DSparkDrafter`). +* **LM head** (`K3LogitsProcessor`): at most `MAX_ROWS` rows of this rank's vocabulary shard on + `gemm/k3_head_gemv` over the target's `K3HeadGemvWorkspace`, then the shards gathered (`comm/allgather`) as the + stock head gathers them. It is the shell's logits processor, so the speculative worker's target logits and the + drafter's logits on the same head take it too. +* **Embedding** (`K3DecodeGemvs.embed_norm`): a decode step's embedding rows written into the attention-residual + bank's slot 0 (layer 0's first snapshot) and layer 0's input RMSNorm applied, in one `norm/k3_embed_norm` launch. +* **Dense MLP** (`K3DecodeGemvs.dense_mlp`): layer 0's MLP at most `MAX_ROWS` rows, split over the whole TP group: + gate_up on `gemm/k3_ctm_gemv_long`, `activation/k3_situ_mul`, down on `gemm/k3_ctm_gemv_long`; the caller then + runs the down projection's all-reduce. + +Each returns None where its kernel does not take the call, and the caller then runs the generic path's module. + +A kernel compiles on its first call for a shape, which must not happen under CUDA-graph capture. +`K3DecodeGemvs.create`, run once the weights are final, runs every site's kernel and the head once, eagerly. The +embedding kernel compiles per token count, on the eager warm-up step before each capture. Under capture, a call +whose kernel has not run eagerly is refused. + +The head's workspace serves every `k3_head_gemv` call of its weight shape, so those calls must be ordered on one +stream: the logits are computed on the model's stream (see the `gemm/k3_head_gemv` contract's State section). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Dict, Iterable, Optional, Set + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.k3_situ_mul import k3_situ_mul +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.allgather import allgather +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv import k3_ctm_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_long import ( + k3_ctm_gemv_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_swiglu import ( + k3_ctm_gemv_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_ctm_gemv_wide import ( + k3_ctm_gemv_wide, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_decode_gemv import k3_decode_gemv +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.k3_head_gemv import ( + K3HeadGemvWorkspace, + k3_head_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.k3_embed_norm import k3_embed_norm + +# The kernels' support predicates: metadata reads only. +from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import op as _ctm_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_decode_gemv import op as _decode_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_embed import op as _embed_op +from tensorrt_llm._torch.cute_dsl_kernels.k3_head_gemv import op as _head_op +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE +from tensorrt_llm._utils import mpi_disabled + +# The row limit of the decode GEMV and head kernels (one token tile), and of k3_ctm_gemv_wide (a decode step of 8 +# requests of 8 tokens). +MAX_ROWS = 8 +WIDE_ROWS = 64 + + +@dataclass(frozen=True) +class Site: + """A call site's weight shape (this target's per-rank shapes) and its kernels: ``small`` at 1..MAX_ROWS rows + ("decode", "wide", "long", "ctm" for k3_ctm_gemv, or "swiglu" for k3_ctm_gemv_swiglu, whose rows are the + ``[gate | up]`` input of the SiLU-and-mul, 2 ``k`` wide), and k3_ctm_gemv_wide at MAX_ROWS+1..``wide_rows`` rows + where ``wide``. Output columns from ``sig_col0`` on are stored through a sigmoid; ``out_fp32``: an fp32 output + (k3_ctm_gemv_wide only). ``split`` / ``ring`` / ``push``: the CTM kernels' CTAs per 128-row weight tile, + k3_ctm_gemv_long's weight-ring stages, and whether the partial sums are pushed to each row's owner.""" + + n: int + k: int + small: str + wide: bool = False + sig_col0: int = -1 + split: int = 0 + ring: int = 0 + push: bool = False + out_fp32: bool = False + wide_rows: int = WIDE_ROWS + + +SITES: Dict[str, Site] = { + # MLA's [W_a; W_g] on a decode step: [q_a 1536 | kv_a 512 | k_pe 64] then the output gate (6 heads x 128), + # the gate rows through a sigmoid. + "mla_ag": Site(2880, 7168, "long", wide=True, sig_col0=2112, split=6, ring=6, push=True), + # KDA's [q | k | v | g | f_a | b] on a decode step (6 heads x 128 each, then 128 and 6, padded to 3208 rows). + "kda_proj": Site(3208, 7168, "long", wide=True, split=5, ring=6, push=True), + # The attention output projection (row parallel; 6 heads x 128 in). + "o_proj": Site(7168, 768, "decode", wide=True), + # The built-in MLA path's projections: kv_a_proj_with_mqa, q_b_proj and the output gate. + "kv_a": Site(2112, 7168, "decode"), + "q_b": Site(1152, 1536, "wide"), + "g_proj": Site(768, 7168, "wide"), + # Layer 0's dense MLP split over the 16-way TP group: gate_up [gate 2112 | up 2112] and down. + "dense_gate_up": Site(4224, 7168, "long", split=4, ring=5), + "dense_down": Site(7168, 2112, "long", split=2, ring=6), + # The MoE decode path (decode_moe.py): this rank's head slice [latent down 224 | router 56] with an fp32 output, + # the shared experts' gate_up, the row-parallel tail [latent up 224 | padding 32 | shared down 384] and the + # replicated tail's latent up projection with an fp32 output. The wide kernel takes the head slice and the tail up + # to 32 rows, where it is faster than the stock GEMM. + "moe_head": Site(280, 7168, "wide", wide=True, out_fp32=True, wide_rows=32), + "moe_shared_gate_up": Site(768, 7168, "wide", wide=True), + "moe_tail": Site(7168, 640, "wide", wide=True, wide_rows=32), + "moe_up": Site(7168, 3584, "wide", out_fp32=True), + # The DSpark drafter's projections (K3DSparkDrafter; 6 query heads, 1 KV head of 64, MLP width 896): the fused + # q / k / v, the output projection (row parallel), gate_up [gate 896 | up 896], and the down projection with the + # SiLU-and-mul (row parallel). + "drafter_qkv": Site(512, 7168, "long", split=8, ring=6, push=True), + "drafter_o": Site(7168, 384, "ctm", split=1, push=True), + "drafter_gate_up": Site(1792, 7168, "long", split=8, ring=6, push=True), + "drafter_down": Site(7168, 896, "swiglu", split=2, push=True), +} + + +def _width(spec: Site) -> int: + """The row width a site's kernel reads: ``k``, or the ``[gate | up]`` input of a SiLU-and-mul site, 2 ``k``.""" + return 2 * spec.k if spec.small == "swiglu" else spec.k + + +def _wide_tile(rows: int) -> int: + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctm_gemv import k3_ctm_gemv_kernel + + return k3_ctm_gemv_kernel.wide_tile(rows) + + +def _run( + spec: Site, kernel: str, x2d: torch.Tensor, weight: torch.Tensor +) -> Optional[torch.Tensor]: + """``spec``'s ``kernel`` on dense rows ``x2d``, or None where it does not take them.""" + if kernel == "decode": + if spec.sig_col0 >= 0 or spec.out_fp32 or not _decode_op.supports(x2d, weight): + return None + return k3_decode_gemv(x2d, weight) + if kernel == "wide": + if not _ctm_op.supports_wide(x2d, weight, spec.sig_col0, spec.out_fp32): + return None + return k3_ctm_gemv_wide(x2d, weight, sig_col0=spec.sig_col0, out_fp32=spec.out_fp32) + if spec.out_fp32: + return None + if kernel == "ctm": + if spec.sig_col0 >= 0 or not _ctm_op.supports(x2d, weight, spec.split): + return None + return k3_ctm_gemv(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + if kernel == "swiglu": + if spec.sig_col0 >= 0 or not _ctm_op.supports_swiglu(x2d, weight, spec.split): + return None + return k3_ctm_gemv_swiglu(x2d, weight, trigger_early=True, split=spec.split, push=spec.push) + # One wave of the GPU's SMs: beyond it the long GEMV loses to the others. + sms = torch.cuda.get_device_properties(x2d.device).multi_processor_count + if math.ceil(spec.n / 128) * spec.split > sms or not _ctm_op.supports_long( + x2d, weight, spec.split, spec.ring + ): + return None + return k3_ctm_gemv_long( + x2d, + weight, + sig_col0=spec.sig_col0, + split=spec.split, + ring=spec.ring, + trigger_early=True, + push=spec.push, + ) + + +def _capturing() -> bool: + return not torch.compiler.is_compiling() and torch.cuda.is_current_stream_capturing() + + +def _dense_rows(x2d: torch.Tensor) -> torch.Tensor: + """``x2d`` itself, or a dense copy: the kernels' TMA descriptors need dense, 16-byte-aligned rows.""" + if x2d.stride() != (x2d.shape[1], 1) or x2d.data_ptr() % 16: + return x2d.clone(memory_format=torch.contiguous_format) + return x2d + + +def _head_takes_module(lm_head: nn.Module) -> bool: + """Whether ``lm_head(rows)`` is a plain vocabulary-parallel GEMM of a bf16 weight whose shards it gathers along + the vocabulary in rank order, with nothing else applied: the stock head this path reproduces.""" + weight = getattr(lm_head, "weight", None) + mapping = getattr(lm_head, "mapping", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and weight.is_contiguous() + and getattr(getattr(lm_head, "tp_mode", None), "name", None) == "COLUMN" + and getattr(lm_head, "gather_output", False) + and getattr(lm_head, "gather_output_sizes", None) is None + and getattr(lm_head, "padding_size", None) == 0 + and getattr(lm_head, "bias", None) is None + and not getattr(lm_head, "has_any_quant", True) + and mapping is not None + and not mapping.enable_attention_dp + ) + + +def _plain_rmsnorm(norm: nn.Module) -> bool: + """Whether ``norm`` is the stock RMSNorm that runs flashinfer's kernel, which ``k3_embed_norm`` reproduces bit for + bit. An unknown module fails closed.""" + weight = getattr(norm, "weight", None) + return ( + IS_FLASHINFER_AVAILABLE + and isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "use_cuda_tile", True) + and not getattr(norm, "return_hp_output", True) + and getattr(norm, "nvfp4_scale", None) is None + and hasattr(norm, "variance_epsilon") + ) + + +class K3DecodeGemvs: + """The decode GEMVs' state for one target: the LM head's `K3HeadGemvWorkspace`, and the calls whose kernels ran + eagerly (compiled), which are the only ones a CUDA-graph capture may take. Built by `create` once the weights are + final; owned by the target.""" + + def __init__(self, head_workspace: Optional[K3HeadGemvWorkspace] = None) -> None: + self.head_workspace = head_workspace + self._ran: Set[tuple] = set() + + @classmethod + def create( + cls, + lm_head: Optional[nn.Module] = None, + sites: Iterable[str] = tuple(SITES), + device: Optional[torch.device] = None, + ) -> "K3DecodeGemvs": + """The state for ``lm_head`` and ``sites`` on ``device`` (default: the head's, else the current one). Eager: + it allocates the head's workspace and runs every kernel of every site once (each wide row class once) on + zero rows of a zero weight of the site's shape, so they compile here and not under a capture. A site or + head whose kernel does not take its shape keeps the generic path.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "K3DecodeGemvs.create allocates and compiles: run it before CUDA-graph capture" + ) + head_weight = getattr(lm_head, "weight", None) if lm_head is not None else None + if device is None: + device = ( + head_weight.device + if isinstance(head_weight, torch.Tensor) + else torch.device("cuda", torch.cuda.current_device()) + ) + state = cls() + for site in sites: + spec = SITES[site] + weight = torch.zeros(spec.n, spec.k, dtype=torch.bfloat16, device=device) + for rows in (1, 16, 32, 64) if spec.wide else (1,): + if rows > spec.wide_rows: + continue + state._project(site, weight.new_zeros(rows, _width(spec)), weight, warm=True) + del weight + if "dense_gate_up" in sites: + gu = torch.zeros(1, SITES["dense_gate_up"].n, dtype=torch.bfloat16, device=device) + if _ctm_op.supports_situ_mul(gu): + for linear_beta in (None, 1.0): + k3_situ_mul(gu, 1.0, linear_beta) + state._ran.add(("situ_mul", linear_beta is not None)) + if lm_head is not None and _head_takes_module(lm_head): + x = head_weight.new_zeros(1, head_weight.shape[1]) + if _head_op.supports(x, head_weight): + workspace = K3HeadGemvWorkspace.create( + head_weight.shape[0], head_weight.shape[1], head_weight.device + ) + k3_head_gemv(x, head_weight, workspace) + state.head_workspace = workspace + state._ran.add(("lm_head",)) + torch.cuda.synchronize(device) + return state + + def project(self, site: str, x: torch.Tensor, weight: torch.Tensor) -> Optional[torch.Tensor]: + """``x @ weight.T`` (``[..., N]``: bf16, the site's sigmoid columns through the sigmoid; fp32 at an + ``out_fp32`` site; on a SiLU-and-mul site, ``silu(gate) * up`` of ``x = [gate | up]`` first) for ``site``'s + weight on its decode kernel, or None where none takes the call: more rows than the site's kernels take, + another shape or dtype, or, under capture, a kernel that has not run eagerly. The caller then runs its + module.""" + return self._project(site, x, weight, warm=False) + + def _project( + self, site: str, x: torch.Tensor, weight: torch.Tensor, warm: bool + ) -> Optional[torch.Tensor]: + spec = SITES[site] + width = _width(spec) + if ( + weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (spec.n, spec.k) + or not weight.is_contiguous() + or x.dtype != torch.bfloat16 + or x.dim() < 1 + or x.shape[-1] != width + ): + return None + rows = x.numel() // width + if 0 < rows <= MAX_ROWS: + kernel = spec.small + elif spec.wide and MAX_ROWS < rows <= spec.wide_rows: + kernel = "wide" + else: + return None + key = (site, kernel, _wide_tile(rows) if kernel == "wide" else 0) + capturing = _capturing() + if capturing and not warm and key not in self._ran: + return None + y = _run(spec, kernel, _dense_rows(x.reshape(rows, width)), weight) + if y is None: + return None + if not capturing: + self._ran.add(key) + return y.view(*x.shape[:-1], spec.n) + + def lm_head_logits( + self, rows: torch.Tensor, lm_head: nn.Module, gather: bool = True + ) -> Optional[torch.Tensor]: + """``lm_head(rows)``, the gathered bf16 logits ``[M, vocab]``, with this rank's shard on + ``gemm/k3_head_gemv`` (``gather`` False: the shard alone, ``[M, vocab / TP]``); None where it does not take + the call (more than `MAX_ROWS` rows, another head, ...).""" + workspace = self.head_workspace + if workspace is None or not _head_takes_module(lm_head): + return None + weight = lm_head.weight + if ( + tuple(weight.shape) != (workspace.n_out, workspace.k_in) + or weight.device != workspace.partials.device + or rows.dim() != 2 + or rows.dtype != torch.bfloat16 + or not 0 < rows.shape[0] <= MAX_ROWS + or rows.shape[1] != workspace.k_in + ): + return None + if _capturing() and ("lm_head",) not in self._ran: + return None + group = lm_head.mapping.tp_group + if gather and len(group) > 1 and mpi_disabled(): + return None + x = _dense_rows(rows) + if not _head_op.supports(x, weight): + return None + local = k3_head_gemv(x, weight, workspace) + if not gather or len(group) == 1: + return local + gathered = allgather(local, None, group) + return torch.cat(torch.split(gathered, rows.shape[0], dim=0), dim=-1) + + def dense_mlp( + self, + x: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + beta: float, + linear_beta: Optional[float], + ) -> Optional[torch.Tensor]: + """Layer 0's dense MLP of at most `MAX_ROWS` rows, before its down projection's all-reduce: gate_up on + ``gemm/k3_ctm_gemv_long``, ``SituAndMul(beta, linear_beta)`` on ``activation/k3_situ_mul``, down on + ``gemm/k3_ctm_gemv_long``. None, with nothing launched, where a kernel does not take the call; the caller then + runs the module.""" + gate_up, down = SITES["dense_gate_up"], SITES["dense_down"] + rows = x.shape[0] if x.dim() == 2 else 0 + if ( + not 0 < rows <= MAX_ROWS + or linear_beta == 0.0 + or tuple(gate_up_weight.shape) != (gate_up.n, gate_up.k) + or tuple(down_weight.shape) != (down.n, down.k) + or gate_up.n != 2 * down.k + ): + return None + if ( + _capturing() + and not { + ("dense_gate_up", "long", 0), + ("dense_down", "long", 0), + ("situ_mul", linear_beta is not None), + } + <= self._ran + ): + return None + gu = self.project("dense_gate_up", x, gate_up_weight) + if gu is None or not _ctm_op.supports_situ_mul(gu): + return None + return self.project("dense_down", k3_situ_mul(gu, beta, linear_beta), down_weight) + + def embed_norm( + self, + input_ids: torch.Tensor, + table: torch.Tensor, + norm: nn.Module, + bank: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Layer 0's normed input for the step's tokens, with their embedding rows written into ``bank[0]``, in one + ``norm/k3_embed_norm`` launch: bit-identical to the embedding followed by ``norm``. None where it does not + apply (another norm, a token count or table the kernel does not take, or, under capture, a token count whose + kernel has not run eagerly).""" + if not _plain_rmsnorm(norm): + return None + ids = input_ids.reshape(-1) + if not _embed_op.supports_norm(ids, table, norm.weight, bank[0]): + return None + key = ("embed_norm", ids.numel(), ids.dtype) + capturing = _capturing() + if capturing and key not in self._ran: + return None + normed = k3_embed_norm(ids, table, norm.weight, norm.variance_epsilon, bank[0]) + if not capturing: + self._ran.add(key) + return normed + + +class _K3Head: + """``lm_head(rows)``, with the rows on ``k3_head_gemv`` where it takes them.""" + + __slots__ = ("_gemvs", "_lm_head") + + def __init__(self, gemvs: K3DecodeGemvs, lm_head: nn.Module) -> None: + self._gemvs = gemvs + self._lm_head = lm_head + + def __call__(self, rows: torch.Tensor) -> torch.Tensor: + logits = self._gemvs.lm_head_logits(rows, self._lm_head) + return self._lm_head(rows) if logits is None else logits + + +class K3LogitsProcessor(nn.Module): + """The shell's logits processor, with its LM head call on ``gemm/k3_head_gemv`` where ``gemvs`` takes the rows. + + It wraps the stock processor instead of repeating it: the row selection and the fp32 conversion stay the stock + processor's. ``gemvs`` is set once the weights are final; until then every call is the stock one. + """ + + def __init__(self, stock: nn.Module) -> None: + super().__init__() + self.stock = stock + self.gemvs: Optional[K3DecodeGemvs] = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: nn.Module, + attn_metadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + head = lm_head if self.gemvs is None else _K3Head(self.gemvs, lm_head) + return self.stock.forward(hidden_states, head, attn_metadata, return_context_logits) + + def lm_head_shard(self, rows: torch.Tensor, lm_head: nn.Module) -> Optional[torch.Tensor]: + """This rank's bf16 vocabulary shard of ``lm_head(rows)`` on ``gemm/k3_head_gemv``, without the gather, or + None where ``gemvs`` does not take the rows. A speculative worker that keeps its logits vocabulary-sharded + computes them here, so they hold the values of the logits this processor gathers.""" + if self.gemvs is None: + return None + return self.gemvs.lm_head_logits(rows, lm_head, gather=False) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_moe.py new file mode 100644 index 000000000000..43f643d49f40 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/decode_moe.py @@ -0,0 +1,408 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode path's latent MoE on the catalog's Kimi K3 MoE entries. + +**At most 8 tokens** (`MAX_TOKENS`), a MoE layer runs: + +* `moe/k3_moe_front`, one kernel: this rank's slice of the MoE head (its latent-down rows and its router rows) as one + GEMV, the slices' all-gather over the TP group's `K3MoeHeadWorkspace`, the top-16 routing, the MXFP8 latent, and + the shared experts' gate_up + SiTU; +* `moe/k3_moe`: this rank's routed partial over its experts (on a `K3MoeState`); +* the latent all-reduce. On a pushing step (`DecodeStep.latent_push`: a pure decode step captured into a CUDA graph + whose attention layers all run the decode kernels) the routed experts run as `moe/k3_moe`'s push form, which stores + this rank's partial into every rank's `K3LatentExchange`, and `comm/k3_latent_reduce` sums the partials in the MNNVL + one-shot's order, so with the same bits, while the experts' grid completes. On every other step it is the routed + experts' all-reduce (one-shot, as `decode_comm.use_decode_one_shot` sets); +* the tail. The latent norm's weight is folded into the latent up projection at load, so + `[RMSNorm(latent) slice | shared activation] @ [latent up columns | shared down]` is this rank's share of the MoE + output. Where the next layer's pre-attention step (or the final norm) reduces it, the layer hands it on unreduced + as a `PendingTail`, which the consumer runs with its all-reduce and residual update as one `comm/k3_sandwich_tail` + kernel (`decode_comm.py`). Elsewhere the replicated tail runs: the latent RMS applied to the fp32 output of one + GEMV with the folded latent up weight, plus the shared experts' down projection and its all-reduce. + +**A wide decode step** (9 to 64 tokens, `WIDE_MAX_TOKENS`) keeps the sharded head and the row-parallel tail on +M-general ops: the head GEMV, `comm/mnnvl_allgather_split`, then `moe/k3_route_quant` and `moe/k3_moe` (on a +`K3MoeWideState`) beside the shared gate_up + SiTU, the latent all-reduce, and one GEMV of +`[RMSNorm(latent) slice | padding | shared activation]` with the tail weight. That is this rank's unreduced share, +which the consumer reduces with a plain all-reduce. + +The GEMVs run on the decode GEMV sites of `decode_gemv.py` where they take the call, else on the stock GEMM ops. + +`K3DecodeMoe` holds what every MoE layer shares: the head workspace and the latent exchange (both collective over +the TP group), the two `k3_moe` builds' scratch, and the TP group's MNNVL workspace (`decode_comm.K3DecodeComm`'s). +`K3DecodeMoeLayer` holds one layer's decode weights and its `k3_moe` counters. The target builds both in +`post_load_weights`, before any CUDA-graph capture, and runs every kernel once there so none compiles under a capture. + +Every MoE layer's push and reduce go to the one exchange, in the stream's order: each push is followed by exactly one +reduce of the same token count before the next push, on every rank in the same order (the exchange's call-order +invariant, `comm/k3_latent_reduce`). Every kernel between a reduce and the next push waits for its predecessor (or +launches without programmatic dependent launch), which the decode kernels a pushing step runs do; pushes run only in +CUDA-graph replays, so no other step's kernels sit between two pushes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.k3_latent_reduce import ( + K3LatentExchange, + k3_latent_reduce, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_allgather_split import ( + mnnvl_allgather_split, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.comm.mnnvl_workspace import ( + MnnvlWorkspace, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe import ( + K3MoeHeadWorkspace, + K3MoeLayer, + K3MoeState, + K3MoeWideState, + is_supported, + k3_moe, + k3_moe_push, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_moe_front import ( + front_weight, + k3_moe_front, + weight_supported, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.k3_route_quant import k3_route_quant +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel + +from .decode_comm import K3DecodeComm, PendingTail, wide_all_reduce + +# The most tokens of the front and the small k3_moe build (one token tile), and of the wide build. +MAX_TOKENS = 8 +WIDE_MAX_TOKENS = 64 + +# The tail weight's latent columns are zero-padded to whole 128-column k-tiles of the tail kernels +# (comm/k3_sandwich_tail takes the TP16 tail weight [7168, 256 + 384]). +TAIL_K_TILE = 256 + + +@dataclass(eq=False) +class K3DecodeMoe: + """What every MoE layer's decode path shares on one device: the TP group's ``K3MoeHeadWorkspace`` (the front's + all-gather), the ``k3_moe`` builds for up to 8 and up to 64 tokens (their scratch; the layers run one at a time on + one stream), the TP group's ``MnnvlWorkspace`` for a wide step's head all-gather, and the TP group's + ``K3LatentExchange`` for a pushing step's latent all-reduce (None: every step keeps the routed experts' + all-reduce). Built by `create`.""" + + head: K3MoeHeadWorkspace + small: K3MoeState + wide: K3MoeWideState + mnnvl: MnnvlWorkspace + exchange: Optional[K3LatentExchange] = None + + @classmethod + def create( + cls, mapping, device, i_tp: int, num_local: int, mnnvl: MnnvlWorkspace, push: bool = True + ) -> "K3DecodeMoe": + """The state for ``mapping``'s TP group on ``device``, for experts of ``i_tp`` intermediate columns per rank, + ``num_local`` of them on this rank. Collective (the head workspace and the latent exchange): every rank of the + group calls it at the same point, eagerly, before any CUDA-graph capture. ``push``: build the latent exchange + (where it takes the TP size: 4, 8 or 16 ranks); without it every step keeps the routed experts' all-reduce.""" + head = K3MoeHeadWorkspace.create(mapping) + exchange = None + if push: + try: + exchange = K3LatentExchange.create(mapping) + # Raised before any collective step, on every rank of the group alike: a TP size the exchange does not take. + except ValueError: + exchange = None + return cls( + head, + K3MoeState(device, i_tp, num_local), + K3MoeWideState(device, i_tp, num_local), + mnnvl, + exchange, + ) + + +def _experts(moe: nn.Module) -> tuple: + """The routed experts' TRTLLM-Gen W4A8_MXFP4_MXFP8 buffers ``k3_moe`` reads in place.""" + backend = moe.routed_experts.backend + return ( + backend.w3_w1_weight, + backend.w3_w1_weight_scale, + backend.w2_weight, + backend.w2_weight_scale, + ) + + +def layout_gaps(moe: nn.Module, tp_size: int, max_snapshots: int) -> list: + """Why the decode path does not take MoE layer ``moe`` (empty when it takes it), read once the routed experts' + buffers exist: the checks in order, up to the first that fails. ``max_snapshots``: the model's snapshot bank + rows.""" + gate, backend = moe.gate, moe.routed_experts.backend + shared = moe.shared_experts + gate_up, down = shared.gate_up_proj.weight, shared.down_proj.weight + if not moe._reduce_routed_output: + return ["a routed output the model does not reduce"] + if (moe.num_experts, moe.top_k, moe.moe_hidden_size, moe.hidden_size) != (896, 16, 3584, 7168): + return [ + f"experts / top-k / latent / hidden {(moe.num_experts, moe.top_k, moe.moe_hidden_size)}" + ] + if (gate.num_expert_group, gate.topk_group) != (1, 1): + return ["grouped routing"] + if moe.moe_hidden_size % (8 * tp_size) or moe.num_experts % (4 * tp_size): + return [f"a head slice of TP {tp_size}"] + latent = (gate.weight, moe.routed_expert_down_proj.weight, moe.routed_expert_up_proj.weight) + if not isinstance(moe.routed_expert_up_proj, nn.Linear) or any( + w.dtype != torch.bfloat16 for w in latent + ): + return ["latent projections or a router other than bf16"] + if ( + gate_up.dtype != torch.bfloat16 + or down.dtype != torch.bfloat16 + or shared.gate_up_proj.bias is not None + or shared.down_proj.bias is not None + or down.shape[0] != moe.hidden_size + ): + return ["a shared expert other than bf16 and unbiased"] + names = ( + "w3_w1_weight", + "w3_w1_weight_scale", + "w2_weight", + "w2_weight_scale", + "expert_size_per_partition", + ) + if ( + not all(hasattr(backend, name) for name in names) + or not is_supported(*_experts(moe), backend.expert_size_per_partition)[0] + ): + return ["routed-expert buffers k3_moe does not read"] + shared_cols, width = gate_up.shape[0] // 2, moe.moe_hidden_size // tp_size + if not weight_supported(tp_size, shared_cols, gate_up.shape[1], gate_up.device): + return ["a MoE front of this TP size and shared width"] + tail_cols = width + (-width % TAIL_K_TILE) + shared_cols + probe = gate_up.new_empty + if not K3DecodeComm.takes_tail( + probe(1, moe.moe_hidden_size), + probe(1, shared_cols), + probe(moe.hidden_size, tail_cols), + max_snapshots, + ): + return ["a row-parallel tail the sandwich tail kernel does not take"] + return [] + + +def fold_latent_norm(moe: nn.Module) -> None: + """Fold the latent RMSNorm's weight into the latent up projection's columns; the norm keeps a weight of ones, so + every step computes the same function and the tails may normalize before slicing.""" + up, norm = moe.routed_expert_up_proj, moe.routed_expert_norm + with torch.no_grad(): + up.weight.mul_(norm.weight.to(up.weight.dtype)[None, :]) + norm.weight.fill_(1) + + +@dataclass(eq=False) +class K3DecodeMoeLayer: + """One MoE layer's decode path: the front weight (this rank's head slice padded to whole tiles, then the shared + gate_up re-ordered), the head slice (a view of it), the row-parallel tail weight + ``[latent up columns lo:lo+width | padding | shared down]``, the routing bias, and its ``k3_moe`` handles on the + shared builds. Built by `create` once the weights are final (and the latent norm folded).""" + + state: K3DecodeMoe + front_weight: torch.Tensor + head_weight: torch.Tensor + tail_weight: torch.Tensor + tail_pad: Optional[torch.Tensor] + # The gate's routing bias detached (a view of the parameter): an op given a tensor that requires grad while + # autograd records, as in post_load_weights, returns outputs that require grad, which k3_moe cannot take. + bias: torch.Tensor + lo: int + width: int + shared_cols: int + small: K3MoeLayer + wide: K3MoeLayer + + @classmethod + def create( + cls, moe: nn.Module, state: K3DecodeMoe, tp_rank: int, tp_size: int + ) -> "K3DecodeMoeLayer": + """Build MoE layer ``moe``'s decode weights and its handles on ``state``'s builds (``layout_gaps`` empty).""" + width = moe.moe_hidden_size // tp_size + experts = moe.num_experts // tp_size + gate_up = moe.shared_experts.gate_up_proj.weight + shared_down = moe.shared_experts.down_proj.weight + up = moe.routed_expert_up_proj.weight + lo = tp_rank * width + pad = -width % TAIL_K_TILE + with torch.no_grad(): + head = torch.cat( + [ + moe.routed_expert_down_proj.weight[lo : lo + width], + moe.gate.weight[tp_rank * experts : (tp_rank + 1) * experts], + ] + ) + front = front_weight(head, gate_up) + parts = [up[:, lo : lo + width]] + if pad: + parts.append(up.new_zeros(up.shape[0], pad)) + parts.append(shared_down) + tail = torch.cat(parts, dim=1).contiguous() + weights = _experts(moe) + return cls( + state=state, + front_weight=front, + head_weight=front[: head.shape[0]], + tail_weight=tail, + tail_pad=up.new_zeros(WIDE_MAX_TOKENS, pad) if pad else None, + bias=moe.gate.e_score_correction_bias.detach(), + lo=lo, + width=width, + shared_cols=gate_up.shape[0] // 2, + small=state.small.layer(*weights), + wide=state.wide.layer(*weights), + ) + + def warm_up(self, moe: nn.Module) -> None: + """One call of every kernel of the decode path on zero inputs (M = 1), so none compiles under a capture: the + front (collective: every rank makes the same call), both ``k3_moe`` builds, ``k3_route_quant`` and, with the + latent exchange, the push build and the reduce (one push + reduce pair, collective). Once per model: every + layer's calls compile the same builds.""" + device = self.front_weight.device + x = torch.zeros(1, moe.hidden_size, dtype=torch.bfloat16, device=device) + ids, weights, x_fp8, x_sf, _ = self._front(moe, x) + offset = moe.routed_experts.backend.slot_start + k3_moe(x_fp8, x_sf, ids, weights, offset, self.small) + exchange = self.state.exchange + if exchange is not None: + # Two calls on one layer must not run back to back (moe/k3_moe): the push build starts once the call + # above has ended. + torch.cuda.synchronize(device) + k3_moe_push(x_fp8, x_sf, ids, weights, offset, self.small, exchange) + k3_latent_reduce(1, exchange) + logits = torch.zeros(1, moe.num_experts, dtype=torch.float32, device=device) + latent = torch.zeros(1, moe.moe_hidden_size, dtype=torch.bfloat16, device=device) + ids, weights, x_fp8, x_sf = k3_route_quant( + logits, self.bias, latent, float(moe.gate.routed_scaling_factor), early_trigger=True + ) + k3_moe(x_fp8, x_sf, ids, weights, offset, self.wide) + torch.cuda.synchronize(device) + + def _front(self, moe: nn.Module, x: torch.Tensor): + return k3_moe_front( + x.contiguous(), + self.front_weight, + self.bias, + float(moe.gate.routed_scaling_factor), + self.shared_cols, + *moe._situ_betas, + self.state.head, + ) + + def takes(self, hidden_states: torch.Tensor, step, partial_tail: bool) -> bool: + """Whether this decode path runs the layer on ``step`` (bf16 rows ``hidden_states``): any step of at most + `MAX_TOKENS` tokens; with ``partial_tail``, also a wide decode step.""" + rows = hidden_states.shape[0] + if hidden_states.dtype != torch.bfloat16: + return False + return 0 < rows <= MAX_TOKENS or (partial_tail and step.wide and rows <= WIDE_MAX_TOKENS) + + def forward( + self, + moe: nn.Module, + hidden_states: torch.Tensor, + gemvs, + partial_tail: bool, + push: bool = False, + ) -> Union[torch.Tensor, PendingTail]: + """The MoE output of ``hidden_states`` (``takes`` holds): with ``partial_tail``, this rank's unreduced share + (a ``PendingTail`` at most `MAX_TOKENS` tokens); else the reduced output. ``gemvs``: the decode GEMVs' state, + or None. ``push`` (the step's ``DecodeStep.latent_push``): at most `MAX_TOKENS` tokens, the latent all-reduce + is the routed experts' push form plus ``comm/k3_latent_reduce`` on the state's exchange (the same bits as the + routed experts' all-reduce); a state without the exchange ignores it.""" + if hidden_states.shape[0] > MAX_TOKENS: + return self._wide(moe, hidden_states, gemvs) + ids, weights, x_fp8, x_sf, shared_act = self._front(moe, hidden_states) + offset = moe.routed_experts.backend.slot_start + exchange = self.state.exchange + if push and exchange is not None: + k3_moe_push(x_fp8, x_sf, ids, weights, offset, self.small, exchange) + latent = k3_latent_reduce(hidden_states.shape[0], exchange) + else: + routed = k3_moe(x_fp8, x_sf, ids, weights, offset, self.small) + latent = moe.routed_experts.all_reduce(routed) + if partial_tail: + return PendingTail( + latent.contiguous(), + shared_act.contiguous(), + self.tail_weight, + self.lo, + float(moe.routed_expert_norm.variance_epsilon), + ) + # The replicated tail: the latent RMS on the fp32 accumulator of the folded latent up projection, plus the + # shared experts' reduced output, rounded to bf16 once. + shared = moe.shared_experts + shared_out = shared.down_proj(shared_act, layer_idx=shared.layer_idx) + up = _gemv(gemvs, "moe_up", latent, moe.routed_expert_up_proj.weight, out_fp32=True) + scale = torch.rsqrt( + latent.float().pow(2).mean(-1, keepdim=True) + moe.routed_expert_norm.variance_epsilon + ) + return (up * scale + shared_out.float()).bfloat16() + + def _wide(self, moe: nn.Module, hidden_states: torch.Tensor, gemvs) -> torch.Tensor: + """A wide decode step's MoE: this rank's unreduced share of the output, ``[M, hidden]`` bf16.""" + x = hidden_states.contiguous() + head = _gemv(gemvs, "moe_head", x, self.head_weight, out_fp32=True) + routed_in, router_logits = mnnvl_allgather_split(head, self.width, self.state.mnnvl) + shared = moe.shared_experts + + def _routed_partial(): + ids, weights, x_fp8, x_sf = k3_route_quant( + router_logits, self.bias, routed_in.contiguous(), + float(moe.gate.routed_scaling_factor), early_trigger=True, + ) # fmt: skip + return k3_moe( + x_fp8, x_sf, ids, weights, moe.routed_experts.backend.slot_start, self.wide + ) + + def _shared_activation(): + return shared._apply_activation( + _gemv(gemvs, "moe_shared_gate_up", x, shared.gate_up_proj.weight) + ) + + routed, shared_act = maybe_execute_in_parallel( + _routed_partial, + _shared_activation, + moe.moe_main_event, + moe.moe_shared_event, + moe.shared_expert_stream, + disable_on_compile=True, + ) + # The latent norm's weight is folded into the tail weight: normalize the whole latent row, keep this rank's + # columns. + normed = moe.routed_expert_norm(wide_all_reduce(moe.routed_experts.all_reduce, routed)) + parts = [normed[:, self.lo : self.lo + self.width]] + if self.tail_pad is not None: + parts.append(self.tail_pad[: x.shape[0]]) + parts.append(shared_act) + return _gemv(gemvs, "moe_tail", torch.cat(parts, dim=1), self.tail_weight) + + +def _gemv( + gemvs, site: str, x: torch.Tensor, weight: torch.Tensor, out_fp32: bool = False +) -> torch.Tensor: + """``x @ weight.T`` on ``site``'s decode GEMV where it takes the call (``decode_gemv.K3DecodeGemvs.project``), + else on the stock GEMM: cuBLAS through ``trtllm::dsv3_router_gemm_op`` for an fp32 output, ``F.linear`` else.""" + y = None if gemvs is None else gemvs.project(site, x, weight) + if y is not None: + return y + if out_fp32: + # The op reads the weight with a leading dimension of K: a strided weight would give wrong values. + if not weight.is_contiguous(): + raise ValueError( + f"the fp32 GEMM needs a contiguous weight, got strides {weight.stride()}" + ) + return torch.ops.trtllm.dsv3_router_gemm_op( + x.contiguous(), weight.t(), bias=None, out_dtype=torch.float32 + ) + return torch.nn.functional.linear(x, weight) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py new file mode 100644 index 000000000000..ec82423bd7cd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py @@ -0,0 +1,3403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 Kimi K3 with fused FP8 attention, shared by SM103 and SM107 targets.""" + +from __future__ import annotations + +import copy +import math +import os +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, NamedTuple, Optional, Tuple, Union + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( + k3_drafter_attn_qknorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify +from tensorrt_llm._torch.attention.backends import AttentionMetadata + +# Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. +from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 +from tensorrt_llm._torch.cute_dsl_kernels.k3_mla.decode_view import k3_mla_decode_view +from tensorrt_llm._torch.distributed import AllReduce +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.models.modeling_dspark import ( + GQADSparkForCausalLM, + draft_is_embedded_in_target, +) +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModel +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import maybe_bcg_kda_core_inplace +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.moe.fused_moe import ( + ConfigurableMoE, + SiTuActivation, + TRTLLMGenFusedMoE, + create_moe, +) +from tensorrt_llm._torch.moe.fused_moe.interface import MoESchedulerKind +from tensorrt_llm._torch.moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.config_utils import is_mla +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.utils import AuxStreamType +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +from . import attention as _fp8_attention +from . import decode_comm as _decode_comm +from . import decode_gemv as _decode_gemv +from . import decode_moe as _decode_moe +from . import spec_worker as _spec_worker +from . import weights as _weights + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +# The GPU architecture this target IS. Routing will not send another one here, but a direct instantiation could, +# and the certification is per arch. + +#: Every trtllm op this target reaches for, in its forward and in the weight load. Declared here, asserted in +#: tests/unittest/_torch/modeling_v2: the K3-specific ops of the generic path (attention residuals, KDA, the router +#: and fused-A GEMMs), then the decode kernels'. +REQUIRED_TRTLLM_OPS = ( + "attn_res_fwd", + "attn_res_rmsnorm_fwd", + "attn_res_add_rmsnorm_fwd", + "attn_res_add_rmsnorm_persistent_fwd", + "kda_prefill", + "kda_decode", + "kda_mtp_decode", + "dsv3_router_gemm_op", + "dsv3_fused_a_gemm_op", + "k3_kda_verify", + "k3_fp8_kda_decode", + "k3_fp8_mla_qkv", + "k3_fp8_mla_attn", + # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). + "k3_decode_gemv", + "k3_ctm_gemv_wide", + "k3_ctm_gemv_long", + "k3_situ_mul", + "k3_head_gemv", + "k3_embed_norm", + "allgather", + # The decode path's collectives (decode_comm.py) and MoE (decode_moe.py). + "mnnvl_allreduce_attn_res", + "k3_sandwich_oproj", + "k3_sandwich_tail", + "k3_moe_front", + "k3_moe", + "k3_latent_reduce", + "k3_route_quant", + "mnnvl_allgather_split", + # The DSpark drafter's decode path (K3DSparkDrafter): its GEMV sites, its block attention, and its all-reduces + # with the residual add and RMSNorm (the split context projection's with hidden_norm). + "k3_ctm_gemv", + "k3_ctm_gemv_swiglu", + "k3_drafter_attn_qknorm", + "k3_sandwich_plain", + "mnnvl_fusion_allreduce", +) + +#: The engine surface the first forward checks before this target relies on it: per object, the attributes read. +#: A renamed field upstream fails here, loudly, instead of reading as a default. +REQUIRED_ENGINE_FIELDS = { + "attn_metadata": ( + "num_contexts", + "num_generations", + "num_seqs", + "num_tokens", + "seq_lens", + "tokens_per_block", + "kv_cache_manager", + "mamba_metadata", + ), + "kv_cache_manager": ("enable_block_reuse", "mamba_layer_cache"), +} + +#: Stock code the generic path runs outside the catalog, declared so it is not consumed silently: every +#: tensorrt_llm import of this module that computes (test_modeling_v2_claims.py checks the list both ways). An entry +#: leaves the list when a catalog entry replaces it. +UNCERTIFIED_GENERIC_CALLS = ( + # The checkpoint load and the engine hooks this target inherits, and the causal LM around the text model. + "tensorrt_llm._torch.models.modeling_kimi_linear.KimiLinearForCausalLM", + "tensorrt_llm._torch.models.modeling_speculative.SpecDecOneEngineForCausalLM", + "tensorrt_llm._torch.models.modeling_utils.DecoderModel", + # The text model's stock modules. + "tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention", + "tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer.maybe_bcg_kda_core_inplace", + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops", + "tensorrt_llm._torch.modules.kimi_k3_mla.KimiK3MLAAttention", + "tensorrt_llm._torch.moe.fused_moe.create_moe", + "tensorrt_llm._torch.moe.fused_moe.ConfigurableMoE", + "tensorrt_llm._torch.moe.fused_moe.TRTLLMGenFusedMoE", + "tensorrt_llm._torch.moe.fused_moe.routing.DeepSeekV3MoeRoutingMethod", + "tensorrt_llm._torch.modules.gated_mlp.GatedMLP", + "tensorrt_llm._torch.modules.situ.SituAndMul", + "tensorrt_llm._torch.modules.rms_norm.RMSNorm", + "tensorrt_llm._torch.distributed.AllReduce", + "tensorrt_llm._torch.modules.multi_stream_utils.maybe_execute_in_parallel", + # The DSpark drafter: the stock GQA drafter K3DSparkDrafter extends (its block forward where the drafter entries + # do not take a block, its context projection without the TP group's collective state, its context k / v, its + # heads and its weight load), and the stock builder's checks for which drafter a checkpoint gets. + "tensorrt_llm._torch.models.modeling_dspark.GQADSparkForCausalLM", + "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", + "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", + "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", +) + +# The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware +# kernels take MAX_REQUESTS generation requests of at most MAX_TOKENS_PER_REQUEST tokens (1 + 7 drafts with DSpark). +DECODE_MAX_TOKENS = 8 +MAX_REQUESTS = 8 +MAX_TOKENS_PER_REQUEST = 8 + +# The MLA generation kernels for K3's 96 query heads exist only at a 64-token page (the built-in model's own +# get_model_defaults sets it for the same reason). +_TOKENS_PER_BLOCK = 64 + +_LANG_PREFIX = "language_model." + + +# ---------------------------------------------------------------------------------------------------------------------- +# The text model: Kimi K3's decoder (93 layers: KDA / MLA attention, attention residuals, the dense layer-0 MLP and +# the latent MoE), its generic path. Each step's classification goes to the attention modules (`K3DecodeKDA`, +# `K3DecodeMLA` below). +# ---------------------------------------------------------------------------------------------------------------------- + +# A/B escape hatch: restore nn.Linear for the K3 latent MoE projections +# instead of the min-latency fused GEMM op (read once at import). +_K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( + os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" +) + + +# Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table +# holds cos=1/sin=0), but the chunked-context path indexes the table by +# absolute position, so it must cover max_position_embeddings (~512MB per +# backend for the 1M-position checkpoint); a smaller table is read out of +# bounds. KIMI_K3_MLA_MAX_POSITIONS overrides the size for short-context +# deployments. +_KIMI_K3_MLA_MAX_POSITIONS_ENV = "KIMI_K3_MLA_MAX_POSITIONS" + + +class KimiK3MoeRoutingMethod(DeepSeekV3MoeRoutingMethod): + """DeepSeek-V3 routing, computed outside the TRTLLM-Gen MoE kernel. + + The kernel's own top-16 of 896 experts is the slow part of a generic MoE call of a few tokens. Outside it, the + scheduler routes with the backend's fused route + MXFP8 quantize (``trtllm::k3_route_quant``) up to 64 tokens and + with ``noaux_tc_op`` above, and the kernel takes the expert ids and weights. Both compute the kernel's arithmetic + and break ties the same way; a bf16 weight can differ by one ulp where it lies within fp32 rounding of a bf16 + rounding boundary. From about a thousand tokens on, ``noaux_tc_op`` takes longer than the kernel's routing. + """ + + @property + def requires_separated_routing(self) -> bool: + return True + + +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + @property + def routing_method(self) -> KimiK3MoeRoutingMethod: + """Return the DeepSeek-V3 router used by ``ConfigurableMoE``, computed outside the kernel.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return KimiK3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + +def _get_text_config(pretrained_config: "PretrainedConfig"): + """Return the Kimi text config, unwrapping a composite kimi_k3 config.""" + if getattr(pretrained_config, "model_type", None) == "kimi_k3" or ( + not hasattr(pretrained_config, "linear_attn_config") + and hasattr(pretrained_config, "text_config") + ): + return pretrained_config.text_config + return pretrained_config + + +def _is_kda_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["kda_layers"] + + +def _is_mla_layer(cfg, layer_idx: int) -> bool: + return (layer_idx + 1) in cfg.linear_attn_config["full_attn_layers"] + + +KIMI_K3_AUX_ATTN_RES_STREAM_ENV = "KIMI_K3_AUX_ATTN_RES_STREAM" + + +_AUX_ATTN_RES_STREAM_ENABLED = os.environ.get(KIMI_K3_AUX_ATTN_RES_STREAM_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" + + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV = "KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS" + + +KIMI_K3_ATTN_RES_TOPOLOGY_ENV = "KIMI_K3_ATTN_RES_TOPOLOGY" + + +_ATTN_RES_TOPOLOGIES = ("per_token", "persistent", "split") + + +# Turning the port on should not require knowing which of the three topologies +# is the right one: "1" selects the measured policy (``split``), so enabling the +# feature and choosing the policy are one action through one variable. +_ATTN_RES_TOPOLOGY_ON = "1" + + +def _read_attn_res_topology() -> str: + """Default ``per_token``: the persistent kernel is opt-in. + + ``1`` is the only accepted on-value and resolves to ``split``. The named + topologies stay available for measurement: ``persistent`` uses the + persistent kernel at every shape it implements, ``per_token`` at none. + """ + raw = os.environ.get(KIMI_K3_ATTN_RES_TOPOLOGY_ENV, "per_token") + if raw == _ATTN_RES_TOPOLOGY_ON: + return "split" + if raw not in _ATTN_RES_TOPOLOGIES: + # Loudly, for the same reason as the token ceiling below: a mistyped A/B + # arm that silently fell back to the default would measure one side + # twice and report no difference. + raise ValueError( + f"{KIMI_K3_ATTN_RES_TOPOLOGY_ENV} must be one of " + f"{_ATTN_RES_TOPOLOGIES} or {_ATTN_RES_TOPOLOGY_ON!r} " + f"(which means 'split'), got {raw!r}" + ) + return raw + + +def _read_fused_attn_res_max_tokens() -> int: + """Resolved after the topology, because its default follows it. + + With the persistent port off -- the default -- the ceiling is 1, the + pre-existing gate: the fused epilogue is taken at the single-token decode + shape and nowhere else. Enabling the port raises it to 32, the top of the + measured range, so that "off" keeps meaning "unchanged". + """ + default = "1" if _ATTN_RES_TOPOLOGY == "per_token" else "32" + raw = os.environ.get(KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV, default) + try: + value = int(raw) + except ValueError: + # Failing loudly matters more than usual here: a mistyped A/B arm that + # silently fell back to the default would measure the candidate twice + # and report no difference. + raise ValueError( + f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be a positive integer, got {raw!r}" + ) from None + if value < 1: + raise ValueError(f"{KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS_ENV} must be >= 1, got {value}") + return value + + +_ATTN_RES_TOPOLOGY = _read_attn_res_topology() + + +_FUSED_ATTN_RES_MAX_TOKENS = _read_fused_attn_res_max_tokens() + + +def _persistent_attn_res_applicable(M: int, H: int, N: int) -> bool: + """Shape gate for the persistent kernel: H == 7168 and 2 <= N <= 9. + + No token ceiling: the persistent grid is sized by the SM count, not by the + token count, so prefill is the case it exists for. + """ + del M # deliberately unused; see above + return H == 7168 and 2 <= N <= 9 + + +def _use_persistent_attn_res(M: int, H: int, N: int, max_fused_tokens: int) -> bool: + """Pick between the two fused kernels for this call site. + + ``persistent`` takes the persistent kernel at every shape it implements; + ``split`` takes it only above ``max_fused_tokens`` tokens, which stands in + for the prefill/decode boundary. Shapes the persistent kernel does not + implement fall through to the caller's existing gate and land on the + unfused path. + """ + if not _persistent_attn_res_applicable(M, H, N): + return False + if _ATTN_RES_TOPOLOGY == "persistent": + return True + return _ATTN_RES_TOPOLOGY == "split" and M > max_fused_tokens + + +def _apply_attn_res_fused( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> Optional[torch.Tensor]: + """Fused attn_res via the in-tree ``trtllm::attn_res_fwd`` op. + + Returns ``None`` when the call falls outside the fused kernel's + contract (dtype/shape/arch) so the caller can use the exact fp32 reference + instead. ``block_residual`` is kept in the kernel-native ``[K, M, H]`` + layout. Candidate order matches the reference: snapshots first, the + running prefix sum last. + """ + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + if K + 1 > 12 or M > 16384 or not (4096 <= H <= 8192 and H % 1024 == 0): + return None + try: + attn_res_op = torch.ops.trtllm.attn_res_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output, _rsigma, _probs, _logits = attn_res_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + ) + return output.reshape(M, H) + + +def _rms_norm_eps(norm: nn.Module) -> float: + if hasattr(norm, "eps"): + return float(norm.eps) + return float(norm.variance_epsilon) + + +def _note_attn_res_fusion(site: str, fused: bool, M: int, H: int, N: int) -> None: + """Report whether the fused path was actually reached, once per shape. + + ``_FUSED_ATTN_RES_ENABLED`` only says the feature is switched on, not that + the shape gate let the call through, and a rejected call looks exactly like + a disabled one in the logs. Emitted at debug level: it is once per distinct + shape, not once per process, so it is a diagnostic rather than a summary. + """ + logger.debug_once( + f"Kimi K3 attn-res fusion [{site}]: " + f"{'FUSED' if fused else 'fallback'} (M={M}, H={H}, N={N})", + key=f"kimi_k3_attn_res_fusion_{site}_{fused}_{M}_{H}_{N}", + ) + + +def _apply_attn_res_rmsnorm_fused( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[torch.Tensor]: + """Fuse attention-residual mixing with its immediately following norm (at most ``max_fused_tokens`` tokens, + default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not block_residual.is_cuda + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # The fused path is taken for M <= max_fused_tokens, H == 7168 and N <= 12, + # which is the measured window; larger token counts have not been measured + # and fall back to the unfused add + attn_res_fwd + RMSNorm path. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + _, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + None, + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("attn_res+norm/persistent", True, M, H, N) + return output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("attn_res+norm", False, M, H, N) + return None + try: + attn_res_rmsnorm_op = torch.ops.trtllm.attn_res_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + output = attn_res_rmsnorm_op( + layer_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("attn_res+norm", True, M, H, N) + return output.reshape(M, H) + + +def _apply_attn_res_add_rmsnorm_fused( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """Fuse ``prefix_sum + addend``, attention-residual, and trailing norm (``max_fused_tokens`` as in + ``_apply_attn_res_rmsnorm_fused``). + + The production residual add produces a BF16 tensor that remains live across + the following MLP. The kernel therefore returns that materialized, + BF16-rounded prefix sum alongside the normalized attention-residual output, + while avoiding a separate add launch and a re-read of the intermediate by + attention-residual selection. + """ + if max_fused_tokens is None: + max_fused_tokens = _FUSED_ATTN_RES_MAX_TOKENS + if ( + prefix_sum.dtype is not torch.bfloat16 + or addend.dtype is not torch.bfloat16 + or not prefix_sum.is_cuda + or not addend.is_cuda + or not block_residual.is_cuda + or prefix_sum.shape != addend.shape + ): + return None + M, H = prefix_sum.shape + K = int(block_residual.shape[0]) + N = K + 1 + # Same measured window as _apply_attn_res_rmsnorm_fused above. + if _use_persistent_attn_res(M, H, N, max_fused_tokens): + try: + persistent_op = torch.ops.trtllm.attn_res_add_rmsnorm_persistent_fwd + except (AttributeError, RuntimeError): + return None + updated_prefix_sum, output = persistent_op( + prefix_sum.reshape(M, 1, H).contiguous(), + addend.reshape(M, 1, H).contiguous(), + block_residual.reshape(K, M, 1, H).contiguous(), + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + ) + _note_attn_res_fusion("add+attn_res+norm/persistent", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + if M > max_fused_tokens or H != 7168 or N > 12: + _note_attn_res_fusion("add+attn_res+norm", False, M, H, N) + return None + try: + attn_res_add_rmsnorm_op = torch.ops.trtllm.attn_res_add_rmsnorm_fwd + except (AttributeError, RuntimeError): + return None + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + addend_kernel = addend.reshape(M, 1, H).contiguous() + block_kernel = block_residual.reshape(K, M, 1, H).contiguous() + updated_prefix_sum, output = attn_res_add_rmsnorm_op( + layer_kernel, + addend_kernel, + block_kernel, + proj.weight.reshape(-1).to(torch.bfloat16).contiguous(), + norm.weight.to(torch.bfloat16).contiguous(), + output_norm.weight.to(torch.bfloat16).contiguous(), + float(norm.eps), + _rms_norm_eps(output_norm), + early_trigger=True, + ) + _note_attn_res_fusion("add+attn_res+norm", True, M, H, N) + return updated_prefix_sum.reshape(M, H), output.reshape(M, H) + + +def _apply_attn_res( + prefix_sum: torch.Tensor, block_residual: torch.Tensor, proj: nn.Linear, norm: KimiK3RMSNorm +) -> torch.Tensor: + """Exact port of HF ``modeling_kimi._apply_attn_res`` (fp32 math). + + prefix_sum: ``[num_tokens, hidden_size]`` + block_residual: ``[num_snapshots, num_tokens, hidden_size]`` + + Unless ``KIMI_K3_FUSED_ATTN_RES=0``, inputs fitting the fused kernel's + contract dispatch directly to the in-tree ``trtllm::attn_res_fwd`` op. + Only the fallback boundary restores the HF ``[M, K, H]`` layout. + """ + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_fused(prefix_sum, block_residual, proj, norm) + if fused is not None: + return fused + block_residual_hf = block_residual.transpose(0, 1) + v = torch.cat((block_residual_hf, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.eps) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +def _apply_attn_res_and_rmsnorm( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> torch.Tensor: + """Apply attention-residual selection and the next RMSNorm. ``max_fused_tokens``: the largest token count the + fused kernel takes (default ``KIMI_K3_FUSED_ATTN_RES_MAX_TOKENS``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_rmsnorm_fused( + prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + return output_norm(_apply_attn_res(prefix_sum, block_residual, proj, norm)) + + +def _apply_attn_res_add_and_rmsnorm( + prefix_sum: torch.Tensor, + addend: torch.Tensor, + block_residual: torch.Tensor, + proj: nn.Linear, + norm: KimiK3RMSNorm, + output_norm: nn.Module, + max_fused_tokens: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Add an attention output to the running residual, then select and norm (``max_fused_tokens`` as in + ``_apply_attn_res_and_rmsnorm``).""" + if _FUSED_ATTN_RES_ENABLED: + fused = _apply_attn_res_add_rmsnorm_fused( + prefix_sum, addend, block_residual, proj, norm, output_norm, max_fused_tokens + ) + if fused is not None: + return fused + updated_prefix_sum = prefix_sum + addend + return updated_prefix_sum, _apply_attn_res_and_rmsnorm( + updated_prefix_sum, block_residual, proj, norm, output_norm, max_fused_tokens + ) + + +# A wide decode step's residual epilogues take the fused add + attn_res + RMSNorm kernels up to this many tokens, where +# they are faster than the add -> attn_res -> RMSNorm chain. +_WIDE_ATTN_RES_MAX_TOKENS = 32 + + +def _attn_res_max_tokens(step: DecodeStep) -> int: + """The most tokens of ``step`` the fused attn_res kernels take: one token tile (``DECODE_MAX_TOKENS``) on a step + ``decode_step`` classifies, ``_WIDE_ATTN_RES_MAX_TOKENS`` on a wide decode step.""" + return _WIDE_ATTN_RES_MAX_TOKENS if step.wide else DECODE_MAX_TOKENS + + +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + + +# The routed experts' module path prefixes. ``exclude_modules`` matches with +# wildcards and walks ancestor prefixes, so an empty prefix would widen what +# matches instead of just missing. +_K3_ROUTED_EXPERT_MODULE_PREFIXES = ("language_model.model.", "model.") + + +@dataclass(frozen=True) +class K3TargetSpec: + """The fixed hardware, expert format and parallelism of a concrete target.""" + + sm: tuple[int, int] + quant_algo: QuantAlgo + moe_tp: int = 4 + moe_ep: int = 4 + + +def _uses_nvfp4(config: ModelConfig) -> bool: + return config.quant_config is not None and config.quant_config.quant_algo in ( + QuantAlgo.MIXED_PRECISION, + QuantAlgo.NVFP4, + ) + + +# Routed-expert quantization used when the checkpoint declares nothing per +# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors +# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that +# checkpoint is what this default has always served. +_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 + + +def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_packed_mxfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), + w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), + w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), + w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), + w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), + w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), + ) + + +def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_streaming_nvfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + **{ + f"{w}_{kind}": get_tensor(f"{base}.{expert_idx}.{w}.{kind}") + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # Whether the layer is finalized after its experts load (the MXFP4 loaders write through). + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.NVFP4: _K3ExpertCkptSpec( + kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"), + loader=_load_nvfp4_expert, + loaded_slots_attr="_streamed_expert_slots", + needs_layer_finalize=True, + ), + QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( + kinds=("weight_packed", "weight_scale"), + loader=_load_packed_mxfp4_expert, + loaded_slots_attr="_packed_mxfp4_loaded_slots", + needs_layer_finalize=False, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + """Build the routed experts and the shared expert for one MoE layer. + + ``cfg`` is the raw ``PretrainedConfig`` rather than anything derived: + the SiTU soft-caps and the routed-expert geometry are Kimi K3 fields + that ``ModelConfig`` does not carry. + + ``aux_stream_dict`` is shared across every layer of the model, so the + streams reached through it are borrowed and must not be synchronized + or reassigned here. + """ + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + self.num_experts = cfg.num_experts + self.top_k = cfg.num_experts_per_token + self.moe_hidden_size = cfg.routed_expert_hidden_size + # ValueError (not assert): these guard unsupported checkpoint + # configurations and must stay active under ``python -O``. + if self.moe_hidden_size is None: + raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)") + if not getattr(cfg, "latent_moe_use_norm", False): + raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") + + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) + self._situ_betas = (situ_beta, situ_linear_beta) + dtype = torch.bfloat16 + + # Routing scores stay fp32; the gate GEMM runs bf16xbf16 with fp32 + # accumulate/output (checkpoint stores the gate weight in bf16; saves a + # per-layer input cast + fp32 splitK pair on the bs1 decode path). + # KIMI_K3_ROUTER_BF16=0 forces the upcast-to-fp32 GEMM. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = _router_bf16_env == "1" if _router_bf16_env is not None else True + self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) + + routed_moe_model_config = self._routed_moe_model_config(model_config) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) + routed_moe_kwargs = dict( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=cfg.moe_intermediate_size, + dtype=dtype, + # Kimi owns the latent reduction so it can order that collective + # after the shared expert's auxiliary-stream reduction. + reduce_results=False, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict, + # Let CommunicationFactory select the best available strategy. + communication_method=None, + activation=SiTuActivation( + gate_softcap=situ_beta, + linear_softcap=situ_linear_beta, + ), + # A request that silently degraded to CUTLASS would be benchmarked + # as if it were the backend that was asked for, and the decline is + # easy to trigger: MegaMoE has its own token / top-k limits and is + # EP-only, and CuteDSL declines on activation shape, SM version and + # the CuTe DSL dependency. As measured once: a CUTEDSL request + # was turned down on every one of the 92 MoE layers, on all 16 + # ranks, and still produced correct text and a zero exit -- the + # only trace was a warning line per layer. Fail in the resolver + # instead, which reports the rejection trail. + # + # CUTLASS is absent on purpose: it is the fallback target, so + # "degraded to CUTLASS" is not a thing that can happen to it. + allow_backend_degradation=routed_moe_model_config.moe_backend + not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", "CUTEDSL"), + ) + self._check_trtllm_situ_quant( + routed_moe_model_config.moe_backend, routed_quant_config.quant_algo + ) + + self.routed_experts = create_moe(**routed_moe_kwargs) + if not isinstance(self.routed_experts, ConfigurableMoE): + raise RuntimeError( + "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." + ) + if self.routed_experts.layer_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "dynamic EPLB or replicated expert slots." + ) + local_expert_ids = list(self.routed_experts.backend.initial_local_expert_ids) + if local_expert_ids != list( + range(local_expert_ids[0], local_expert_ids[0] + len(local_expert_ids)) + ): + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming currently requires a " + "contiguous static expert partition." + ) + self.local_expert_ids = tuple(local_expert_ids) + self.experts_per_rank = len(local_expert_ids) + self.expert_lo = local_expert_ids[0] + self.expert_hi = self.expert_lo + self.experts_per_rank + + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. + use_shared_tp = model_config.mapping.tp_size > 1 + self._reduce_routed_output = ( + use_shared_tp + and self.routed_experts.backend.scheduler_kind != MoESchedulerKind.FUSED_COMM + ) + if self._reduce_routed_output and self.routed_experts.all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self.shared_experts = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=shared_model_config, + reduce_output=use_shared_tp, + layer_idx=layer_idx, + is_shared_expert=True, + ) + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.moe_main_event = torch.cuda.Event() + self.moe_shared_event = torch.cuda.Event() + self.routed_expert_down_proj = nn.Linear( + cfg.hidden_size, self.moe_hidden_size, bias=False, dtype=dtype + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype + ) + # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer + # fallback is the same fp32-variance eager math as KimiK3RMSNorm). + self.routed_expert_norm = RMSNorm( + hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + # The decode path (decode_moe.py) and the decode GEMVs' state, set by the target's post_load_weights. + self.decode_moe: Optional[_decode_moe.K3DecodeMoeLayer] = None + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, projection.weight.t(), None, None + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """The expert split, validated against the concrete deployment target.""" + return mapping.moe_tp_size, mapping.moe_ep_size + + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Resolve per-layer NVFP4 metadata or the compressed-tensors MXFP4 default. + + An exclusion outranks the default and is checked against actual module names. + """ + if _uses_nvfp4(model_config): + global_quant = model_config.quant_config + if any( + global_quant.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in ("language_model.model.", "model.", "") + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + raise ValueError( + f"NVFP4 experts at layer {layer_idx} are excluded from quantization" + ) + declarations = model_config.quant_config_dict or {} + matches = [ + declarations[name] + for prefix in ("language_model.model.", "model.", "") + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + if (name := f"{prefix}layers.{layer_idx}.{suffix}") in declarations + ] + if matches: + if any( + quant.quant_algo != matches[0].quant_algo + or quant.group_size != matches[0].group_size + for quant in matches[1:] + ): + raise ValueError( + f"Conflicting expert quantization aliases at layer {layer_idx}" + ) + return matches[0] + if global_quant.quant_algo == QuantAlgo.NVFP4: + return global_quant + raise ValueError(f"Missing NVFP4 expert declaration at layer {layer_idx}") + quant_config = model_config.quant_config + if quant_config is not None and any( + quant_config.is_module_excluded_from_quantization( + f"{prefix}layers.{layer_idx}.{suffix}" + ) + for prefix in _K3_ROUTED_EXPERT_MODULE_PREFIXES + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES + ): + logger.debug( + "Kimi K3 layer %d routed experts: excluded from quantization, " + "keeping them unquantized", + layer_idx, + ) + return QuantConfig(kv_cache_quant_algo=quant_config.kv_cache_quant_algo) + + logger.debug( + "Kimi K3 layer %d routed experts: no per-layer quant config in the " + "checkpoint, defaulting to %s", + layer_idx, + _K3_DEFAULT_ROUTED_QUANT_ALGO, + ) + return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) + + @staticmethod + def _check_trtllm_situ_quant(moe_backend: str, quant_algo: Optional[QuantAlgo]) -> None: + """Reject formats without a fused SiTU cubin in the resolved TRTLLM backend. + + Read the supported formats from the backend so model construction and + kernel selection agree, including when AUTO resolves to TRTLLM. + """ + situ_supported = TRTLLMGenFusedMoE.situ_supported_quant_algos() + if moe_backend != "TRTLLM" or quant_algo in situ_supported: + return + supported = ", ".join(sorted(algo.name for algo in situ_supported)) + raise ValueError( + f"Kimi K3 routed experts are quantized as {quant_algo}, which the " + "TRTLLM (trtllm-gen) MoE backend cannot serve: fused SiTu cubins " + f"exist only for {supported}. Set moe_config.backend to CUTLASS " + "or MEGAMOE_CUTEDSL." + ) + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config, with the split of ``_select_moe_tp_ep``.""" + # Every backend here declares ``ActivationType.SiTu`` in its + # ``activation_support``; the list is not a preference order. CUTEDSL + # joined once its act-fusion kernel grew the SiTU epilogue. + supported_backends = { + "CUTLASS", + "TRTLLM", + "CUTEDSL", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "CUTEDSL, MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " + f"got {model_config.moe_backend!r}." + ) + if model_config.moe_load_balancer is not None: + raise NotImplementedError( + "Kimi K3 packed-checkpoint streaming does not yet support " + "EPLB or replicated expert slots." + ) + mapping = model_config.mapping + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + logger.info_once( + f"Kimi K3 routed MoE parallelism: moe_tp={moe_tp}, " + f"moe_ep={moe_ep} (tp_size={mapping.tp_size})", + key="kimi_k3_moe_tp_ep_split", + ) + + mapping_dict = mapping.to_dict() + mapping_dict["moe_cluster_size"] = 1 + mapping_dict["moe_tp_size"] = moe_tp + mapping_dict["moe_ep_size"] = moe_ep + routed_mapping = Mapping.from_dict(mapping_dict) + + routed_model_config = copy.copy(model_config) + routed_model_config._frozen = False + routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) + routed_model_config.mapping = routed_mapping + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend in { + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + }: + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) + if configured_moe_max_num_tokens < default_moe_max_num_tokens: + logger.info_once( + "Kimi K3 MegaMoE raises moe_max_num_tokens from " + f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} " + "because the global DP SymmBuffer requires capacity for " + "max_num_tokens * dp_size.", + key=( + "kimi_k3_megamoe_capacity_override_" + f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}" + ), + ) + routed_model_config.moe_max_num_tokens = max( + configured_moe_max_num_tokens, + default_moe_max_num_tokens, + ) + routed_model_config._frozen = True + return routed_model_config + + def tail_rp_eligible(self, hidden_states: torch.Tensor, step: DecodeStep) -> bool: + """Whether this layer's forward on ``step`` can hand its output on as this rank's unreduced row-parallel + share (``partial_tail``): the decode path takes the step with that tail.""" + return self.decode_moe is not None and self.decode_moe.takes(hidden_states, step, True) + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens=None, + partial_tail: bool = False, + *, + step: DecodeStep, + ) -> Union[torch.Tensor, _decode_comm.PendingTail]: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16. ``step``: the step's classification + (``decode_step``); the decode path (``decode_moe.py``) runs the steps it takes. ``partial_tail`` (only where + ``tail_rp_eligible`` holds): return this rank's unreduced share of the output instead, a ``PendingTail`` at + most 8 tokens, a tensor on a wide decode step.""" + decode = self.decode_moe + if decode is not None and decode.takes(hidden_states, step, partial_tail): + return decode.forward( + self, hidden_states, self.decode_gemvs, partial_tail, push=step.latent_push + ) + if partial_tail: + raise RuntimeError("the row-parallel MoE tail needs the decode path to take the step") + return self.forward_generic(hidden_states, all_rank_num_tokens) + + def forward_generic( + self, hidden_states: torch.Tensor, all_rank_num_tokens=None + ) -> torch.Tensor: + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce if self._reduce_routed_output else None + + def _routed_output(): + # Latent down/up projections via the min-latency fused GEMM op: + # at <=16 tokens (decode graphs) it runs a single pipelined + # bf16 kernel per projection instead of cuBLAS's split-K GEMV + + # splitKreduce pair (~17+3.6us -> ~8us for 7168->3584 at M=1); + # for larger token counts the op falls back to cuBLAS internally. + # TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ=1 restores nn.Linear + # (A/B escape hatch). When the FP8 weight-read conversion has + # replaced the projection module, call it directly: its weight is + # an e4m3 buffer the bf16 dsv3 op must not read, and its forward + # is already a single fused GEMM (fp8_swap_ab_gemm). + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + if self._reduce_routed_output: + return y + # Communication-backed paths return a complete routed result. + y = self.routed_expert_norm(y) + return self._routed_projection(y, self.routed_expert_up_proj) + + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. The shared GatedMLP includes its output all-reduce on + # the auxiliary stream. The join below must precede the routed + # all-reduce: concurrent collectives on different streams can corrupt + # SYMM_MEM all-reduce state. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.shared_expert_stream, + disable_on_compile=True, + ) + if self._reduce_routed_output: + routed_latent = moe_all_reduce(routed_out) + routed_latent = self.routed_expert_norm(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + return routed_out + shared_out + + +def resolve_attention_quant_config( + config: ModelConfig | None, layer_idx: int, projection: str +) -> QuantConfig: + """Resolve a checkpoint projection, including mixed-precision exclusions.""" + if config is None: + return QuantConfig() + global_config = config.quant_config or QuantConfig() + names = [ + f"{prefix}layers.{layer_idx}.self_attn.{projection}" + for prefix in ("language_model.model.", "model.", "") + ] + if any(global_config.is_module_excluded_from_quantization(name) for name in names): + return QuantConfig(kv_cache_quant_algo=global_config.kv_cache_quant_algo) + declarations = config.quant_config_dict or {} + matches = [declarations[name] for name in names if name in declarations] + if matches: + selected = matches[0] + if any( + (match.quant_algo, match.group_size) != (selected.quant_algo, selected.group_size) + for match in matches[1:] + ): + raise ValueError(f"Conflicting Kimi K3 quantization aliases for {names[0]}") + elif global_config.quant_algo == QuantAlgo.MIXED_PRECISION: + selected = QuantConfig() + else: + selected = global_config + if selected.quant_algo not in (None, QuantAlgo.FP8_BLOCK_SCALES): + raise ValueError( + f"Kimi K3 attention projection {names[0]} has unsupported checkpoint " + f"quantization {selected.quant_algo}" + ) + if selected.quant_algo == QuantAlgo.FP8_BLOCK_SCALES and selected.group_size not in (None, 128): + raise ValueError(f"Kimi K3 attention requires 128x128 FP8 blocks for {names[0]}") + result = copy.copy(selected) + result.kv_cache_quant_algo = global_config.kv_cache_quant_algo + return result + + +class KimiMLARuntime(nn.Module): + """Wraps K3 MLA and applies its external TP output reduction.""" + + def __init__( + self, + cfg: "PretrainedConfig", + layer_idx: int, + model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ) -> None: + super().__init__() + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the + # final output reduction in this wrapper so the output gate remains + # between attention and the row-parallel o_proj. + mapping = model_config.mapping + reduce_output = mapping.tp_size > 1 + self._o_allreduce = ( + AllReduce( + mapping=mapping, + strategy=model_config.allreduce_strategy, + dtype=torch.bfloat16, + ) + if reduce_output + else None + ) + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in ( + "q_a_proj", + "kv_a_proj_with_mqa", + "q_b_proj", + "kv_b_proj", + "g_proj", + "o_proj", + ) + } + attention_config.use_cute_dsl_blockscaling_bmm = True + attention_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + attention_config._frozen = model_config._frozen + mla_cls = K3DecodeMLA + self.mixer = mla_cls( + hidden_size=cfg.hidden_size, + num_heads=cfg.num_attention_heads, + q_lora_rank=cfg.q_lora_rank, + kv_lora_rank=cfg.kv_lora_rank, + qk_nope_head_dim=cfg.qk_nope_head_dim, + qk_rope_head_dim=cfg.qk_rope_head_dim, + v_head_dim=cfg.v_head_dim, + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + use_output_gate=cfg.mla_use_output_gate, + max_position_embeddings=max_positions, + model_config=attention_config, + aux_stream_dict=aux_stream_dict, + ) + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the mixer runs ``step`` on the decode kernels (``K3DecodeMLA.will_run_decode_branch``).""" + return self.mixer.will_run_decode_branch(attn_metadata, step) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce); ``project_output=False`` the + mixer's gated attention output before ``o_proj``, only where ``will_run_decode_branch`` holds.""" + # MLA.forward takes position_ids first; K3 is NoPE, so pass None. + out = self.mixer( + None, hidden_states, attn_metadata, step=step, project_output=project_output + ) + return self._reduce(out) if project_output and reduce_output else out + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + out = self.mixer.forward_generic(None, hidden_states, attn_metadata) + return self._reduce(out) if reduce_output else out + + def _reduce(self, out: torch.Tensor) -> torch.Tensor: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + return self._o_allreduce(out) if self._o_allreduce is not None else out + + +class KimiLinearDecoderLayer(nn.Module): + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + dtype = torch.bfloat16 + + self.is_kda = _is_kda_layer(cfg, layer_idx) + is_mla = _is_mla_layer(cfg, layer_idx) + if self.is_kda == is_mla: + raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") + + if self.is_kda: + projection_names = ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj") + attention_config = copy.copy(model_config) + attention_config._frozen = False + attention_config.quant_config_dict = { + name: resolve_attention_quant_config(model_config, layer_idx, name) + for name in projection_names + } + attention_config._frozen = model_config._frozen + kda_cls = K3DecodeKDA + self.linear_attn = kda_cls( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + model_config=attention_config, + ) + else: + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + model_config=model_config, + aux_stream_dict=aux_stream_dict, + ) + + self.is_moe = ( + cfg.num_experts is not None + and layer_idx >= cfg.first_k_dense_replace + and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 + ) + if self.is_moe: + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + # Over MNNVL (one NVLink domain across the nodes, where a cross-node all-reduce costs what a node's + # does) the MLP stays split over the whole TP group, the per-rank shapes its decode GEMVs take + # (decode_gemv.SITES); otherwise it stays within one node. + spans_nodes = self._mnnvl_allreduce() is not None + if self.mlp_tp_size > model_config.mapping.gpus_per_node and not spans_nodes: + self.mlp_tp_size = math.gcd(self.mlp_tp_size, model_config.mapping.gpus_per_node) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. + self.mlp = GatedMLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=dtype, + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, + ) + self._situ = (situ_beta, situ_linear_beta) + # The decode GEMVs' state (decode_gemv.py), set by the target's cache_derived_state. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics + # are drop-in for KimiK3RMSNorm (fp32 variance, weight applied + # after downcast, use_gemma=False). + self.input_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + + # Attention residual scheme (always on for K3). The res norms stay + # KimiK3RMSNorm: they are consumed field-wise (.weight/.eps) by + # _apply_attn_res and the fused attn_res op, never called as + # modules. + self.attn_res_block_size = cfg.attn_res_block_size + assert self.attn_res_block_size is not None, ( + "Kimi K3 runtime expects attn_res_block_size to be set" + ) + self.self_attention_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.mlp_res_norm = KimiK3RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + # The decode path's collectives (decode_comm.py), and whether their sandwich takes this layer's o_proj; set + # by the target's post_load_weights. + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + self.sandwich_oproj = False + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + step: DecodeStep, + capture: Optional[Tuple[Any, int]] = None, + prenormed: bool = False, + pending_moe_partial: Optional[Union[torch.Tensor, _decode_comm.PendingTail]] = None, + defer_moe_tail: bool = False, + ) -> Union[Tuple[torch.Tensor, int], Tuple[torch.Tensor, int, Any]]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + ``block_residual`` is a preallocated snapshot bank in kernel-native + ``[K_max, M, H]`` layout. Returns the running prefix sum and the + number of valid bank rows. + + ``capture`` is ``(spec_metadata, layer_id)`` and taps the DSpark aux + stream for the layer BEFORE this one: the aggregated stream for layer j + is by definition what its next consumer sees, so the mixture computed + below already is it. Reading it here beats recomputing it, and is only + possible because K3 asserts pp_size == 1 -- layer j+1 is always local. + PP support would need a recompute at the rank boundary. + + ``step`` is the step's classification (``decode_step``), handed to the + attention module. + + ``prenormed`` (layer 0 on a decode step): ``hidden_states`` already is + this layer's input norm, and the layer's input, the step's embedding, + already is in ``block_residual[0]`` (``K3DecodeGemvs.embed_norm``). + + The post-attention step of at most ``AR_ATTN_RES_MAX_TOKENS`` tokens + runs the attention's all-reduce and the residual update in one + collective once the target has built its ``decode_comm``: with o_proj + (``K3DecodeComm.sandwich_oproj``) where the sandwich takes the layer, + the step and the attention's decode branch, else on the attention's + unreduced o_proj output (``K3DecodeComm.allreduce_attn_res``). + + ``pending_moe_partial`` (only where ``accepts_moe_partial`` held): the + previous layer's MoE output, unreduced (``defer_moe_tail``); then + ``hidden_states`` is the prefix sum without it, and this layer's + pre-attention step reduces and adds it: a ``PendingTail`` in one + ``K3DecodeComm.sandwich_tail`` call, a wide decode step's tensor by its + all-reduce and the fused add + attn_res + RMSNorm. On a tapped layer + the ``sandwich_tail`` call also stores the tap into the capture slot + where the speculative metadata exposes it (``capture_view``); other + tapped layers take the tap after the step. + + ``defer_moe_tail``: return ``(prefix_sum, num_snapshots, partial)`` + instead, ``partial`` this layer's MoE output unreduced, for the next + consumer to reduce. + """ + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + attn_res_max_tokens = _attn_res_max_tokens(step) + tail = ( + pending_moe_partial + if isinstance(pending_moe_partial, _decode_comm.PendingTail) + else None + ) + # A snapshot layer whose pre-attention step is the sandwich tail has the kernel store the running prefix sum + # straight into the bank row it snapshots. + snapshot_row = None + if tail is not None and self.layer_idx % self.attn_res_block_size == 0: + snapshot_row = block_residual[num_snapshots] + # A tapped layer whose pre-attention step is the sandwich tail has the kernel store the tap straight into the + # layer's capture slot, where the speculative metadata exposes that slot as a view (``capture_view``). + tap_view = None + if tail is not None and capture is not None: + tap_view = _spec_worker.capture_view(capture[0], capture[1], prefix_sum.shape[0]) + + if prenormed: + assert num_snapshots == 0 and self.layer_idx % self.attn_res_block_size == 0 + assert pending_moe_partial is None and capture is None + elif tail is not None: + hidden_states, prefix_sum = self.decode_comm.sandwich_tail( + tail, + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + updated_out=snapshot_row, + tap=tap_view, + tap_updated=not _AUX_ATTN_RES_STREAM_ENABLED, + ) + elif pending_moe_partial is not None: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + _decode_comm.wide_all_reduce(self._o_allreduce(), pending_moe_partial), + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + elif capture is not None: + # The mixture tap needs the PRE-norm value, which the fused + # attn-res + RMSNorm kernel does not expose. Keep the two steps + # split on captured layers only and fuse everywhere else. + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + # A property of the DRAFTER checkpoint, not a knob: a mismatch only lowers + # acceptance, silently. hidden_states is the pre-norm attn_res mixture; + # prefix_only wants the running prefix, already in hand as prefix_sum. + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + attn_res_max_tokens, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if capture is not None and pending_moe_partial is not None and tap_view is None: + # The tapped layer handed its MoE output on and no kernel wrote the tap: the step above reduced the output + # into prefix_sum. Tap that value's pre-norm attn_res mixture, what the split path captures. + tapped = ( + _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + if _AUX_ATTN_RES_STREAM_ENABLED + else prefix_sum + ) + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + + if self.layer_idx % self.attn_res_block_size == 0: + if not prenormed and snapshot_row is None: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, step): + if ( + self.sandwich_oproj + and step.num_tokens <= _decode_comm.SANDWICH_MAX_TOKENS + and attention.will_run_decode_branch(attn_metadata, step) + ): + core = attention(hidden_states, attn_metadata, step=step, project_output=False) + hidden_states, prefix_sum = comm.sandwich_oproj( + core, + self._o_proj().weight, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + if comm is not None and step.wide: + partial = attention(hidden_states, attn_metadata, step=step, reduce_output=False) + attention_output = _decode_comm.wide_all_reduce(self._o_allreduce(), partial) + else: + attention_output = attention(hidden_states, attn_metadata, step=step) + if prefix_sum is None: + prefix_sum = attention_output + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention_output, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + attn_res_max_tokens, + ) + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + if self.is_moe and defer_moe_tail: + partial = self.block_sparse_moe( + hidden_states, all_rank_num_tokens, partial_tail=True, step=step + ) + return prefix_sum, num_snapshots, partial + if self.is_moe: + hidden_states = self.block_sparse_moe(hidden_states, all_rank_num_tokens, step=step) + else: + hidden_states = self._dense_mlp(hidden_states, step) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, num_snapshots + + def forward_generic( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + attn_metadata: AttentionMetadata, + capture: Optional[Tuple[Any, int]] = None, + ) -> Tuple[torch.Tensor, int]: + prefix_sum = hidden_states + valid_block_residual = block_residual[:num_snapshots] + if capture is not None: + if num_snapshots > 0: + hidden_states = _apply_attn_res( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + tapped = hidden_states if _AUX_ATTN_RES_STREAM_ENABLED else prefix_sum + capture[0].maybe_capture_hidden_states(capture[1], tapped, None) + hidden_states = self.input_layernorm(hidden_states) + elif num_snapshots > 0: + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + self.input_layernorm, + ) + else: + hidden_states = self.input_layernorm(hidden_states) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual[num_snapshots].copy_(prefix_sum) + num_snapshots += 1 + valid_block_residual = block_residual[:num_snapshots] + prefix_sum = None + attention = self.linear_attn if self.is_kda else self.self_attn + comm = self.decode_comm + if comm is not None and comm.takes_post_attention(hidden_states, None): + partial = attention.forward_generic(hidden_states, attn_metadata, reduce_output=False) + hidden_states, prefix_sum = comm.allreduce_attn_res( + partial, + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + elif prefix_sum is None: + prefix_sum = attention.forward_generic(hidden_states, attn_metadata) + hidden_states = _apply_attn_res_and_rmsnorm( + prefix_sum, + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + else: + prefix_sum, hidden_states = _apply_attn_res_add_and_rmsnorm( + prefix_sum, + attention.forward_generic(hidden_states, attn_metadata), + valid_block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + self.post_attention_layernorm, + ) + if self.is_moe: + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + hidden_states = self.block_sparse_moe.forward_generic( + hidden_states, all_rank_num_tokens + ) + else: + hidden_states = self.mlp(hidden_states) + return prefix_sum + hidden_states, num_snapshots + + def _dense_mlp(self, hidden_states: torch.Tensor, step: DecodeStep) -> torch.Tensor: + """The dense MLP: on a step of at most DECODE_MAX_TOKENS tokens, its GEMVs and activation on the decode + kernels (``K3DecodeGemvs.dense_mlp``), then the down projection's all-reduce; else the module.""" + gemvs = self.decode_gemvs + if gemvs is not None and step.small: + out = gemvs.dense_mlp( + hidden_states, + self.mlp.gate_up_proj.weight, + self.mlp.down_proj.weight, + *self._situ, + ) + if out is not None: + return self.mlp.down_proj.all_reduce(out) if self.mlp_tp_size > 1 else out + return self.mlp(hidden_states) + + def _mnnvl_allreduce(self): + """The MNNVL all-reduce of this layer's attention output, or None.""" + attention = self.linear_attn if self.is_kda else self.self_attn + return getattr(getattr(attention, "_o_allreduce", None), "mnnvl_allreduce", None) + + def _o_proj(self) -> nn.Module: + """This layer's attention output projection (row parallel).""" + return self.linear_attn.o_proj if self.is_kda else self.self_attn.mixer.o_proj + + def _o_allreduce(self) -> AllReduce: + """The all-reduce module of this layer's attention output; it also reduces a wide step's MoE partial.""" + return (self.linear_attn if self.is_kda else self.self_attn)._o_allreduce + + def accepts_moe_partial(self, num_snapshots: int) -> bool: + """Whether this layer's pre-attention step can reduce the previous layer's unreduced MoE output: the target + built its collective state, and the snapshot bank is not empty.""" + return num_snapshots > 0 and self.decode_comm is not None + + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + + +class KimiLinearModel(DecoderModel): + def __init__(self, model_config: ModelConfig): + super().__init__(model_config) + cfg = _get_text_config(model_config.pretrained_config) + self._text_cfg = cfg + dtype = torch.bfloat16 + + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream() for _ in range(4)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + } + + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) + self.layers = nn.ModuleList( + [ + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) + for layer_idx in range(cfg.num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size=cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype) + + # KimiK3RMSNorm (not RMSNorm): consumed field-wise (.weight/.eps) + # by _apply_attn_res and the fused attn_res op. + self.output_attn_res_norm = KimiK3RMSNorm( + cfg.hidden_size, eps=cfg.rms_norm_eps, dtype=dtype + ) + self.output_attn_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype) + self.num_attn_res_snapshots = ( + cfg.num_hidden_layers + cfg.attn_res_block_size - 1 + ) // cfg.attn_res_block_size + # The decode path's GEMVs and embedding (decode_gemv.py), built by the target's cache_derived_state once the + # weights are final. None: every step embeds on the generic path. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + # Which convention the drafter tap is on is not recoverable from the + # served output -- a mismatch only lowers acceptance -- so state it once + # at construction rather than leaving it to be inferred from an AL. + logger.info_once( + "Kimi K3 aux hidden capture: mode=" + f"{'attn_res_stream' if _AUX_ATTN_RES_STREAM_ENABLED else 'prefix_only'} " + f"({KIMI_K3_AUX_ATTN_RES_STREAM_ENV}={int(_AUX_ATTN_RES_STREAM_ENABLED)})", + key="kimi_k3_aux_capture_mode", + ) + + @property + def kda_draft_records(self) -> bool: + """Whether the hybrid cache manager keeps the last verify's draft records (``kda_state_tok``), from which + ``ssm/k3_kda_verify`` and ``ssm/k3_kda_attn`` replay the accepted drafts. The engine reads it once the weights + are loaded, to build the manager: DFlash / DSpark drafts of an even verify width up to 8, every KDA layer + taking the K3 kernels. Otherwise the KDA verify replays the accepted drafts from the replay caches (the + built-in verify).""" + spec_config = getattr(self.model_config, "spec_config", None) + return bool( + spec_config is not None + and (spec_config.spec_dec_mode.is_dflash() or spec_config.spec_dec_mode.is_dspark()) + and spec_config.tokens_per_gen_step in (2, 4, 6, 8) + and all(layer.linear_attn.takes_k3_kernels for layer in self.layers if layer.is_kda) + ) + + def _latent_push(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether the MoE layers push on ``step`` (``latent_push``): read only while a CUDA graph is being captured, + from what the graph's key fixes (the step's shape, the attention metadata the decode kernels read) and from + load-time state, so every rank and every replay of the graph decides alike.""" + if not torch.cuda.is_current_stream_capturing(): + return False + kda = [layer.linear_attn for layer in self.layers if layer.is_kda] + mla = [layer.self_attn for layer in self.layers if not layer.is_kda] + return latent_push( + step, + capturing=True, + breakable=is_in_breakable_cuda_graph(), + kda_draft_records=self.kda_draft_records, + kda_decode_kernels=all(m.takes_k3_kernels and m.fp8_decode_ready for m in kda), + mla_decode_branch=all(m.will_run_decode_branch(attn_metadata, step) for m in mla[:1]), + ) + + def _defer_moe_tail( + self, + i: int, + hidden_states: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + step: DecodeStep, + ) -> bool: + """Whether layer ``i`` hands its MoE output on unreduced (the row-parallel tail, ``decode_moe.py``): its MoE + takes the step with that tail, and its consumer, the next layer's pre-attention step or the final norm after + the last layer, accepts it. A layer DSpark taps defers too: the next layer's step reduces the output, then taps + its pre-norm mixture. An unknown capture set (every layer tapped) and a tapped last layer keep the replicated + tail.""" + layer = self.layers[i] + if not (layer.is_moe and layer.block_sparse_moe.tail_rp_eligible(hidden_states, step)): + return False + if spec_metadata is not None and ( + capture_set is None or (layer.layer_idx in capture_set and i == len(self.layers) - 1) + ): + return False + return self.layers[min(i + 1, len(self.layers) - 1)].accepts_moe_partial(num_snapshots) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata=None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + num_tokens = (input_ids if inputs_embeds is None else inputs_embeds).shape[0] + step = decode_step(attn_metadata, num_tokens) + if step is None: + return self._generic_forward(attn_metadata, input_ids, inputs_embeds, spec_metadata) + if self._latent_push(attn_metadata, step): + step = replace(step, latent_push=True) + # A decode step embeds and norms for layer 0 in one launch, the embedding written as layer 0's first snapshot. + prenormed = None + if ( + inputs_embeds is None + and self.decode_gemvs is not None + and len(self.layers) > 0 + and self.num_attn_res_snapshots > 0 + ): + table = self.embed_tokens.weight + block_residual = table.new_empty( + self.num_attn_res_snapshots, num_tokens, table.shape[1] + ) + prenormed = self.decode_gemvs.embed_norm( + input_ids, table, self.layers[0].input_layernorm, block_residual + ) + if prenormed is not None: + hidden_states = prenormed + else: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + # A MoE layer's output handed on unreduced, which the next layer's pre-attention step (or the final norm) + # reduces: a decode_comm.PendingTail, or a wide decode step's tensor. + pending_moe_partial = None + for i, layer in enumerate(self.layers): + defer_moe_tail = self._defer_moe_tail( + i, hidden_states, num_snapshots, spec_metadata, capture_set, step + ) + outputs = layer( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + step=step, + capture=self._capture(spec_metadata, capture_set, i), + prenormed=i == 0 and prenormed is not None, + pending_moe_partial=pending_moe_partial, + defer_moe_tail=defer_moe_tail, + ) + if defer_moe_tail: + hidden_states, num_snapshots, pending_moe_partial = outputs + else: + (hidden_states, num_snapshots), pending_moe_partial = outputs, None + + if isinstance(pending_moe_partial, _decode_comm.PendingTail): + normed, _ = self.layers[-1].decode_comm.sandwich_tail( + pending_moe_partial, + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + ) + return normed + if pending_moe_partial is not None: + _, normed = _apply_attn_res_add_and_rmsnorm( + hidden_states, + _decode_comm.wide_all_reduce(self.layers[-1]._o_allreduce(), pending_moe_partial), + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + _attn_res_max_tokens(step), + ) + return normed + return self._final_norm( + hidden_states, + block_residual, + num_snapshots, + spec_metadata, + capture_set, + _attn_res_max_tokens(step), + ) + + def _generic_forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor], + inputs_embeds: Optional[torch.FloatTensor], + spec_metadata, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds + block_residual = hidden_states.new_empty( + self.num_attn_res_snapshots, + hidden_states.shape[0], + hidden_states.shape[1], + ) + num_snapshots = 0 + capture_set = getattr(spec_metadata, "_capture_layer_set", None) + for i, layer in enumerate(self.layers): + hidden_states, num_snapshots = layer.forward_generic( + hidden_states, + block_residual, + num_snapshots, + attn_metadata, + capture=self._capture(spec_metadata, capture_set, i), + ) + return self._final_norm( + hidden_states, block_residual, num_snapshots, spec_metadata, capture_set + ) + + def _capture(self, spec_metadata, capture_set, i: int) -> Optional[Tuple[Any, int]]: + # DFlash/DSpark hidden-state capture. The drafter is distilled on + # the aggregated stream value -- the pre-norm softmax mixture its + # next consumer sees -- not on the raw prefix sum a layer returns, + # which is SGLang's fallback for models without the + # attention-residual scheme. Capturing the prefix sum costs 4.5pt + # of draft acceptance on K3 + RadixArk DSpark (AR 66.9% -> 71.4%). + # The tap fires inside layer i+1, which computes that tensor + # anyway; see its forward docstring. Ground truth: SGLang + # kimi_k3.py:2697 _dspark_capture_stream, attn_residual.py:313 + # aggregate_stream_torch. + if ( + spec_metadata is not None + and i > 0 + and (capture_set is None or self.layers[i - 1].layer_idx in capture_set) + ): + return (spec_metadata, self.layers[i - 1].layer_idx) + return None + + def _final_norm( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + num_snapshots: int, + spec_metadata, + capture_set, + max_fused_tokens: Optional[int] = None, + ) -> torch.Tensor: + # The last layer has no successor, so this one recompute is + # unavoidable -- output-side score weights, matching SGLang's + # layer_idx + 1 >= end_layer branch. Unreachable for K3's capture set + # against 93 layers; kept so a set that does include the final layer + # gets the right tensor rather than the raw prefix sum. + if spec_metadata is not None and len(self.layers) > 0: + last = self.layers[-1] + if capture_set is None or last.layer_idx in capture_set: + tail = ( + _apply_attn_res( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + ) + if num_snapshots > 0 and _AUX_ATTN_RES_STREAM_ENABLED + else hidden_states + ) + spec_metadata.maybe_capture_hidden_states(last.layer_idx, tail, None) + + return _apply_attn_res_and_rmsnorm( + hidden_states, + block_residual[:num_snapshots], + self.output_attn_res_proj, + self.output_attn_res_norm, + self.norm, + max_fused_tokens, + ) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The attention modules: the built-in KDA and MLA modules, with the steps the decode kernels take on those kernels. +# ---------------------------------------------------------------------------------------------------------------------- + + +class K3DecodeKDA(KimiKDALinearAttention): + """FP8 QKVG projection followed by fused gate, convolution and delta recurrence. + + Verification consumes the same FP8 projection rows and writes the per-token + draft records used by the target's speculative worker. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.fp8_decode_ready = False + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + + @property + def takes_k3_kernels(self) -> bool: + return self.qkvg_proj is not None + + def finalize_decode_weights(self) -> None: + super().finalize_decode_weights() + if self.qkvg_proj is None: + raise ValueError("NVFP4 fused KDA requires checkpoint FP8 Q/K/V/G projections") + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + """Whether ``forward`` runs ``step`` on the decode branch: outside a breakable CUDA graph.""" + return not is_in_breakable_cuda_graph() + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + reduce_output: bool = True, + project_output: bool = True, + ) -> torch.Tensor: + """``forward_generic`` where ``will_run_decode_branch`` does not hold. On the decode branch: the plain decode + on ``ssm/k3_kda_decode_attn`` on a decode step of one token per request, else the built-in dispatch; then + ``o_proj`` on its decode GEMV site and the TP all-reduce. + + ``reduce_output=False`` returns ``o_proj``'s TP partial (no all-reduce), and, on the decode branch only, + ``project_output=False`` the post-o_norm core ``[N, H * 128]`` (no ``o_proj``).""" + if not self.will_run_decode_branch(attn_metadata, step): + return self.forward_generic(hidden_states, attn_metadata, reduce_output) + if ( + step.decode + and step.tokens_per_request == 1 + and self.takes_k3_kernels + and self.fp8_decode_ready + ): + core = self._k3_decode(hidden_states[: step.num_tokens], attn_metadata) + else: + core = self._forward_impl(hidden_states, attn_metadata) + if not project_output: + return core.reshape(-1, self.proj_size) + return self._k3_project_output(core, reduce_output) + + def forward_generic( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + reduce_output: bool = True, + ) -> torch.Tensor: + if reduce_output: + return super().forward(hidden_states, attn_metadata) + return self.o_proj( + self._builtin_core(hidden_states, attn_metadata).reshape(-1, self.proj_size) + ) + + def _builtin_core( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """The built-in forward's core ``[N, H, 128]``: inside a breakable CUDA graph from the eager + ``maybe_bcg_kda_core_inplace``, else from ``_forward_impl``.""" + if self.register_to_config and is_in_breakable_cuda_graph(): + core = hidden_states.new_empty( + (hidden_states.shape[0], self.num_heads, self.head_dim), dtype=torch.bfloat16 + ) + maybe_bcg_kda_core_inplace(hidden_states, self.layer_idx_str, core) + return core + return self._forward_impl(hidden_states, attn_metadata) + + def _k3_project_output(self, core: torch.Tensor, reduce_output: bool = True) -> torch.Tensor: + """``o_proj`` on the ``o_proj`` decode GEMV site where it takes the rows (else the module), then, with + ``reduce_output``, the TP all-reduce.""" + core2d = core.reshape(-1, self.proj_size) + out = None + if self.decode_gemvs is not None: + out = self.decode_gemvs.project("o_proj", core2d, self.o_proj.weight) + if out is None: + if reduce_output: + return self._project_output(core) + out = self.o_proj(core2d) + if reduce_output and self._o_allreduce is not None: + out = self._o_allreduce(out) + return out + + def _k3_decode(self, x: torch.Tensor, attn_metadata: AttentionMetadata) -> torch.Tensor: + """``ssm/k3_kda_decode_attn``: the core output ``[R, H, 128]`` of one token of each of the step's R requests; + each slot's conv window and state advance in place.""" + mamba_metadata = attn_metadata.mamba_metadata + slots = getattr(mamba_metadata, "generation_state_indices", None) + if slots is None: + slots = mamba_metadata.state_indices[: x.shape[0]] + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + w_q, w_k, w_v = self._get_mtp_conv_weights() + qkvg = self.qkvg_proj(x) + bfa = torch.nn.functional.linear(x, self._bfa_proj_weight) + core = _fp8_attention.kda_decode( + qkvg, + bfa, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.conv, + layer_cache.temporal, + slots, + float(self.gate_lower_bound), + self.head_k_dim**-0.5, + float(self.o_norm.eps), + ) + self._sync_kda_replay_conv_window(layer_cache, slots, layer_cache.conv) + return core + + def forward_verify( + self, + x2d, + num_steps, + layer_cache, + conv_pool, + ssm_pool, + slot_indices, + output: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The built-in verify, unless the cache manager keeps the per-token states; then the decode kernels'.""" + if not (layer_cache.has_kda_replay_caches and layer_cache.kda_state_tok is not None): + return super().forward_verify( + x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices, output=output + ) + assert self.takes_k3_kernels and self.fp8_decode_ready, ( + f"Kimi K3 KDA layer {self.layer_idx}: the cache manager keeps the per-token verify states, which only " + "the decode kernels write, and this layer has no fused projection weight or buffers" + ) + if ssm_pool.dtype != torch.float32: + raise ValueError("Fused KDA verification requires an FP32 recurrent-state pool") + core = self._k3_verify(x2d, num_steps, layer_cache, ssm_pool, slot_indices) + return self._store_core(core, output) + + def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch.Tensor: + """The core output ``[N T, H, 128]`` of N requests of T verify tokens: ``ssm/k3_kda_attn`` for one request of + 8 tokens, else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it + takes them). Each slot's state after the golden token, its drafts' states and its conv window are written in + place.""" + num_requests = x.shape[0] // num_steps + slots = slot_indices[:num_requests] + w_q, w_k, w_v = self._get_mtp_conv_weights() + constants = (float(self.gate_lower_bound), self.head_k_dim**-0.5, float(self.o_norm.eps)) + rows = torch.cat( + (self.qkvg_proj(x), torch.nn.functional.linear(x, self._bfa_proj_weight)), dim=-1 + ) + out = k3_kda_verify( + rows, + self.f_b_proj.weight, + w_q, + w_k, + w_v, + self._A_log_f32, + self._dt_bias_f32, + self._onorm_w_f32, + layer_cache.kda_conv_q, + layer_cache.kda_conv_k, + layer_cache.kda_conv_v, + ssm_pool, + layer_cache.kda_state_tok, + slots, + layer_cache.prev_num_accepted_tokens, + num_steps - 1, + *constants, + ) + return out.view(-1, self.num_heads, self.head_dim) + + +class K3DecodeMLA(KimiK3MLAAttention): + """FP8 projections with fused query/cache and attention/V-B/gate decode.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self._fp8_qb_weight = self._fp8_qb_scale = None + self._fp8_decode_ready = False + + def prepare_fp8_decode_weights(self) -> None: + """Retain checkpoint Q-B codes before the GEMM backend transforms them.""" + linear = self.q_b_proj + if linear.weight.dtype != torch.float8_e4m3fn: + raise ValueError("NVFP4 fused MLA requires checkpoint FP8 Q-B projections") + scale = linear.weight_scale + expected = (math.ceil(linear.out_features / 128), math.ceil(linear.in_features / 128)) + if tuple(scale.shape) != expected or scale.dtype != torch.float32: + raise ValueError( + f"FP8 Q-B requires raw 128x128 block scales, got {scale.shape} {scale.dtype}" + ) + self._fp8_qb_weight = linear.weight.detach().clone().contiguous() + self._fp8_qb_scale = scale.detach().clone().contiguous() + + def post_load_weights(self) -> None: + super().post_load_weights() + if self._fp8_qb_weight is None: + raise RuntimeError("FP8 MLA decode weights were not prepared during checkpoint loading") + self._fp8_decode_ready = True + + def will_run_decode_branch(self, attn_metadata: AttentionMetadata, step: DecodeStep) -> bool: + return self._k3_step_view(attn_metadata, step) is not None + + def _k3_step_view(self, metadata: AttentionMetadata, step: DecodeStep) -> Optional[dict]: + if not (step.decode and self._fp8_decode_ready and not is_in_breakable_cuda_graph()): + return None + if not self.fuse_qkv_a_proj or self.apply_rotary_emb or self.mapping.has_cp_helix(): + return None + if ( + not self.use_output_gate + or self.sparse_attn_hooks is not None + or ( + self.q_lora_rank, + self.kv_lora_rank, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + ) + != (1536, 512, 128, 64, 128) + ): + return None + view = k3_mla_decode_view(self.mqa, metadata, step.num_tokens) + return None if isinstance(view, str) else view + + def forward( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + step: DecodeStep, + project_output: bool = True, + ) -> torch.Tensor: + view = self._k3_step_view(attn_metadata, step) + if view is None: + if not project_output: + raise ValueError("project_output=False requires fused MLA decode") + return self.forward_generic(position_ids, hidden_states, attn_metadata) + x = hidden_states[: step.num_tokens] + ag = torch.cat((self.kv_a_proj_with_mqa(x), self.g_proj(x)), dim=-1) + fused_q = _fp8_attention.mla_qkv( + ag, + self.q_a_layernorm.weight, + self._fp8_qb_weight, + self._fp8_qb_scale, + self.k_b_proj_trans, + self.k_b_proj_trans_scale, + self.kv_a_layernorm.weight, + view["pool"], + view["page_table"], + view["seq_len"], + view["row_stride"], + view["page_offset"], + view["tokens_per_request"], + float(self.q_a_layernorm.variance_epsilon), + float(self.kv_a_layernorm.variance_epsilon), + ) + core = _fp8_attention.mla_attn( + fused_q, + view["pool"], + view["page_table"], + view["seq_len"], + self.v_b_proj, + self.v_b_proj_scale, + ag, + view["row_stride"], + view["page_offset"], + view["tokens_per_request"], + view["softmax_scale"], + self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim, + ) + if not project_output: + return core + return self._project_output([core], position_ids, attn_metadata, None) + + def forward_generic( + self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + return super().forward(position_ids, hidden_states, attn_metadata) + + +# ---------------------------------------------------------------------------------------------------------------------- +# The DSpark drafter: the stock GQA drafter, with a decode step's block on the drafter entries. +# ---------------------------------------------------------------------------------------------------------------------- + +#: The (requests, block tokens) splits attention/k3_drafter_attn_qknorm certifies; R x 7 is DSpark's block at +#: max_draft_len 7. +DRAFTER_ATTN_SPLITS = frozenset( + {(1, 1), (1, 8), (2, 4), (4, 2), (8, 1), (3, 1), (2, 8), (8, 8)} | {(r, 7) for r in range(1, 9)} +) +# The drafter layout it certifies: 6 query heads and 1 KV head of 64 per rank, context pages of 64 rows, the q / k +# RMSNorm's epsilon and the NeoX RoPE's base. +_DRAFTER_HEADS = (6, 1) +_DRAFTER_HEAD_DIM = 64 +_DRAFTER_PAGE = 64 +_DRAFTER_EPS = 1e-5 +_DRAFTER_ROPE_BASE = 10000.0 + + +def fc_columns(in_features: int, tp_size: int, tp_rank: int) -> Optional[Tuple[int, int]]: + """Rank ``tp_rank``'s input columns ``[start, end)`` of the drafter's context projection ``fc`` split by input + feature over ``tp_size`` ranks: equal contiguous blocks in rank order. None where they do not split evenly.""" + if tp_size < 1 or not 0 <= tp_rank < tp_size or in_features <= 0 or in_features % tp_size: + return None + width = in_features // tp_size + return tp_rank * width, (tp_rank + 1) * width + + +class K3FcSlice(nn.Module): + """This rank's block of the DSpark drafter's context projection ``fc`` (`fc_columns`): ``weight`` is + ``fc.weight[:, start:end]``, contiguous. Its output is this rank's partial product; the sum over the TP group's + ranks is ``fc``'s output.""" + + def __init__(self, weight: torch.Tensor, start: int, end: int) -> None: + super().__init__() + self.weight = nn.Parameter(weight, requires_grad=False) + self.start = start + self.end = end + + @classmethod + def of(cls, fc_weight: torch.Tensor, tp_size: int, tp_rank: int) -> Optional["K3FcSlice"]: + """Rank ``tp_rank``'s block of the full ``fc_weight`` ``[out_features, in_features]``: a copy of its columns + (on one rank, the weight itself); None where the input columns do not split evenly over ``tp_size`` ranks.""" + columns = fc_columns(fc_weight.shape[1], tp_size, tp_rank) + if columns is None: + return None + start, end = columns + return cls(fc_weight.detach()[:, start:end].contiguous(), start, end) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states[:, start:end] @ weight.T`` for the full features ``hidden_states`` ``[N, in_features]``: the + GEMM reads the column block in place, through the rows' stride.""" + return torch.nn.functional.linear(hidden_states[:, self.start : self.end], self.weight) + + +class K3DSparkDrafter(GQADSparkForCausalLM): + """Kimi K3's DSpark drafter: the stock GQA drafter, with a decode step's block forward on the drafter entries. + + A block the entries take runs, per layer: + + * the fused q / k / v projection on the ``drafter_qkv`` decode GEMV site; + * ``attention/k3_drafter_attn_qknorm``: the q / k RMSNorm, the NeoX RoPE, and each request's block attending to + its paged context and to its own k / v, in one launch. The block's k / v are not stored in the cache: the + worker writes the accepted rows' context k / v before a block reads them; + * the output projection on ``drafter_o``, then the module's all-reduce; + * gate / up on ``drafter_gate_up`` and the down projection with the SiLU-and-mul on ``drafter_down``, then the + module's all-reduce. + + The norms and the residual adds are the stock modules' unless the drafter runs on the TP group's collective state + (below); a projection whose site does not take its rows runs its module. Every other block runs the stock + ``dflash_forward``: a split `DRAFTER_ATTN_SPLITS` does not list, another attention backend, head layout, RoPE or + normalization, a cache the kernel does not read, or a compile key of the attention that has not run eagerly, under + CUDA-graph capture. + + Where the target hands the drafter the TP group's collective state (``use_decode_comm``): + + * the context projection's ``fc`` holds only this rank's block of input columns (`K3FcSlice`), and + ``project_target_hidden`` sums the ranks' partial products with ``hidden_norm`` in one all-reduce; + * a block the entries take runs each residual add and RMSNorm in the all-reduce before it (o_proj's with the + post-attention norm, the MLP's with the next layer's input norm, the last one's with the final norm), the + block's input serving as the first residual, uncopied. Up to 8 rows ``comm/k3_sandwich_plain`` runs the + projection, its all-reduce and the norm in one launch (o_proj, and with the SiLU-and-mul the down projection + after ``drafter_gate_up``); otherwise the projection (its site, or its module without the all-reduce) is + followed by ``comm/mnnvl_fusion_allreduce``. A block whose rows the MNNVL workspace does not hold, or layers + whose norms or projections the fused all-reduces do not reproduce, keep the stock all-reduces and norms. + + Otherwise ``fc`` stays replicated, with the stock context projection, and the norms and residual adds are the + stock modules'. The worker that calls the drafter is this target's ``KimiK3DSparkWorker`` (``spec_worker.py``); + the context k / v and the Markov head stay upstream code. + """ + + def __init__( + self, draft_config: ModelConfig, *, dflash_attention_backend: str = "AUTO" + ) -> None: + super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The decode GEMVs' state (decode_gemv.py), shared by the target's layers; set by the target once built. + self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None + # Whether every layer is one the drafter entries take; decided on the first block, the weights loaded. + self._k3_take: Optional[bool] = None + # The attention's compile keys that ran eagerly ((more than one request, page stride)); a capture takes + # only these. + self._k3_attn_ran: set = set() + # The TP group's collective state (decode_comm.py), set by the target where it builds one (use_decode_comm). + self.decode_comm: Optional[_decode_comm.K3DecodeComm] = None + # The zero residual of the split context projection's all-reduce, up to a decode step's rows. + self._k3_zero_rows: Optional[torch.Tensor] = None + # Whether a block's residual adds and RMSNorms run in its all-reduces, and the comm/k3_sandwich_plain forms + # ("o_proj", "down") compiled for the layers' shapes; set with the collective state. + self._k3_norms_fuse = False + self._k3_sandwich_forms: frozenset = frozenset() + + def use_decode_comm(self, comm: _decode_comm.K3DecodeComm) -> None: + """Run the drafter's collectives on ``comm``, the TP group's decode state (``decode_comm.K3DecodeComm``) the + target builds where every attention all-reduce runs over MNNVL. The target calls it once the weights are + loaded, before any CUDA-graph capture. Collective: every rank of the group calls it at the same point. + + * The context projection's ``fc`` keeps only this rank's block of input columns (`K3FcSlice`), where they + split evenly over the group and the drafter's layers have their TP all-reduce: ``project_target_hidden`` + sums the ranks' partial products and applies ``hidden_norm`` in one ``comm/mnnvl_fusion_allreduce`` call. + * Where the layers' norms and projections take it (`_k3_norms_take_comm`), a block's residual adds and + RMSNorms run in its all-reduces (``_k3_block_forward``). Each form of ``comm/k3_sandwich_plain`` whose + shape every layer shares compiles here, with one call on a zero row of a zero weight; a form that does not + compile is not used.""" + self.decode_comm = comm + self._k3_split_fc() + self._k3_norms_fuse = self._k3_norms_take_comm() + forms = set() + if self._k3_norms_fuse: + layers = self.model.layers + for form, swiglu, weights in ( + ("o_proj", False, [layer.self_attn.o_proj.weight for layer in layers]), + ("down", True, [layer.mlp.down_proj.weight for layer in layers]), + ): + if len({tuple(w.shape) for w in weights}) == 1 and comm.compile_plain( + weights[0], swiglu=swiglu + ): + forms.add(form) + self._k3_sandwich_forms = frozenset(forms) + logger.info( + "Kimi K3 DSpark drafter: residual adds and RMSNorms " + + ( + f"in the all-reduces (k3_sandwich_plain: {sorted(forms) or 'none'}, else " + "mnnvl_fusion_allreduce)" + if self._k3_norms_fuse + else "on the stock modules (a layer norm or projection the fused all-reduces do not reproduce)" + ) + ) + + def _k3_norms_take_comm(self) -> bool: + """Whether a block's residual adds and RMSNorms can run in its all-reduces: every layer's input and + post-attention norms and the final norm are plain bf16 RMSNorms of the hidden width (no Gemma offset, no + quantized output), and every output and down projection is row parallel with its all-reduce.""" + hidden = self.config.hidden_size + + def plain(norm: nn.Module) -> bool: + weight = getattr(norm, "weight", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (hidden,) + and hasattr(norm, "variance_epsilon") + and not getattr(norm, "use_gemma", True) + and not getattr(norm, "is_nvfp4", True) + and not getattr(norm, "return_hp_output", True) + ) + + def reduces(linear: nn.Module) -> bool: + return ( + getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" + and getattr(linear, "reduce_output", False) + and getattr(linear, "all_reduce", None) is not None + ) + + return plain(self.model.norm) and all( + plain(layer.input_layernorm) + and plain(layer.post_attention_layernorm) + and reduces(layer.self_attn.o_proj) + and reduces(layer.mlp.down_proj) + for layer in self.model.layers + ) + + def _k3_tp_all_reduce(self) -> Optional[nn.Module]: + """The drafter's own TP all-reduce (its first output projection's row-parallel all-reduce module), or None.""" + o_proj = self.model.layers[0].self_attn.o_proj + if getattr(getattr(o_proj, "tp_mode", None), "name", None) != "ROW": + return None + return getattr(o_proj, "all_reduce", None) + + def _k3_split_fc(self) -> None: + """Replace the replicated ``fc`` with this rank's block of its input columns (`K3FcSlice`) and size the zero + residual of the projection's all-reduce, where the drafter runs on the TP group's collective state, ``fc`` is + a bias-free bf16 projection whose input columns split evenly over the group, and the drafter has its TP + all-reduce (for rows the MNNVL workspace does not hold). Otherwise ``fc`` stays replicated.""" + fc = getattr(self, "fc", None) + if self.decode_comm is None or fc is None or isinstance(fc, K3FcSlice): + return + mapping = self.model_config.mapping + weight = getattr(fc, "weight", None) + sliced = None + if ( + isinstance(weight, torch.Tensor) + and weight.dim() == 2 + and weight.dtype == torch.bfloat16 + and getattr(fc, "bias", None) is None + and self._k3_tp_all_reduce() is not None + ): + sliced = K3FcSlice.of(weight, mapping.tp_size, mapping.tp_rank) + if sliced is None: + logger.info( + "Kimi K3 DSpark drafter: fc stays replicated (it is not a bias-free bf16 projection whose input " + f"columns split evenly over TP{mapping.tp_size}, or the layers have no TP all-reduce)" + ) + return + self.fc = sliced + self._k3_zero_rows = weight.new_zeros( + MAX_REQUESTS * MAX_TOKENS_PER_REQUEST, weight.shape[0] + ) + logger.info( + f"Kimi K3 DSpark drafter: fc split by input feature over TP{mapping.tp_size}: rank {mapping.tp_rank} " + f"holds columns [{sliced.start}, {sliced.end}), hidden_norm in the all-reduce" + ) + + def load_weights(self, weights, weight_mapper=None, **kwargs): + """The stock load; on the TP group's collective state, ``fc`` is split again (`_k3_split_fc`).""" + result = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + self._k3_split_fc() + return result + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_norm(fc(hidden_states))`` of the captured target features ``[N, in_features]``. + + With ``fc`` split (`K3FcSlice`): this rank's partial product of its columns, then the sum over the TP group + with ``hidden_norm`` applied in one ``comm/mnnvl_fusion_allreduce`` call (a zero residual) up to a decode + step's rows the MNNVL workspace holds; more rows (a prefill) go through the drafter's TP all-reduce, then + ``hidden_norm``. Otherwise the stock projection.""" + fc = self.fc + if not isinstance(fc, K3FcSlice): + return super().project_target_hidden(hidden_states) + partial = fc(hidden_states.to(fc.weight.dtype)) + rows = partial.shape[0] + zeros = self._k3_zero_rows + if rows <= zeros.shape[0] and self.decode_comm.takes_allreduce_norm(rows, partial.shape[1]): + normed, _ = self.decode_comm.allreduce_norm(partial, zeros[:rows], self.hidden_norm) + return normed + if rows == 0: + return self.hidden_norm(partial) + return self.hidden_norm(self._k3_tp_all_reduce()(partial)) + + def dflash_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_k_cache: torch.Tensor, + ctx_v_cache: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor] = None, + ctx_page_table: Optional[torch.Tensor] = None, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + """The block's hidden states ``[B * block, hidden]``: on the drafter entries where they take it (the class + docstring), else the stock forward. ``ctx_rows_start``: where ``ctx_cache_batch_idx`` is the contiguous rows + ``[ctx_rows_start, ctx_rows_start + B)`` of ``ctx_page_table``, their start; the entries then read those rows + as a view instead of gathering them.""" + keys = self._k3_block_keys(noise_embedding, ctx_kv_cache, ctx_page_table) + if keys is None: + return super().dflash_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_k_cache, + ctx_v_cache, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ) + out = self._k3_block_forward( + noise_embedding, + query_positions, + num_ctx_per_req, + ctx_cache_batch_idx, + ctx_kv_cache, + ctx_page_table, + ctx_rows_start, + ) + if not torch.cuda.is_current_stream_capturing(): + self._k3_attn_ran |= keys + return out + + def _k3_layers_take(self) -> bool: + """Whether every layer is one the drafter entries take: plain NeoX RoPE with the q / k RMSNorm at the head + layout, epsilon and RoPE base the attention certifies, bias-free projections, a plain SiLU-and-mul MLP, + non-causal unwindowed block attention, and no block convolution or post-attention gate.""" + if self._k3_take is None: + if self._fused_kv_weight is None: + self._build_fused_kv_buffers() + take = ( + self._use_fused_qk_norm_rope + and not self.has_block_conv + and type(self)._post_attention_gate is DFlashForCausalLM._post_attention_gate + ) + for layer_idx, layer in enumerate(self.model.layers): + attn, mlp = layer.self_attn, layer.mlp + pos = getattr(attn, "pos_embd_params", None) + rope = getattr(pos, "rope", None) + norms = (attn.q_norm, attn.k_norm) + take = take and ( + (attn.num_heads, attn.num_key_value_heads) == _DRAFTER_HEADS + and attn.head_dim == _DRAFTER_HEAD_DIM + and getattr(attn, "is_qk_norm", False) + and not getattr(attn, "use_gemma_rms_norm", True) + and rope is not None + and pos.is_neox + and getattr(pos, "mrope_section", None) is None + and rope.theta == _DRAFTER_ROPE_BASE + and all( + norm.variance_epsilon == _DRAFTER_EPS + and norm.weight.dtype == torch.bfloat16 + and tuple(norm.weight.shape) == (_DRAFTER_HEAD_DIM,) + for norm in norms + ) + and attn.qkv_proj.bias is None + and attn.o_proj.bias is None + and mlp.activation is torch.nn.functional.silu + and mlp.swiglu_limit is None + and mlp.swiglu_alpha in (None, 1.0) + and mlp.swiglu_beta in (None, 0.0) + and mlp.gate_up_proj.bias is None + and mlp.down_proj.bias is None + and self._resolve_block_attention(layer_idx) == (False, (-1, -1)) + ) + self._k3_take = bool(take) + logger.info( + f"Kimi K3 DSpark drafter: {len(self.model.layers)} layers, decode blocks " + + ( + "on k3_drafter_attn_qknorm and the drafter GEMV sites" + if self._k3_take + else "on the stock block forward (a layer the drafter entries do not take)" + ) + ) + return self._k3_take + + def _k3_block_keys( + self, + noise_embedding: torch.Tensor, + ctx_kv_cache: Optional[torch.Tensor], + ctx_page_table: Optional[torch.Tensor], + ) -> Optional[frozenset]: + """The attention's compile keys for this block when the drafter entries take it, else None.""" + if ( + self.dflash_attention_backend != "TRTLLM" + or ctx_kv_cache is None + or ctx_page_table is None + or noise_embedding.dtype != torch.bfloat16 + or tuple(noise_embedding.shape[:2]) not in DRAFTER_ATTN_SPLITS + or is_in_breakable_cuda_graph() + or not self._k3_layers_take() + ): + return None + num_kv_heads = self.model.layers[0].self_attn.num_key_value_heads + page = (2, num_kv_heads, _DRAFTER_PAGE, _DRAFTER_HEAD_DIM) + dense = ( + num_kv_heads * _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + _DRAFTER_PAGE * _DRAFTER_HEAD_DIM, + ) + strides = set() + for layer_idx in range(len(self.model.layers)): + cache = ctx_kv_cache[layer_idx] + if ( + cache.dtype != torch.bfloat16 + or cache.dim() != 5 + or tuple(cache.shape[1:]) != page + or cache.stride()[1:] != (*dense, _DRAFTER_HEAD_DIM, 1) + or cache.stride(0) % 8 + ): + return None + strides.add(cache.stride(0)) + keys = frozenset((noise_embedding.shape[0] > 1, stride) for stride in strides) + if torch.cuda.is_current_stream_capturing() and not keys <= self._k3_attn_ran: + return None + return keys + + def _k3_block_forward( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + ctx_cache_batch_idx: torch.Tensor, + ctx_kv_cache: torch.Tensor, + ctx_page_table: torch.Tensor, + ctx_rows_start: Optional[int] = None, + ) -> torch.Tensor: + batch, block = noise_embedding.shape[:2] + rows = batch * block + ctx_len = num_ctx_per_req[:batch].to(torch.int32) + if ctx_rows_start is None: + page_table = ctx_page_table.index_select(0, ctx_cache_batch_idx.to(torch.long)) + else: + # The batch's page-table rows are one contiguous run: a view, no gather. + page_table = ctx_page_table[ctx_rows_start : ctx_rows_start + batch] + positions = query_positions.reshape(-1).contiguous() + hidden = noise_embedding.reshape(rows, -1) + layers = self.model.layers + comm = self._k3_fused_comm(hidden) + residual = None + if comm is not None: + # The fused all-reduces read the residual and return the updated one: the block's input serves as the + # first residual, uncopied. + residual = hidden + normed = layers[0].input_layernorm(hidden) + for layer_idx, layer in enumerate(layers): + attn = layer.self_attn + if comm is None: + if residual is None: + residual = hidden.clone() + normed = layer.input_layernorm(hidden) + else: + normed, residual = layer.input_layernorm(hidden, residual) + qkv = self._k3_project("drafter_qkv", normed, attn.qkv_proj) + out = torch.empty(rows, attn.q_size, dtype=torch.bfloat16, device=qkv.device) + k3_drafter_attn_qknorm( + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + positions, + attn.q_norm.variance_epsilon, + attn.pos_embd_params.rope.theta, + ctx_kv_cache[layer_idx], + page_table, + ctx_len, + attn.num_heads, + attn.num_key_value_heads, + out, + ) + if comm is not None: + # o_proj's all-reduce applies the post-attention norm; the MLP's the next layer's input norm, after + # the last layer the final norm. + next_norm = ( + layers[layer_idx + 1].input_layernorm + if layer_idx + 1 < len(layers) + else self.model.norm + ) + normed, residual = self._k3_project_norm( + comm, out, attn.o_proj, residual, layer.post_attention_layernorm + ) + normed, residual = self._k3_mlp_norm(comm, layer.mlp, normed, residual, next_norm) + continue + hidden = self._k3_project("drafter_o", out, attn.o_proj) + hidden, residual = layer.post_attention_layernorm(hidden, residual) + hidden = self._k3_mlp(layer.mlp, hidden) + if comm is not None: + return normed + out, _ = self.model.norm(hidden, residual) + return out + + def _k3_fused_comm(self, hidden: torch.Tensor) -> Optional[_decode_comm.K3DecodeComm]: + """The TP group's collective state where a block of rows ``hidden`` runs its residual adds and RMSNorms in its + all-reduces: the drafter runs on it (``use_decode_comm``) with layers that take it, and its MNNVL workspace + holds an all-reduce of the block's rows; else None (the stock all-reduces and norms).""" + comm = self.decode_comm + if comm is None or not self._k3_norms_fuse or not comm.takes_allreduce_norm(*hidden.shape): + return None + return comm + + def _k3_project_norm( + self, + comm: _decode_comm.K3DecodeComm, + x: torch.Tensor, + linear: nn.Module, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + linear(x)`` for the row-parallel output projection, + the residual add and ``norm`` in its all-reduce: ``comm/k3_sandwich_plain`` where it takes the rows (the + projection in the ``drafter_o`` site's arithmetic, in the same launch), else the projection on ``drafter_o`` + where that takes the rows, or the module without its all-reduce, then ``comm/mnnvl_fusion_allreduce``.""" + if "o_proj" in self._k3_sandwich_forms and comm.takes_plain( + x, linear.weight, residual, norm + ): + return comm.sandwich_plain(x, linear.weight, residual, norm) + gemvs = self.decode_gemvs + partial = None if gemvs is None else gemvs.project("drafter_o", x, linear.weight) + if partial is None: + partial = linear(x, all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_mlp_norm( + self, + comm: _decode_comm.K3DecodeComm, + mlp: nn.Module, + x: torch.Tensor, + residual: torch.Tensor, + norm: nn.Module, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """``(norm(updated), updated)``, ``updated = residual + mlp(x)``, the residual add and ``norm`` in the down + projection's all-reduce: gate / up on ``drafter_gate_up``, then ``comm/k3_sandwich_plain``'s SiLU-and-mul + form where it takes the rows (the ``drafter_down`` site's arithmetic), else the SiLU-and-mul and down + projection on ``drafter_down``; where those sites do not take the rows, the module without its all-reduce; + then ``comm/mnnvl_fusion_allreduce``.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + partial = None + if gate_up is not None: + if "down" in self._k3_sandwich_forms and comm.takes_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ): + return comm.sandwich_plain( + gate_up, mlp.down_proj.weight, residual, norm, swiglu=True + ) + partial = gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + if partial is None: + partial = mlp(x, final_all_reduce_params=_decode_comm.skip_all_reduce()) + return comm.allreduce_norm(partial, residual, norm) + + def _k3_project(self, site: str, x: torch.Tensor, linear: nn.Module) -> torch.Tensor: + """``linear(x)``, its GEMM on the ``site`` decode GEMV where that takes the rows, then a row-parallel + projection's all-reduce.""" + y = None if self.decode_gemvs is None else self.decode_gemvs.project(site, x, linear.weight) + if y is None: + return linear(x) + return _row_parallel_reduce(linear, y) + + def _k3_mlp(self, mlp: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``mlp(x)``: gate / up on ``drafter_gate_up`` and the SiLU-and-mul with the down projection on + ``drafter_down`` where both take the rows, then the down projection's all-reduce.""" + gemvs = self.decode_gemvs + gate_up = ( + None if gemvs is None else gemvs.project("drafter_gate_up", x, mlp.gate_up_proj.weight) + ) + down = ( + None + if gate_up is None + else gemvs.project("drafter_down", gate_up, mlp.down_proj.weight) + ) + if down is None: + return mlp(x) + return _row_parallel_reduce(mlp.down_proj, down) + + +def _row_parallel_reduce(linear: nn.Module, partial: torch.Tensor) -> torch.Tensor: + """``partial`` all-reduced as ``linear`` reduces its output: a row-parallel projection's all-reduce, if any.""" + all_reduce = getattr(linear, "all_reduce", None) + if getattr(getattr(linear, "tp_mode", None), "name", None) == "ROW" and all_reduce is not None: + return all_reduce(partial) + return partial + + +# ---------------------------------------------------------------------------------------------------------------------- +# The target: step classification, the construction checks and the class route returns. +# ---------------------------------------------------------------------------------------------------------------------- + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The language model's ModelConfig: the checkpoint's text_config, with quant exclusions renamed to match. + + The checkpoint names its language-model modules `language_model.`; the text model's are ``, with + `layers.*` under `model.`. + """ + config = model_config.pretrained_config + text = copy.copy(model_config) + text._frozen = False + text.pretrained_config = config.text_config + excluded = text.quant_config.exclude_modules + if excluded: + text.quant_config = copy.copy(text.quant_config) + renamed = [] + for name in excluded: + if name.startswith(_LANG_PREFIX): + name = name[len(_LANG_PREFIX) :] + if name.startswith("layers."): + name = "model." + name + renamed.append(name) + text.quant_config.exclude_modules = renamed + text.skip_create_weights_in_init = True + text._frozen = True + return text + + +@dataclass(frozen=True) +class DecodeStep: + """A step the Kimi K3 decode kernels take: ``num_tokens`` rows and, on a pure decode step, ``num_requests`` + generation requests of ``tokens_per_request`` tokens each (None on a step with context requests).""" + + num_tokens: int + num_requests: Optional[int] = None + tokens_per_request: Optional[int] = None + #: Whether the MoE layers' latent all-reduce is the routed experts' push form plus ``comm/k3_latent_reduce`` + #: (``decode_moe.py``): set by ``KimiLinearModel.forward`` from ``latent_push``. + latent_push: bool = False + + @property + def small(self) -> bool: + """Whether the step fits one token tile of the token-count kernels.""" + return self.num_tokens <= DECODE_MAX_TOKENS + + @property + def decode(self) -> bool: + """Whether the step is a pure decode step the request-aware kernels take.""" + return self.num_requests is not None + + @property + def wide(self) -> bool: + """Whether the step is a pure decode step of more than one token tile: its token-count work keeps the decode + layout's MoE head and tail, on M-general ops.""" + return self.decode and not self.small + + +def latent_push( + step: DecodeStep, + *, + capturing: bool, + breakable: bool, + kda_draft_records: bool, + kda_decode_kernels: bool, + mla_decode_branch: bool, +) -> bool: + """Whether a step's MoE layers push their routed partials (``DecodeStep.latent_push``): a pure decode step of at + most 8 tokens whose attention layers all run the decode kernels, captured into a CUDA graph and not inside a + breakable one. The KDA layers run them on one token per request with every KDA layer taking the K3 kernels + (``kda_decode_kernels``), and on every verify width with the per-token states (``kda_draft_records``); the MLA + layers where their decode branch takes the step (``mla_decode_branch``, the same for every MLA layer of a step). + Every other step keeps the routed experts' all-reduce: the exchange's call-order invariant needs every kernel + between a reduce and the next push to wait for its predecessor, which only those kernels were checked for, and a + graph launch orders every pushing replay behind whatever ran before it.""" + if not (step.decode and step.small) or not capturing or breakable: + return False + if not (kda_draft_records or (step.tokens_per_request == 1 and kda_decode_kernels)): + return False + return mla_decode_branch + + +def decode_step(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[DecodeStep]: + """The step's shape if any Kimi K3 decode kernel takes it, else None (the generic path runs). + + ``num_tokens`` is the step's token count (the rows of the model input). Read on the host from per-step integers + and the host copy of the sequence lengths only. A CUDA graph is captured per decode batch shape, and every input + here is fixed by that shape, so a captured step and its replays are classified alike. + """ + if num_tokens <= 0: + return None + requests = _decode_requests(attn_metadata, num_tokens) + if requests is not None: + return DecodeStep(num_tokens, requests, num_tokens // requests) + if num_tokens <= DECODE_MAX_TOKENS: + return DecodeStep(num_tokens) + return None + + +def _decode_requests(attn_metadata: AttentionMetadata, num_tokens: int) -> Optional[int]: + """R when the step is R <= 8 generation requests of the same T <= 8 tokens and no context request, else None.""" + if attn_metadata.num_contexts != 0: + return None + requests = attn_metadata.num_generations + if not 0 < requests <= MAX_REQUESTS or num_tokens % requests != 0: + return None + tokens = num_tokens // requests + if not 0 < tokens <= MAX_TOKENS_PER_REQUEST: + return None + seq_lens = getattr(attn_metadata, "seq_lens", None) + if seq_lens is not None and seq_lens.device.type == "cpu": + lens = seq_lens[:requests] + if lens.numel() != requests or bool((lens != tokens).any()): + return None + return requests + + +def _check_construction(model_config: ModelConfig, target: K3TargetSpec) -> None: + """Validate the concrete target before constructing modules or workspaces.""" + capability = torch.cuda.get_device_capability() + assert capability == target.sm, ( + f"this target requires sm_{target.sm[0]}{target.sm[1]}, running on " + f"sm_{capability[0]}{capability[1]}" + ) + mapping = model_config.mapping + topology = ( + mapping.world_size, + mapping.tp_size, + mapping.pp_size, + mapping.moe_tp_size, + mapping.moe_ep_size, + mapping.enable_attention_dp, + ) + parallel = f"tp16_moetp{target.moe_tp}ep{target.moe_ep}" + assert topology == (16, 16, 1, target.moe_tp, target.moe_ep, False), ( + f"the {parallel} target needs world_size 16, tensor_parallel_size 16, " + f"pipeline_parallel_size 1, moe_tensor_parallel_size {target.moe_tp}, " + f"moe_expert_parallel_size {target.moe_ep} and enable_attention_dp false; " + f"the engine built (world, tp, pp, moe_tp, moe_ep, attention_dp) = {topology}" + ) + if target.moe_tp == 16: + assert mapping.moe_tp_ep_user_specified, ( + f"the {parallel} target needs moe_tensor_parallel_size 16 and " + "moe_expert_parallel_size 1 set explicitly" + ) + assert getattr(model_config, "spec_config", None) is None, ( + f"the {parallel} target decodes without speculation; for DSpark, DFlash " + "or SA set moe_tensor_parallel_size 4 and moe_expert_parallel_size 4" + ) + assert model_config.torch_dtype == torch.bfloat16, ( + f"this target computes in bf16; the engine resolved dtype {model_config.torch_dtype}" + ) + if target.quant_algo == QuantAlgo.NVFP4: + for layer_idx in range(1, 93): + quant = KimiK3MoERuntime._resolve_routed_quant_config(model_config, layer_idx) + if quant.quant_algo != target.quant_algo or quant.group_size != 16: + raise ValueError( + f"NVFP4 target requires group-size-16 NVFP4 experts at layer {layer_idx}; " + f"got {quant.quant_algo}, group_size={quant.group_size}" + ) + else: + quant = model_config.quant_config + assert quant.kv_cache_quant_algo is None, ( + "this target's MLA kernels read a bf16 KV pool; " + f"kv_cache_config.dtype resolved to {quant.kv_cache_quant_algo}" + ) + assert quant.quant_algo == QuantAlgo.W4A16_MXFP4 and not model_config.quant_config_dict, ( + "this target loads the MXFP4 checkpoint with no per-layer declarations; " + f"the engine read {quant.quant_algo} with " + f"{len(model_config.quant_config_dict or {})} per-layer declarations" + ) + strategy = model_config.allreduce_strategy + assert strategy in (AllReduceStrategy.AUTO, AllReduceStrategy.MNNVL), ( + f"this target runs its all-reduces over MNNVL; allreduce_strategy is {strategy.name}" + ) + + +class ModelingV2KimiK3(KimiLinearForCausalLM): + """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. + + It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through + weights.py, the KDA metadata class, the model defaults), which walk the model by its module names; the text model + keeps the built-in one's. + """ + + target_spec: K3TargetSpec + + @classmethod + def get_preferred_kv_cache_manager_version(cls, pretrained_config: Any = None) -> Literal["V2"]: + """The V2 hybrid manager holds the KDA states; the step contract requires it.""" + return "V2" + + def __init__(self, model_config: ModelConfig): + config = model_config.pretrained_config + assert getattr(config, "text_config", None) is not None, ( + "this target loads the KimiK3ForConditionalGeneration checkpoint, whose language model is its " + "text_config" + ) + _check_construction(model_config, self.target_spec) + text = _text_model_config(model_config) + spec_config = getattr(text, "spec_config", None) + assert ( + spec_config is None + or spec_config.spec_dec_mode.is_sa() + or spec_config.spec_dec_mode.is_dflash() + or spec_config.spec_dec_mode.is_dspark() + ), "Kimi K3 supports speculative decoding only with SA, DFlash or DSpark" + # The inherited loader reads these: this target has neither helix context parallelism nor the fp8 + # weight-read conversion of the shared / latent MLPs. + self._fp8_weight_read_moe_mlp = False + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + cfg = text.pretrained_config + SpecDecOneEngineForCausalLM.__init__( + self, + KimiLinearModel(text), + text, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self._own_spec_branch(text) + self._step_checked = False + self._decode_state_ready = False + # The LM head on gemm/k3_head_gemv at decode size, once cache_derived_state has built its state. It is the + # processor the speculative worker's target logits and a parallel drafter's logits call too. + stock_logits_processor = self.logits_processor + self.logits_processor = _decode_gemv.K3LogitsProcessor(stock_logits_processor) + draft_model = getattr(self, "draft_model", None) + if getattr(draft_model, "logits_processor", None) is stock_logits_processor: + draft_model.logits_processor = self.logits_processor + # The executor reads generation settings (eos_token_id, ...) off the model config the engine holds, which + # must therefore be the text config, as the built-in wrapper leaves it. + model_config._frozen = False + model_config.pretrained_config = self.config + model_config._frozen = True + + def _own_spec_branch(self, model_config: ModelConfig) -> None: + """Put this target's drafter and worker in place of the shell's stock ones: DSpark's drafter from a + standalone GQA checkpoint is `K3DSparkDrafter`, and the DFlash / standalone DSpark worker is + `spec_worker.py`'s. Under the engine's meta initialization the stock drafter it replaces holds no weights.""" + spec_config = model_config.spec_config + if spec_config is None: + return + mode = spec_config.spec_dec_mode + standalone_dspark = mode.is_dspark() and not draft_is_embedded_in_target(model_config) + if ( + standalone_dspark + and self.draft_config is not None + and is_mla(self.draft_config.pretrained_config) + ): + # The MLA drafter's compressed context cache requires its stock worker. + return + if ( + standalone_dspark + and self.draft_config is not None + and not is_mla(self.draft_config.pretrained_config) + ): + drafter = K3DSparkDrafter( + self.draft_config, dflash_attention_backend=spec_config.attention_backend + ) + if mode.is_parallel_draft(): + drafter.logits_processor = self.logits_processor + self.epilogue[self.epilogue.index(self.draft_model)] = drafter + self.draft_model = drafter + if mode.is_dflash(): + worker_cls = _spec_worker.KimiK3DFlashWorker + elif standalone_dspark: + worker_cls = _spec_worker.KimiK3DSparkWorker + else: + return + worker = worker_cls(spec_config, model_config.mapping, self.use_separate_draft_kv_cache) + worker.set_draft_model(self.draft_model) + self.epilogue[self.epilogue.index(self.spec_worker)] = worker + self.spec_worker = worker + + def load_weights(self, weights, *args, **kwargs): + _weights.load(self, weights) + + def _load_attention_fp8(self, weights: Dict[str, torch.Tensor], prefix: str) -> int: + count = super()._load_attention_fp8(weights, prefix) + for layer in self.model.layers: + if not layer.is_kda: + layer.self_attn.mixer.prepare_fp8_decode_weights() + return count + + def cache_derived_state(self) -> None: + """Build the decode GEMVs' state once the weights are final: the LM head's workspace, and one eager call of + every decode GEMV kernel at its site's shape (decode_gemv.SITES), so none compiles under capture. Built once: + a later call keeps it, since CUDA graphs captured in between hold its workspace. The DSpark drafter's sites + are among them, and the drafter gets the state too.""" + super().cache_derived_state() + if self.model.decode_gemvs is not None: + return + sites = tuple(_decode_gemv.SITES) + if self.target_spec.quant_algo == QuantAlgo.NVFP4: + sites = ("dense_gate_up", "dense_down") + if isinstance(self.draft_model, K3DSparkDrafter): + sites += ("drafter_qkv", "drafter_o", "drafter_gate_up", "drafter_down") + gemvs = _decode_gemv.K3DecodeGemvs.create(self.lm_head, sites=sites) + self.model.decode_gemvs = gemvs + self.logits_processor.gemvs = gemvs + for layer in self.model.layers: + if not layer.is_moe: + layer.decode_gemvs = gemvs + if isinstance(self.draft_model, K3DSparkDrafter): + self.draft_model.decode_gemvs = gemvs + + def post_load_weights(self) -> None: + """Enable fused FP8 decode after weight preparation and build the TP collective state. + + Every rank initializes its MNNVL state before CUDA Graph capture. The + output projections retain FP8 GEMMs and use the attention-residual + all-reduce without the BF16 projection sandwich. + """ + super().post_load_weights() + if self._decode_state_ready: + return + if self.target_spec.quant_algo == QuantAlgo.NVFP4: + layers = self.model.layers + for layer in layers: + if layer.is_kda: + layer.linear_attn.fp8_decode_ready = True + comm = None + if all(layer._mnnvl_allreduce() is not None for layer in layers): + comm = _decode_comm.K3DecodeComm.create(self.model_config.mapping, None) + for layer in layers: + layer.decode_comm = comm + layer.sandwich_oproj = False + _decode_comm.use_decode_one_shot(self) + self._gate_spec_worker_kernels(comm) + self._gate_drafter_comm(comm) + self._decode_state_ready = True + logger.info( + f"Kimi K3 NVFP4 SM{self.target_spec.sm[0]}{self.target_spec.sm[1]}: " + "fused FP8 attention and NVFP4 experts; " + "BF16 decode GEMVs and DSpark specializations; " + f"attention-residual MNNVL fusion={comm is not None}" + ) + return + + def _gate_spec_worker_kernels(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Turn the DFlash / DSpark worker's Kimi K3 decode kernels (its ``k3_decode``: ``trtllm::k3_spec_accept``, + ``k3_ctx_kv`` and ``k3_markov``, on target and draft logits kept vocabulary-sharded) on only when every input + that path needs exists: the TP group's collective state over MNNVL (``comm``) and the LM head's + ``gemm/k3_head_gemv`` workspace. ``K3LogitsProcessor.lm_head_shard`` needs that workspace to produce the + vocabulary shard, so the workspace is a precondition of the path, not a policy choice. The worker still checks + each step's own conditions. Returns the setting (False without such a worker).""" + worker = getattr(self, "spec_worker", None) + if not hasattr(worker, "k3_decode"): + return False + gemvs = self.model.decode_gemvs + worker.k3_decode = ( + comm is not None and gemvs is not None and gemvs.head_workspace is not None + ) + logger.info( + f"Kimi K3 decode kernels: {type(worker).__name__} k3_spec_accept, k3_ctx_kv and k3_markov " + + ( + "on" + if worker.k3_decode + else "off (no MNNVL decode state or no k3_head_gemv LM head)" + ) + ) + return worker.k3_decode + + def _gate_drafter_comm(self, comm: Optional[_decode_comm.K3DecodeComm]) -> bool: + """Hand the DSpark drafter (`K3DSparkDrafter`) the TP group's collective state ``comm`` where this target built + it (every attention all-reduce over MNNVL; TP16 is a construction assert): the drafter then runs its context + projection split over the group, ``hidden_norm`` in the all-reduce, and its blocks' residual adds and RMSNorms + in their all-reduces (``K3DSparkDrafter.use_decode_comm``, collective: it compiles the drafter's sandwich on + the group's workspace). Without ``comm`` it keeps the stock replicated ``fc``, all-reduces and norms. The gate + requires exactly what this path uses: ``comm`` and nothing else, so not the LM head's ``k3_head_gemv`` + workspace, which only the speculative worker's path (`_gate_spec_worker_kernels`) reads. Returns whether the + drafter took the state (False without such a drafter).""" + drafter = getattr(self, "draft_model", None) + if comm is None or not isinstance(drafter, K3DSparkDrafter): + return False + drafter.use_decode_comm(comm) + return True + + def _check_step_contract(self, attn_metadata: AttentionMetadata) -> None: + """First-forward checks of the engine surface and the per-engine settings.""" + objects = { + "attn_metadata": attn_metadata, + "kv_cache_manager": attn_metadata.kv_cache_manager, + } + missing = [ + f"{owner}.{name}" + for owner, names in REQUIRED_ENGINE_FIELDS.items() + for name in names + if not hasattr(objects[owner], name) + ] + assert not missing, f"engine fields this target reads are missing: {missing}" + manager = attn_metadata.kv_cache_manager + assert isinstance(manager, MambaHybridCacheManagerV2), ( + "this target needs the V2 hybrid KV / state cache manager " + "(kv_cache_config.use_kv_cache_manager_v2); the engine built " + f"{type(manager).__name__}" + ) + if self.target_spec.quant_algo != QuantAlgo.NVFP4: + assert attn_metadata.tokens_per_block == _TOKENS_PER_BLOCK, ( + f"this target needs kv_cache_config.tokens_per_block {_TOKENS_PER_BLOCK}; the engine built " + f"{attn_metadata.tokens_per_block}" + ) + assert not manager.enable_block_reuse, ( + "this target runs with kv_cache_config.enable_block_reuse false; the engine enabled it" + ) + kda_layer = next(layer.layer_idx for layer in self.model.layers if layer.is_kda) + state_dtype = manager.mamba_layer_cache(kda_layer).temporal.dtype + assert state_dtype == torch.float32, ( + "this target's KDA kernels keep fp32 recurrent states " + "(kv_cache_config.mamba_ssm_cache_dtype float32 or " + f"auto); the engine built a {state_dtype} state pool" + ) + self._step_checked = True + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + return_context_logits: bool = False, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + if kwargs.pop("multimodal_params", None): + raise ValueError( + "this Kimi K3 target is text only: it loads no vision tower, and a request carried image input" + ) + if not self._step_checked: + self._check_step_contract(attn_metadata) + return super().forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + return_context_logits, + spec_metadata, + resource_manager, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/spec_worker.py new file mode 100644 index 000000000000..abc1a1388213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/spec_worker.py @@ -0,0 +1,1343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import functools +import inspect +from typing import Optional + +import torch +from torch import nn + +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.speculative.dflash import ( + DFlashSpecMetadata, + DFlashWorker, + dflash_allocated_ctx_limit, + dflash_noise_block_embedding, +) +from tensorrt_llm._torch.speculative.dspark import DSparkWorker +from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase +from tensorrt_llm.logger import logger + + +@functools.lru_cache(maxsize=None) +def _takes_ctx_rows_start_cls(cls) -> bool: + return "ctx_rows_start" in inspect.signature(cls.dflash_forward).parameters + + +def _takes_ctx_rows_start(draft_model) -> bool: + """Whether the draft model's ``dflash_forward`` accepts ``ctx_rows_start``.""" + return _takes_ctx_rows_start_cls(type(draft_model)) + + +def dflash_ctx_rows_kwargs(draft_model, ctx_rows_start: Optional[int]) -> dict: + """The ``ctx_rows_start`` keyword (from ``prepare_1st_drafter_inputs``) for a draft model whose ``dflash_forward`` + takes it; no keyword for another draft model or without a start.""" + if ctx_rows_start is None or not _takes_ctx_rows_start(draft_model): + return {} + return {"ctx_rows_start": ctx_rows_start} + + +def capture_view( + spec_metadata: Optional[SpecMetadata], layer_id: int, num_tokens: int +) -> Optional[torch.Tensor]: + """``layer_id``'s slot of the DFlash capture buffer for ``num_tokens`` rows (a strided view a kernel can write the + tap into), or None when ``spec_metadata`` is not DFlash's or the layer is not captured.""" + if ( + not isinstance(spec_metadata, DFlashSpecMetadata) + or spec_metadata.captured_hidden_states is None + ): + return None + i = spec_metadata._layer_to_idx.get(layer_id) + if i is None: + return None + return spec_metadata.captured_hidden_states[ + :num_tokens, i * spec_metadata.hidden_size : (i + 1) * spec_metadata.hidden_size + ] + + +class KimiK3DFlashWorker(DFlashWorker): + # Set by a Kimi K3 target running its decode kernels: decode steps then run the acceptance and the drafter's + # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context + # K/V with trtllm::k3_ctx_kv. + k3_decode = False + + def _mask_token_id(self, draft_model) -> int: + """The drafter's mask token id, resolved once.""" + if self._resolved_mask_token_id is None: + if ( + hasattr(self.spec_config, "mask_token_id") + and self.spec_config.mask_token_id is not None + ): + self._resolved_mask_token_id = self.spec_config.mask_token_id + elif hasattr(draft_model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.mask_token_id + elif hasattr(draft_model.model, "mask_token_id"): + self._resolved_mask_token_id = draft_model.model.mask_token_id + else: + raise ValueError( + "DFlash requires mask_token_id to be set. Please set it in DFlashDecodingConfig " + "or ensure the draft model config has 'dflash_config.mask_token_id' or 'mask_token_id'." + ) + return self._resolved_mask_token_id + + @staticmethod + def _trained_mask_embedding(draft_model, mask_token_id: int) -> Optional[torch.Tensor]: + """The drafter's own trained mask row, if it kept one for the mask id in use.""" + if mask_token_id != getattr(draft_model, "mask_token_id", None): + return None + return getattr(draft_model, "mask_token_embedding", None) + + def _k3_accept_applies( + self, logits, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Whether this step's acceptance and the drafter's inputs run as ``trtllm::k3_spec_accept``. + + That mode: a decode step (no context requests, at most 8 gen requests) with greedy strict acceptance (no + rejection sampling, penalties or guided decoding, the base draft-token and logits layouts), fp32 target + logits of a vocabulary the kernel splits, the V2 Mamba manager's KDA replay record, the draft pool's block + table, and an unsharded bf16 draft embedding. Any other step keeps the torch path. + """ + if not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return False + K = spec_metadata.runtime_draft_len + if ( + logits.dim() != 2 + or logits.dtype != torch.float32 + or logits.shape[0] != num_gens * (K + 1) + ): + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + embed = draft_model.draft_model_full.model.embed_tokens + return accept_op.supports( + logits.shape[1], num_gens, self._compute_block_size, K, embed.weight.shape[1] + ) + + def _k3_accept_step_applies( + self, attn_metadata, spec_metadata, draft_model, num_contexts: int, num_gens: int + ) -> bool: + """Everything ``_k3_accept_applies`` checks but the target logits.""" + if not self.k3_decode: + return False + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + + if num_contexts != 0 or not 0 < num_gens <= 8 or self.guided_decoder is not None: + return False + if ( + not spec_metadata.is_all_greedy_sample + or self._can_use_rejection_sampling(spec_metadata) + or getattr(spec_metadata, "enable_penalty", False) + or type(self)._reshape_draft_tokens_for_accept + is not SpecWorkerBase._reshape_draft_tokens_for_accept + or type(self)._reshape_logits_for_accept + is not SpecWorkerBase._reshape_logits_for_accept + ): + return False + K = spec_metadata.runtime_draft_len + draft_tokens = spec_metadata.draft_tokens + if ( + K <= 0 + or draft_tokens is None + or draft_tokens.dtype != torch.int32 + or draft_tokens.numel() != num_gens * K + ): + return False + mgr = attn_metadata.kv_cache_manager + mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) + if ( + not isinstance(mgr, MambaHybridCacheManagerV2) + or type(mgr).update_mamba_states is not MambaHybridCacheManagerV2.update_mamba_states + or not mgr.use_kda_replay_update + or getattr(mgr, "prev_num_accepted_tokens", None) is None + or getattr(mgr, "_dummy_request_mask", None) is None + or mamba_metadata is None + or mamba_metadata.state_indices.dtype != torch.int32 + ): + return False + if ( + self._ctx_block_tables is None + or getattr(attn_metadata, "draft_kv_cache_block_offsets", None) is None + or getattr(attn_metadata, "kv_lens_cuda", None) is None + ): + return False + embed = draft_model.draft_model_full.model.embed_tokens + return getattr(embed, "tp_size", 1) == 1 and embed.weight.dtype == torch.bfloat16 + + def target_logits( + self, hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) -> torch.Tensor: + """This rank's bf16 vocabulary shard of the target logits, [rows, vocab / TP] (``_k3_head_shard``: the head + GEMM without its all-gather and fp32 cast), for a step whose acceptance runs as the sharded + ``trtllm::k3_spec_accept`` (see ``_k3_logits_shard``); the gathered fp32 logits otherwise. At such a step the + acceptance is their only reader: the one-model spec sampler takes the worker's tokens and rejects requests for + logits or log probabilities, so the shard is also what the step returns as ``"logits"``.""" + self._k3_step_shard = None + shard = self._k3_logits_shard( + hidden_states, lm_head, attn_metadata, spec_metadata, draft_model + ) + if shard is None: + return super().target_logits( + hidden_states, lm_head, logits_processor, attn_metadata, spec_metadata, draft_model + ) + logits = self._k3_head_shard(logits_processor, lm_head, hidden_states) + self._k3_step_shard = (logits, shard) + return logits + + @staticmethod + def _k3_head_shard(logits_processor, lm_head, rows: torch.Tensor) -> torch.Tensor: + """This rank's bf16 vocabulary shard of ``lm_head(rows)``, without the head's all-gather: from the logits + processor's own head kernel where it has one that takes the rows (``lm_head_shard``; the Kimi K3 target's + ``gemm/k3_head_gemv``), so the shard holds the values of the processor's gathered logits; else from the head's + own GEMM.""" + head_shard = getattr(logits_processor, "lm_head_shard", None) + logits = None if head_shard is None else head_shard(rows, lm_head) + if logits is None: + logits = lm_head.apply_linear(rows, lm_head.bias) + return logits + + def _k3_logits_shard(self, hidden_states, lm_head, attn_metadata, spec_metadata, draft_model): + """``(workspace, first column)`` when this step's target logits stay vocabulary-sharded, else None. + + That needs a step ``_k3_accept_step_applies`` takes, plain TP whose all-reduces own an MNNVL workspace for + this mapping, and an unquantized, bias-free, unpadded, evenly split column-parallel bf16 lm_head with shards + the kernel splits. The exchange is collective, so every rank must decide alike: every condition is the + batch's or the configuration's. The workspace is allocated (collectively) and the kernel compiled outside + CUDA-graph capture only, so a graph whose warmup allocated no workspace is captured on the gathered path. + """ + if not self.k3_decode: + return None + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + + mapping = self.mapping + if getattr(self, "_k3_shard_head", None) is not lm_head: + reason = self._k3_logits_shard_reason(lm_head) + if reason is None and MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None: + # Not kept: re-checked on later steps, since the model's MNNVL workspace may not exist yet. + logger.info_once( + "DFlash: the target logits are all-gathered (the TP all-reduces are not MNNVL)", + key="dflash_k3_logits_no_mnnvl", + ) + return None + if mapping is not None and mapping.tp_size > 1: + if reason is None: + logger.info( + f"DFlash: decode-step target logits stay vocabulary-sharded ({lm_head.weight.shape[0]} " + f"columns per rank); trtllm::k3_spec_accept exchanges the row maxima, and a request's " + f"logits post-processors get them gathered" + ) + else: + logger.info_once( + f"DFlash: the target logits are all-gathered ({reason})", key=reason + ) + self._k3_shard_head, self._k3_shard_ok = lm_head, reason is None + if not self._k3_shard_ok: + return None + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + K = spec_metadata.runtime_draft_len + if hidden_states.shape[0] != num_gens * (K + 1) or not self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + return None + embed = draft_model.draft_model_full.model.embed_tokens + columns = lm_head.weight.shape[0] + if not accept_op.supports( + columns, num_gens, self._compute_block_size, K, embed.weight.shape[1], mapping.tp_size + ): + return None + if torch.cuda.is_current_stream_capturing(): + ws = accept_op.existing_workspace(mapping) + if ws is None: + return None + else: + ws = accept_op.workspace(mapping) + return ws, mapping.tp_rank * columns + + def _k3_logits_shard_reason(self, lm_head): + """Why the target logits cannot stay vocabulary-sharded under this configuration and head (None if they can, + given an MNNVL workspace for the mapping, which ``_k3_logits_shard`` checks).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + if mapping is None or mapping.tp_size <= 1: + return "no TP" + if mapping.enable_attention_dp or mapping.has_cp() or mapping.has_pp(): + return "attention DP, CP or PP" + if ( + getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "padding_size", 0) != 0 + or getattr(lm_head, "gather_output_sizes", None) is not None + ): + return "the lm_head is not an evenly split, unpadded column-parallel head" + if ( + lm_head.bias is not None + or lm_head.has_any_quant + or lm_head.weight.dtype != torch.bfloat16 + ): + return "the lm_head has a bias, is quantized or is not bf16" + if not accept_op.supports_columns(lm_head.weight.shape[0], mapping.tp_size): + columns = lm_head.weight.shape[0] + return f"trtllm::k3_spec_accept cannot split {columns} columns over {mapping.tp_size} ranks" + return None + + def _k3_static_arange(self, lo: int, hi: int) -> torch.Tensor: + """``arange(lo, hi)`` (int64), built once outside CUDA-graph capture.""" + cache = self.__dict__.setdefault("_k3_aranges", {}) + t = cache.get((lo, hi)) + if t is None: + t = torch.arange(lo, hi, dtype=torch.long, device="cuda") + if not torch.cuda.is_current_stream_capturing(): + cache[(lo, hi)] = t + return t + + def _k3_mask_row(self, draft_model) -> torch.Tensor: + """The noise block's mask embedding, bf16 [hidden], cached: the drafter's own trained mask row if it kept one + for the mask id in use, else the embedding's row (``dflash_noise_block_embedding`` reads the same row).""" + row = getattr(self, "_k3_mask_embedding", None) + if row is None: + embed = draft_model.draft_model_full.model.embed_tokens + mask_token_id = self._mask_token_id(draft_model) + trained = self._trained_mask_embedding(draft_model, mask_token_id) + if trained is not None: + row = trained.to(embed.weight.dtype).reshape(-1).contiguous() + else: + row = embed.weight[mask_token_id].contiguous() + if torch.cuda.is_current_stream_capturing(): + return row + self._k3_mask_embedding = row + return row + + def _k3_ctx_pool_key(self): + """The identity of the pool the ctx cache is bound to now (KV cache estimation rebinds it to a new one).""" + return ( + tuple(t.data_ptr() for t in self._ctx_kv_buf) if self._ctx_kv_buf is not None else None + ) + + def _k3_ctx_kv_applies(self, draft_model, projected: torch.Tensor) -> bool: + """Whether ``trtllm::k3_ctx_kv`` writes this step's context K/V: the manager-bound paged pool with K and V, + bf16, the fused K/V weight without bias or context input norm, k_norm, NeoX RoPE from flashinfer's fp32 + cache, a context the op can split (up to 64 tokens: B <= 8 requests of K + 1 <= 8, ``ctx_op.pick_split``), one + allocation for every layer's pool. Checked once per token count and + bound pool; the pool view is kept for the bound pool only, so a rebind never leaves the kernel writing the + replaced pool.""" + pool_key = self._k3_ctx_pool_key() + if getattr(self, "_k3_ctx_pool_bound", None) != pool_key: + self._k3_ctx_pool_bound = pool_key + self._k3_ctx_pool = None + self._k3_ctx_kv_ok = {} + cache = self._k3_ctx_kv_ok + n = projected.shape[0] + ok = cache.get(n) + if ok is not None: + return ok + from tensorrt_llm._torch.cute_dsl_kernels.k3_ctx_kv import op as ctx_op + from tensorrt_llm._torch.models import modeling_dflash + + reason = None + if not (self._ctx_paged and self._ctx_block_tables is not None): + reason = "the context pool is not the manager-bound paged pool" + elif projected.dtype != torch.bfloat16 or self._ctx_kv_buf[0].dtype != torch.bfloat16: + reason = "not bf16" + else: + if draft_model._fused_kv_weight is None: + draft_model._build_fused_kv_buffers() + view = ctx_op.pool_view(list(self._ctx_kv_buf)) + if ( + draft_model._fused_kv_bias is not None + or getattr(draft_model, "_input_ln_eps", None) is not None + ): + reason = "K/V bias or context input norm" + elif draft_model._k_norm_stacked is None or not draft_model._is_neox: + reason = "no k_norm or not NeoX RoPE" + elif ( + modeling_dflash._flashinfer_rope is None + or draft_model._get_cos_sin_cache().shape[1] != 64 + ): + reason = "RoPE is not flashinfer's full-rotary fp32 cache" + elif view is None or self._ctx_kv_buf[0].size(1) != 2: + reason = "the layers' pools are not K/V views of one allocation" + elif not ctx_op.pick_split( + draft_model._fused_kv_weight.shape[0], + projected.shape[1], + draft_model._num_kv_heads, + self._draft_tokens_per_req, + n, + projected.device, + ): + reason = f"{n} context tokens (at most 64) or the K/V weight shape" + else: + self._k3_ctx_pool = view + ok = reason is None + if not ok: + logger.info(f"DFlash: context K/V on the Python path ({reason})") + if not torch.cuda.is_current_stream_capturing(): + cache[n] = ok + return ok + + def _k3_ctx_kv( + self, draft_model, projected, ctx_positions, num_accepted, slots, rows + ) -> torch.Tensor: + """``trtllm::k3_ctx_kv``: every drafter layer's context K/V of this step into the paged pool and + ``_ctx_len += num_accepted`` (clamped); returns the context length each request may advertise.""" + if self._k3_ctx_pool is None or self._k3_ctx_pool_bound != self._k3_ctx_pool_key(): + raise RuntimeError("k3_ctx_kv: the context pool was rebound after the path was chosen") + flat, layer_off, page_stride, kv_stride, head_stride = self._k3_ctx_pool + return torch.ops.trtllm.k3_ctx_kv( + projected.contiguous(), + draft_model._fused_kv_weight, + draft_model._k_norm_stacked, + draft_model._get_cos_sin_cache(), + ctx_positions, + num_accepted, + self._ctx_len, + slots, + rows, + self._ctx_block_tables, + self._ctx_block_counts, + flat, + layer_off, + page_stride, + kv_stride, + head_stride, + draft_model._k_norm_eps, + self._max_ctx, + self._ctx_page_size, + self._compute_block_size, + draft_model._num_kv_heads, + ) + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Called with this step's acceptance when ``k3_spec_accept`` produced it (a hook for drafter families).""" + + def _k3_accept(self, logits, attn_metadata, spec_metadata, draft_model, shard=None): + """``trtllm::k3_spec_accept`` for a step ``_k3_accept_applies`` accepted: the acceptance, the block table, + the KDA replay record, kv_lens + 1 and the drafter's inputs in one launch. ``shard``: the exchange + workspace and first column when ``logits`` are this rank's vocabulary shard (``_k3_logits_shard``). + Returns (accepted tokens, num accepted, {rewind, bonus, qpos, cpos, noise}).""" + from tensorrt_llm._torch.cute_dsl_kernels.k3_spec_accept import op as accept_op + + num_gens = attn_metadata.num_seqs + K = spec_metadata.runtime_draft_len + mgr = attn_metadata.kv_cache_manager + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + force = 0.0 if is_warmup else float(self.force_num_accepted_tokens) + mode, _, _ = accept_op.force_mode(force, K) + if mode == accept_op._kernel_module().FORCE_FRAC: + self._ensure_force_accept_rng_state(logits.device) + pool, counter = self._force_accept_rng_pool, self._force_accept_rng_counter + else: + dummy_rng = getattr(self, "_k3_dummy_rng", None) + if dummy_rng is None: + dummy_rng = self._k3_dummy_rng = ( + torch.zeros(4, dtype=torch.float32, device=logits.device), + torch.zeros(2, dtype=torch.int64, device=logits.device), + ) + pool, counter = dummy_rng + embed = draft_model.draft_model_full.model.embed_tokens + accepted, num_acc, rewind, bonus, qpos, cpos, noise = torch.ops.trtllm.k3_spec_accept( + logits, + spec_metadata.draft_tokens.reshape(num_gens, K), + attn_metadata.draft_kv_cache_block_offsets, + self._ctx_pool_idx, + self._ctx_block_divisor, + self._ctx_block_counts, + self._ctx_block_tables, + mgr.prev_num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices, + mgr._dummy_request_mask, + attn_metadata.kv_lens_cuda, + self._batch_to_slot, + self._ctx_len, + self._max_ctx, + embed.weight, + self._k3_mask_row(draft_model), + pool, + counter, + force, + self._compute_block_size, + *self._k3_shard_args(shard), + ) + self._on_acceptance(accepted, num_acc, attn_metadata, spec_metadata) + return ( + accepted, + num_acc, + dict(rewind=rewind, bonus=bonus, qpos=qpos, cpos=cpos, noise=noise), + ) + + @staticmethod + def _k3_shard_args(shard) -> tuple: + """``trtllm::k3_spec_accept``'s exchange arguments for a vocabulary-sharded step (none otherwise).""" + if shard is None: + return () + ws, first_column = shard + return ( + ws["uc"], + ws["mc"], + ws["flags"], + ws["rank"], + ws["slots"], + ws["push_copies"], + first_column, + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + # Set when target_logits kept this step's logits vocabulary-sharded (only k3_spec_accept reads them). + step_shard, self._k3_step_shard = getattr(self, "_k3_step_shard", None), None + shard = None + if step_shard is not None: + if step_shard[0] is not logits: + raise RuntimeError( + "DFlash: the target logits are not the vocabulary shard target_logits returned" + ) + shard = step_shard[1] + + raw_logits = logits + K = spec_metadata.runtime_draft_len + + if K == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + # Lazy init buffers and attach worker reference for prepare() + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + # A decode step of the supported mode runs the acceptance and the drafter's inputs as one kernel + # (trtllm::k3_spec_accept), which also decodes the block table below. + if shard is None: + fused_accept = self._k3_accept_applies( + logits, attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ) + elif self._k3_accept_step_applies( + attn_metadata, spec_metadata, draft_model, num_contexts, num_gens + ): + fused_accept = True + else: + raise RuntimeError( + "DFlash: the target logits are vocabulary-sharded, but the fused acceptance no longer applies" + ) + # Before any store: prefill and decode both address pages through it. + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + if not fused_accept: + self._refresh_ctx_block_tables(attn_metadata, batch_size) + + # Save context lengths so both warmup and a failed forward can roll + # back the in-place _ctx_len updates made during drafting. + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # Never allocate the snapshot while capturing a CUDA graph: the + # clone would live in the graph memory pool, and its replay-time + # writes could alias blocks reused by later captures. Rollback is + # only meaningful for eager/warmup forwards anyway; a failure + # during capture aborts the graph itself, and captured ops do not + # mutate _ctx_len until replay. + self._saved_ctx_len = self._ctx_len.clone() + self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_req_ctx_pos = dict(self._req_ctx_pos) + self._ctx_len_restore_pending = True + + self._execute_guided_decoder_if_present(logits) + + prep = None + if fused_accept: + # Saved before the kernel updates kv_lens (warmup also snapshots kv_lens_cuda here). + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + accepted_tokens, num_accepted_tokens, prep = self._k3_accept( + logits, attn_metadata, spec_metadata, draft_model, shard + ) + else: + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + + # Opt-in acceptance recording (env-gated; eager-mode measurement + # runs only). Skipped for CUDA-graph batches (capture/replay/warmup + # use synthetic requests and forbid the host sync). + if ( + self._accept_stats is not None + and num_gens > 0 + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + ): + self._accept_stats.on_accept( + spec_metadata.request_ids[num_contexts:batch_size], + num_accepted_tokens[num_contexts:batch_size].tolist(), + ) + + if fused_accept: + # The kernel recorded the KDA replay acceptance and advanced kv_lens_cuda. + self._kv_rewind_amount = prep["rewind"] + self._kv_rewind_nc = num_contexts + self._kv_rewind_bs = batch_size + self._kv_rewind_pending = True + attn_metadata.update_for_spec_dec() + else: + # Update GDN/Mamba recurrent states to the accepted token's state. + if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) + self._prepare_kv_for_draft_forward( + attn_metadata, num_accepted_tokens, num_contexts, batch_size + ) + + # Collapse mrope [3, 1, N] to 1D by taking the first (temporal) dimension. + # The draft model uses standard 1D RoPE, so only scalar positions are needed. + if position_ids.ndim == 3: + position_ids = position_ids[0, 0] + else: + position_ids = position_ids.squeeze(0) + + # Get total tokens processed by target model (for hidden state extraction) + total_target_tokens = input_ids.shape[0] + + # Capture prefill (context) hidden states for future gen steps. + # This gives the draft model the full prompt context, not just gen tokens. + if num_contexts > 0: + self._store_prefill_context( + draft_model, spec_metadata, attn_metadata, position_ids, total_target_tokens + ) + # Rebuild batch_to_slot after prefill assigns new slots + if self._ctx_buf_inited and spec_metadata.request_ids: + self._rebuild_batch_to_slot(spec_metadata.request_ids) + + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=total_target_tokens, + prep=prep, + ) + + if num_gens > 0: + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + hidden_states_out = draft_model.dflash_forward( + noise_embedding=inputs["noise_embedding"], + query_positions=inputs["query_positions"], + num_ctx_per_req=inputs["num_ctx_per_req"], + ctx_k_cache=inputs["ctx_k_cache"], + ctx_v_cache=inputs["ctx_v_cache"], + ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + ctx_kv_cache=inputs["ctx_kv_cache"], + ctx_page_table=inputs["ctx_page_table"], + **dflash_ctx_rows_kwargs(draft_model, inputs["ctx_rows_start"]), + ) + + # Gather K logits per gen request from the block outputs. + # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. + # Which block slots carry them is a drafter-family convention, + # resolved through _draft_slot_ids. + block_size = self._compute_block_size + gen_hidden_states = self._draft_block_hidden_states( + draft_model, hidden_states_out, num_gens, block_size, K + ) + gen_logits = self._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + vocab_size = gen_logits.shape[-1] + gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + + gen_logits = self._refine_block_logits( + draft_model, gen_logits, inputs, spec_metadata + ) + + # DFlash 2: replace the independent per-position picks with one + # coherent path through the block (absent for plain DFlash). + if getattr(draft_model, "has_candidate_selector", False): + gen_logits = self._apply_dflash2_selector( + draft_model, + gen_logits, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + spec_metadata, + ) + vocab_size = gen_logits.shape[-1] + + gen_draft_tokens = self.sample_draft_tokens( + gen_logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + ) + + # Opt-in confidence-calibration recording (env-gated). The + # confidence values come from an optional provider on the + # recorder (None here -> no calibration rows collected; the + # DSpark confidence-scheduled verification MR supplies the + # real provider, see accept_stats.py). Guarded off for + # CUDA-graph batches and d2t vocab-mapped drafters. + if ( + self._accept_stats is not None + and self._accept_stats.confidence_provider is not None + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + and self._d2t is None + ): + self._accept_stats.record_draft_confidence( + spec_metadata.request_ids[num_contexts:batch_size], + draft_model, + gen_hidden_states.reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + gen_draft_tokens, + ) + + else: + gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") + + # Context requests are not drafted by the block worker (zero placeholder + # token); fill their draft-prob slot rows with a one-hot placeholder so + # they are a legal distribution when they become gen requests next iter. + gen_vocab = vocab_size if num_gens > 0 else None + self.write_context_onehot_draft_probs(spec_metadata, num_contexts, num_gens, K, gen_vocab) + + if num_contexts > 0 and num_gens > 0: + ctx_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + next_draft_tokens = torch.cat([ctx_draft_tokens, gen_draft_tokens], dim=0) + elif num_contexts > 0: + next_draft_tokens = torch.zeros((num_contexts, K), dtype=torch.int32, device="cuda") + else: + next_draft_tokens = gen_draft_tokens + + self._restore_attn_metadata_from_spec_dec(attn_metadata) + self._apply_kv_rewind_after_draft(attn_metadata, spec_metadata) + + self._rollback_guided_decoder_after_verify(num_accepted_tokens) + + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + # Restore context lengths after warmup; real runs keep the updates. + if is_warmup: + self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_host() + self._ctx_len_restore_pending = False + + outputs = { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + if shard is not None: + # "logits" is this rank's vocabulary shard: the engine gathers it for any logits post-processor. + outputs["logits_vocab_shard"] = True + return outputs + + def _draft_block_hidden_states( + self, + draft_model, + hidden_states_out: torch.Tensor, + num_gens: int, + block_size: int, + num_draft_tokens: int, + ) -> torch.Tensor: + """The block-output rows that produce the K draft logits per gen request (``_draft_slot_ids``). + + The slot ids depend only on the shapes, so they are built once per shape, outside CUDA-graph + capture, and cached. When they are one contiguous run of rows (a single request, or the + shift_label convention with a block of K slots, whose requests' slots follow each other) the + rows are returned as a view, without a gather; otherwise they are gathered. + """ + rows = hidden_states_out.shape[0] + key = (num_gens, block_size, num_draft_tokens, rows) + cache = self.__dict__.setdefault("_draft_block_rows", {}) + entry = cache.get(key) + if entry is None: + ids = self._draft_slot_ids(draft_model, num_gens, block_size, num_draft_tokens) + # Shields only the last request: at block_size == K with + # shift_label off, slots run 1..K, so every request reads the + # next one's slot 0 and the last overruns. Degrades, never raises. + ids = ids.clamp(max=rows - 1) + if torch.cuda.is_current_stream_capturing(): + return hidden_states_out[ids] + host = ids.tolist() + if host and host == list(range(host[0], host[0] + len(host))): + entry = host[0] + else: + entry = ids + cache[key] = entry + if isinstance(entry, int): + return hidden_states_out[entry : entry + num_gens * num_draft_tokens] + return hidden_states_out[entry] + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata, + ) -> torch.Tensor: + """fp32 logits of the block positions over the full vocabulary (the TP lm_head all-gathers its + shards). A drafter family whose refinement and draft sampler reduce across ranks themselves may + return this rank's vocabulary shard instead.""" + return draft_model.logits_processor( + gen_hidden_states, draft_model.lm_head, attn_metadata, True + ) + + def _ctx_rows_start(self, num_contexts: int, num_gens: int) -> Optional[int]: + """The first row of the gen requests' page-table rows where those rows are one contiguous run, else None. + + The manager's block table is keyed by batch position, so the gen requests' rows are [num_contexts, + num_contexts + num_gens): a drafter can view them instead of gathering ``ctx_cache_batch_idx``'s rows. The + private arena's table is keyed by slot, which need not be contiguous. + """ + if num_gens > 0 and self._ctx_block_tables is not None: + return num_contexts + return None + + def prepare_1st_drafter_inputs( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, + hidden_states: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + attn_metadata: AttentionMetadata, + spec_metadata: DFlashSpecMetadata, + draft_model: nn.Module, + total_target_tokens: int = 0, + prep: Optional[dict] = None, + ): + """Prepare inputs for DFlash's draft forward. + + ``prep``: ``trtllm::k3_spec_accept``'s bonus tokens, positions and noise embedding for a decode step it + took (see ``_k3_accept``); computed here otherwise. + + For gen requests, builds: + - noise_embedding: token embeddings for [accepted + mask] tokens + - query_positions: position IDs for the query tokens + - num_ctx_per_req: per-request context length in the pool + - ctx_k_cache / ctx_v_cache / ctx_cache_batch_idx: slot-indexed + views of the persistent per-layer K/V pool. + - ctx_rows_start: where ctx_cache_batch_idx is the contiguous rows + [ctx_rows_start, ctx_rows_start + num_gens) of ctx_page_table, their + start (``_ctx_rows_start``), else None. + """ + num_contexts = attn_metadata.num_contexts + batch_size = attn_metadata.num_seqs + num_gens = batch_size - num_contexts + + mask_token_id = self._mask_token_id(draft_model) + + # Get the embed_tokens layer from the draft model + embed_tokens = draft_model.draft_model_full.model.embed_tokens + hidden_dim = ( + spec_metadata.hidden_size if spec_metadata.hidden_size > 0 else hidden_states.shape[-1] + ) + + if num_gens > 0: + gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] + gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] + + total_tokens_per_req = self._draft_tokens_per_req # K+1 + K = spec_metadata.runtime_draft_len + + # Get captured multi-layer hidden states from spec_metadata + captured_hs = spec_metadata.get_hidden_states(total_target_tokens) + has_target_features = ( + captured_hs is not None + and hasattr(draft_model, "fc") + and hasattr(draft_model, "hidden_norm") + ) + + # Use cached block_size (resolved once on first call) + block_size = self._compute_block_size + query_tokens_per_req = block_size + + # Get slots for gen requests from pre-computed mapping + slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] + K_plus_1 = K + 1 + + if prep is not None: + gen_rows_out = self._k3_static_arange(num_contexts, num_contexts + num_gens) + offsets_kp1 = self._k3_static_arange(0, K_plus_1) + bonus = prep["bonus"] + query_position_ids = prep["qpos"] + ctx_position_ids = prep["cpos"] + noise_embed_2d = prep["noise"] + else: + gen_rows_out = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + + bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) + bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() + + ctx_len_gen = self._ctx_len[slots] + j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") + offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") + + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, + bonus, + mask_token_id, + query_tokens_per_req, + self._trained_mask_embedding(draft_model, mask_token_id), + ) + + # Accumulate new accepted features into context buffers + fused_ctx_kv = False + if has_target_features: + gen_start = attn_metadata.num_ctx_tokens + # Target now processes exactly K+1 tokens per gen req, so the + # captured slice is already the full set we need to project. + gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] + gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) + projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) + fused_ctx_kv = prep is not None and self._k3_ctx_kv_applies( + draft_model, projected_to_store + ) + if fused_ctx_kv: + num_ctx_fused = self._k3_ctx_kv( + draft_model, + projected_to_store, + ctx_position_ids, + gen_num_accepted, + slots, + gen_rows_out, + ) + # k3_ctx_kv adds the accepted tokens in the kernel; empty the dummy slot after + # it, as _advance_ctx_len does after the add on the torch path. + self._ctx_len[self._dummy_slot].zero_() + if has_target_features and not fused_ctx_kv: + gen_num_accepted_long = gen_num_accepted.long() + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) + if self._ctx_block_tables is not None: + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) + else: + col_idx = col_idx.clamp(max=self._max_ctx - 1) + + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) + col_flat = col_idx.reshape(-1) + proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] + pos_flat = ctx_position_ids.long().reshape(-1) + mask_1d = write_mask.reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = ( + self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + ) + k_new, v_new = draft_model.precompute_context_kv( + proj_flat.to(cache_dtype), pos_flat + ) + mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + slot_long = slot_flat.long() + col_long = col_flat.long() + if self._ctx_paged: + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) + else: # VANILLA DFlash backend (FlashAttention) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + self._advance_ctx_len(slots, gen_num_accepted_long) + + if fused_ctx_kv: + num_ctx_per_req_t = num_ctx_fused + else: + num_ctx_per_req_t = self._ctx_len[slots] + if not fused_ctx_kv and self._ctx_block_tables is not None: + # The write above clamps columns to the request's allocation, so + # a context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size + ) + num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + noise_embedding = noise_embed_2d + query_positions = query_position_ids.long() + + # Update seq_lens for gen requests to K+1 (the number of tokens + # the target forward actually processed). + attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = ( + total_tokens_per_req + ) + attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req + else: + noise_embedding = hidden_states.new_empty(0, 0, hidden_dim) + query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") + num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") + slots = torch.empty(0, dtype=torch.long, device="cuda") + gen_rows_out = slots + bonus = torch.empty(0, dtype=torch.long, device="cuda") + + return { + "noise_embedding": noise_embedding, + "query_positions": query_positions, + "num_ctx_per_req": num_ctx_per_req_t, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": bonus, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), + "ctx_rows_start": self._ctx_rows_start(num_contexts, num_gens), + } + + +class KimiK3DSparkWorker(KimiK3DFlashWorker, DSparkWorker): + # Whether this step's block logits are the Kimi K3 decode path's vocabulary shard (see ``_draft_block_logits``). + _k3_sharded_block_logits = False + + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Block logits the Kimi K3 decode path kept vocab-sharded (see ``_draft_block_logits``) run the + whole chain as ``trtllm::k3_markov``, which also returns the greedy draft tokens and + next_new_tokens. Other TP-sharded logits slice markov_w2's rows to this rank's contiguous + shard and chain through the TP-aware global argmax. + """ + # The d2t guard lives in set_draft_model: it is model-static, so raising + # it here would surface a load-time config error per decode step. + # Unlike the d2t guard this one cannot move to set_draft_model: it + # keys on the runtime logits width, and reproducing that at init would + # duplicate the draft head's sharding rules. A standalone drafter + # borrows the target lm_head, whose gather_output defaults to True, so + # the logits normally arrive full-vocab and this branch is skipped. + self._k3_markov = None + k3_sharded, self._k3_sharded_block_logits = self._k3_sharded_block_logits, False + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + if k3_sharded: + return self._k3_markov_chain( + draft_model, gen_logits, first_prev_tokens, vocab_slice + ) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) + + def _keep_draft_logits_sharded(self, draft_model, spec_metadata, num_gens: int) -> bool: + """Keep the draft logits vocab-sharded: ``trtllm::k3_markov`` reduces across the ranks itself. + + The chain's global argmaxes are the greedy draft tokens, so the lm_head's all-gather of the + full logits and the draft sampler's gather are both skipped. Needs a Kimi K3 target running its + decode kernels (``k3_decode``), plain TP, greedy drafting (rejection sampling reads + full-vocabulary probabilities), an unquantized bias-free column-parallel bf16 lm_head whose + shards tile the Markov vocabulary, bf16 Markov weights of the kernel's rank, MNNVL (the model's + all-reduces own an MNNVL workspace for this mapping), and a block, shard and batch of + ``num_gens`` requests the kernel can split. + """ + if not self.k3_decode: + return False + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce + from tensorrt_llm._torch.modules.linear import TensorParallelMode + + mapping = self.mapping + lm_head = getattr(draft_model, "lm_head", None) + if ( + lm_head is None + or not getattr(draft_model, "has_markov_head", False) + or mapping is None + or mapping.tp_size <= 1 + or mapping.enable_attention_dp + or spec_metadata.wants_advanced_draft_sampling + or getattr(lm_head, "tp_mode", None) != TensorParallelMode.COLUMN + or not getattr(lm_head, "gather_output", False) + or getattr(lm_head, "bias", None) is not None + or lm_head.weight.dtype != torch.bfloat16 + or draft_model.markov_w1.dtype != torch.bfloat16 + or draft_model.markov_w2.dtype != torch.bfloat16 + or lm_head.weight.shape[0] * mapping.tp_size != draft_model.markov_w2.shape[0] + or MNNVLAllReduce.allreduce_mnnvl_workspaces.get(mapping) is None + ): + return False + shard = lm_head.weight.shape[0] + block = spec_metadata.runtime_draft_len + rank = k3_markov_op._kernel_module().MARKOV_RANK + if ( + not 0 < block <= k3_markov_op.WORKSPACE_MAX_BLOCK + or draft_model.markov_w1.dim() != 2 + or draft_model.markov_w1.shape[1] != rank + or tuple(draft_model.markov_w2.shape[1:]) != (rank,) + ): + return False + if k3_markov_op.pick_grid(shard, block, num_gens) == 0: + logger.warning_once( + f"DSpark Markov head: trtllm::k3_markov cannot split a {shard}-row vocab shard for " + f"{num_gens} requests; the draft logits are all-gathered and the chain runs unfused.", + key=f"dspark_k3_markov_shard_{num_gens}", + ) + return False + return True + + def _draft_block_logits( + self, + draft_model, + gen_hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + ) -> torch.Tensor: + """This rank's bf16 shard of the block logits when ``k3_markov`` takes them (it converts them to fp32 + exactly, as ``.float()`` would), from the drafter's logits processor's head kernel where it takes the rows, + else the head's own GEMM (``_k3_head_shard``); otherwise the base class' fp32 logits.""" + num_gens = gen_hidden_states.shape[0] // max(spec_metadata.runtime_draft_len, 1) + self._k3_sharded_block_logits = self._keep_draft_logits_sharded( + draft_model, spec_metadata, num_gens + ) + if self._k3_sharded_block_logits: + return self._k3_head_shard( + draft_model.logits_processor, draft_model.lm_head, gen_hidden_states + ) + return super()._draft_block_logits( + draft_model, gen_hidden_states, attn_metadata, spec_metadata + ) + + def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): + """The base acceptance; under ``k3_decode`` its outputs are kept for ``k3_markov``'s next_new_tokens.""" + accepted_tokens, num_accepted_tokens = super().sample_and_accept_draft_tokens( + logits, attn_metadata, spec_metadata + ) + if self.k3_decode: + self._on_acceptance(accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata) + return accepted_tokens, num_accepted_tokens + + def _on_acceptance( + self, accepted_tokens, num_accepted_tokens, attn_metadata, spec_metadata + ) -> None: + """Keeps this step's acceptance for ``k3_markov``'s next_new_tokens (and the metadata whose KV lengths it + rewinds).""" + self._k3_acceptance = ( + accepted_tokens, num_accepted_tokens, attn_metadata.num_contexts, spec_metadata, attn_metadata + ) # fmt: skip + + def _k3_markov_chain( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + vocab_slice: slice, + ) -> torch.Tensor: + """The Markov chain as ``trtllm::k3_markov`` on this rank's shard of the block logits. + + Returns the corrected logits (fp32); keeps the kernel's greedy tokens for + ``sample_draft_tokens`` and its next_new_tokens (built from this step's acceptance) for + ``_prepare_next_new_tokens``. + """ + from tensorrt_llm._torch.cute_dsl_kernels.k3_markov import op as k3_markov_op + + acceptance = getattr(self, "_k3_acceptance", None) + if acceptance is None: + raise RuntimeError("DSpark Markov head: k3_markov needs this step's acceptance first") + accepted_tokens, num_accepted_tokens, num_contexts, spec_metadata, attn_metadata = ( + acceptance + ) + num_gens = gen_logits.shape[0] + # The draft forward's pending KV-length rewind (see _apply_kv_rewind_after_draft) runs in the kernel: every + # reader of kv_lens_cuda in the draft forward precedes it. In warmup the restore of kv_lens_cuda that + # follows overwrites it, as it would the Python rewind. + rewind = getattr(self, "_kv_rewind_amount", None) + fold = ( + getattr(self, "_kv_rewind_pending", False) + and rewind is not None + and getattr(attn_metadata, "kv_lens_cuda", None) is not None + and rewind.dtype == torch.int32 + and self._kv_rewind_bs - self._kv_rewind_nc == rewind.numel() <= num_gens + ) + corrected, tokens, next_new = k3_markov_op.markov_chain( + self.mapping, + gen_logits.contiguous(), + first_prev_tokens.long(), + draft_model.markov_w1, + draft_model.markov_w2[vocab_slice], + vocab_slice.start, + accepted_tokens, + num_accepted_tokens[num_contexts : num_contexts + num_gens], + spec_metadata.batch_indices_cuda[num_contexts : num_contexts + num_gens], + kv_lens=attn_metadata.kv_lens_cuda if fold else None, + rewind=rewind if fold else None, + rewind_first=self._kv_rewind_nc if fold else 0, + ) + if fold: + self._kv_rewind_amount = None + self._kv_rewind_pending = False + self._k3_markov = (corrected, tokens, next_new, accepted_tokens, num_accepted_tokens) + return corrected + + def sample_draft_tokens( + self, + logits, + spec_metadata, + batch_size, + *, + num_contexts=0, + draft_step=None, + mapping_lm_head_tp=None, + ): + """Greedy block drafts straight from ``k3_markov``. + + When ``logits`` are the corrected logits the kernel just returned, its per-position global + argmax (first maximum, lowest vocabulary index among equal values) is the token the + TP-gathered greedy sampler would pick. Anything else goes to the base sampler. + """ + chain = getattr(self, "_k3_markov", None) + if ( + chain is not None + and chain[0] is logits + and mapping_lm_head_tp is None + and not spec_metadata.wants_advanced_draft_sampling + ): + self._k3_markov_next = (chain[1], chain[2], chain[3], chain[4]) + return chain[1] + self._k3_markov_next = None + return super().sample_draft_tokens( + logits, + spec_metadata, + batch_size, + num_contexts=num_contexts, + draft_step=draft_step, + mapping_lm_head_tp=mapping_lm_head_tp, + ) + + def _prepare_next_new_tokens( + self, + accepted_tokens, + next_draft_tokens, + batch_indices_cuda, + batch_size, + num_accepted_tokens, + ): + """``k3_markov``'s next_new_tokens when the drafts are its tokens for the whole batch (no context + requests); otherwise the base assembly.""" + chain = getattr(self, "_k3_markov_next", None) + # The step ends here: release its acceptance (which holds the step's attention and spec metadata) and the + # kernel's outputs, so they do not keep a captured graph's pool or metadata alive after the step. + self._k3_markov_next = None + self._k3_acceptance = None + self._k3_markov = None + if ( + chain is not None + and chain[0] is next_draft_tokens + and chain[2] is accepted_tokens + and chain[3] is num_accepted_tokens + and next_draft_tokens.shape[0] == batch_size + ): + return chain[1] + return super()._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, batch_indices_cuda, batch_size, num_accepted_tokens + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/weights.py new file mode 100644 index 000000000000..915a6df08011 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/weights.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight loading for the Kimi K3 text targets. + +The checkpoint is the vision-language wrapper's. `language_model.*` holds the language model; `vision_tower.*` and +`mm_projector.*` hold the vision tower and its projector. This target is text only, so the keys split three ways: + +* `language_model.*` goes to the built-in text model's loader with the prefix stripped. That loader streams the + routed experts one at a time, keeps this rank's slice (a quarter of the width of a quarter of the experts under + moe_tp 4 x moe_ep 4), and checks its own key coverage. +* The vision tower and the projector are a predicted non-load: listed here, never read. +* Any other key fails the load, naming it, rather than being dropped. +""" + +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.models.modeling_utils import filter_weights + +_LANG_PREFIX = "language_model." + +# The vision tower's and the projector's key families: in the checkpoint, not loaded by this text-only target. +PREDICTED_NON_LOAD = ("vision_tower.", "mm_projector.") + + +def load(model, weights) -> None: + """Load the language model's weights into `model`, the target shell, and check every other key is predicted.""" + unknown = sorted( + k + for k in weights.keys() + if not k.startswith(_LANG_PREFIX) and not k.startswith(PREDICTED_NON_LOAD) + ) + assert not unknown, ( + f"{len(unknown)} checkpoint key(s) are neither language-model weights nor a predicted non-load, " + f"first {unknown[:5]}" + ) + lm_weights = ConsumableWeightsDict(filter_weights(_LANG_PREFIX[:-1], weights)) + assert len(lm_weights), f"the checkpoint has no {_LANG_PREFIX}* keys" + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = _LANG_PREFIX + KimiLinearForCausalLM.load_weights(model, lm_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py index 107683470e2d..d98271c207ac 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/modeling.py @@ -5,8 +5,8 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..modeling import K3TargetSpec, ModelingV2KimiK3 -from ..modeling import _check_construction as _check_target_construction +from ..kimi_k3_nvfp4.implementation import K3TargetSpec, ModelingV2KimiK3 +from ..kimi_k3_nvfp4.implementation import _check_construction as _check_target_construction _TARGET = K3TargetSpec(sm=(10, 3), quant_algo=QuantAlgo.NVFP4, moe_tp=4, moe_ep=4) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py index 24233d194709..723db3326e18 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_103__tp16_moetp4ep4/weights.py @@ -2,5 +2,5 @@ # SPDX-License-Identifier: Apache-2.0 """The shared Kimi K3 text checkpoint loader.""" -from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD -from ..weights import load as load +from ..kimi_k3_nvfp4.weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..kimi_k3_nvfp4.weights import load as load diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py index 55d2c599ebd3..d1b9d9269d56 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/modeling.py @@ -5,8 +5,8 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..modeling import K3TargetSpec, ModelingV2KimiK3 -from ..modeling import _check_construction as _check_target_construction +from ..kimi_k3_nvfp4.implementation import K3TargetSpec, ModelingV2KimiK3 +from ..kimi_k3_nvfp4.implementation import _check_construction as _check_target_construction _TARGET = K3TargetSpec(sm=(10, 7), quant_algo=QuantAlgo.NVFP4, moe_tp=4, moe_ep=4) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py index 24233d194709..723db3326e18 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__tp16_moetp4ep4/weights.py @@ -2,5 +2,5 @@ # SPDX-License-Identifier: Apache-2.0 """The shared Kimi K3 text checkpoint loader.""" -from ..weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD -from ..weights import load as load +from ..kimi_k3_nvfp4.weights import PREDICTED_NON_LOAD as PREDICTED_NON_LOAD +from ..kimi_k3_nvfp4.weights import load as load diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index c794e45d5903..783e5e0354ed 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -503,6 +503,7 @@ l0_b200: - unittest/_torch/modeling_v2/ssm - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py new file mode 100644 index 000000000000..d42ced111f03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Numerical and paged-cache parity for NVFP4's FP8 fused attention kernels.""" + +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + + +def _ops(): + """Load the standalone kernels without requiring an initialized inference engine.""" + name = "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4.attention" + if name not in sys.modules: + path = Path(__file__).resolve().parents[5] / ( + "tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/" + "kimi_k3_nvfp4/attention.py" + ) + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return sys.modules[name] + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _rand(*shape, scale=1.0): + return (torch.randn(shape, device="cuda") * scale).bfloat16() + + +def _quantize(x): + groups = x.float().unflatten(-1, (-1, 128)) + scale = groups.abs().amax(-1, keepdim=True).clamp_min(4.48e-10) / 448 + codes = (groups / scale).to(torch.float8_e4m3fn) + return (codes.float() * scale).flatten(-2) + + +@pytest.mark.parametrize("batch", [1, 2, 4, 8]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +def test_fp8_kda_fused_decode(batch, state_dtype): + torch.manual_seed(413) + h, d, slots_count = 6, 768, batch + 3 + slots = torch.randperm(slots_count, device="cuda")[:batch].int() + qkvg = _rand(batch, 4 * d, scale=0.2) + bfa = _rand(batch, 136, scale=0.2) + fb = _rand(d, 128, scale=0.05) + weights = [_rand(d, 4, scale=0.3).float() for _ in range(3)] + a_log = torch.zeros(h, device="cuda") + dt = torch.zeros(d, device="cuda") + norm = torch.ones(128, device="cuda") + conv = _rand(slots_count, 3 * d, 3, scale=0.2) + # The live pool has dense rows and padding between slots. + state = torch.randn(slots_count, h * 16384 + 16, device="cuda", dtype=state_dtype)[ + :, : h * 16384 + ].view(slots_count, h, 128, 128) + expected_conv, expected_state = conv.clone(), state.clone() + initial_conv, initial_state = conv.clone(), state.clone() + output = [] + for row, slot in enumerate(slots.tolist()): + parts = [] + for section in range(3): + raw = qkvg[row, section * d : (section + 1) * d].float() + old = expected_conv[slot, section * d : (section + 1) * d].float() + value = (old * weights[section][:, :3]).sum(-1) + raw * weights[section][:, 3] + parts.append(F.silu(value).view(h, 128)) + expected_conv[slot, section * d : (section + 1) * d] = torch.cat( + (old[:, 1:], raw[:, None]), -1 + ) + q, k, v = parts + raw_q = parts[0] + q = raw_q * torch.rsqrt(raw_q.square().sum(-1, keepdim=True) + 1.0e-6) / 128**0.5 + k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1.0e-6) + gate = F.linear(bfa[row, :128].float(), fb.float()).bfloat16().float().view(h, 128) + decay = torch.exp(-5 * torch.sigmoid(gate)) + s = state[slot].float() * decay[:, None, :] + beta = bfa[row, 128 : 128 + h].float().sigmoid() + innovation = (v - (s * k[:, None, :]).sum(-1)) * beta[:, None] + s += innovation[:, :, None] * k[:, None, :] + expected_state[slot] = s + o = (s * q[:, None, :]).sum(-1) + o *= torch.rsqrt(o.square().mean(-1, keepdim=True) + 1.0e-6) + o *= qkvg[row, 3 * d :].float().view(h, 128).sigmoid() + output.append(o.bfloat16()) + actual = _ops().kda_decode( + qkvg, bfa, fb, *weights, a_log, dt, norm, conv, state, slots, -5.0, 128**-0.5, 1.0e-6 + ) + torch.testing.assert_close(actual, torch.stack(output), atol=0.008, rtol=0.01) + torch.testing.assert_close(conv, expected_conv, atol=0, rtol=0) + torch.testing.assert_close( + state.float(), + expected_state.float(), + atol=0.016 if state_dtype == torch.bfloat16 else 2.0e-5, + rtol=0.01 if state_dtype == torch.bfloat16 else 2.0e-5, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = _ops().kda_decode( + qkvg, + bfa, + fb, + *weights, + a_log, + dt, + norm, + conv, + state, + slots, + -5.0, + 128**-0.5, + 1.0e-6, + ) + conv.copy_(initial_conv) + state.copy_(initial_state) + graph.replay() + torch.testing.assert_close(graph_output, actual, atol=0, rtol=0) + torch.testing.assert_close(conv, expected_conv, atol=0, rtol=0) + + +@pytest.mark.parametrize("batch,tokens", [(1, 1), (2, 1), (4, 1), (1, 4), (4, 8)]) +def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): + torch.manual_seed(422) + torch.backends.cuda.matmul.allow_tf32 = False + h, m = 6, batch * tokens + ag = _rand(m, 2880, scale=0.2) + qa_norm, kv_norm = _rand(1536, scale=0.2), _rand(512, scale=0.2) + qb = _rand(h * 192, 1536, scale=0.05).to(torch.float8_e4m3fn) + qb_scale = torch.rand(9, 12, device="cuda") * 0.2 + 0.1 + kb = _rand(h, 512, 128, scale=0.05).to(torch.float8_e4m3fn) + kb_scale = torch.rand(h, 4, 1, device="cuda") * 0.2 + 0.1 + vb = _rand(h, 128, 512, scale=0.05).to(torch.float8_e4m3fn) + vb_scale = torch.rand(h, 1, 4, device="cuda") * 0.2 + 0.1 + lengths = torch.arange(batch, device="cuda", dtype=torch.int32) * 67 + tokens + 5 + pages_per_request = (int(lengths.max()) + 63) // 64 + pages = ( + torch.randperm(batch * pages_per_request, device="cuda") + .int() + .view(batch, pages_per_request) + ) + page_storage = torch.zeros(batch, 2 * pages_per_request, device="cuda", dtype=torch.int32) + page_storage[:, :pages_per_request] = pages + pages = page_storage[:, :pages_per_request] + page_offset = 2 + pool = _rand(batch * pages_per_request + page_offset, 64, 576, scale=0.1) + expected_pool = pool.clone() + qa = ag[:, :1536].float() + qa = ( + qa * torch.rsqrt(qa.square().mean(-1, keepdim=True) + 1.0e-6) * qa_norm.float() + ).bfloat16() + qb_dequant = qb.float() * qb_scale.repeat_interleave(128, 0).repeat_interleave(128, 1) + q = (_quantize(qa) @ qb_dequant.T).bfloat16().view(m, h, 192) + kb_dequant = kb.float() * kb_scale.repeat_interleave(128, 1).repeat_interleave(128, 2) + absorbed = torch.einsum("mhk,hvk->mhv", _quantize(q[..., :128]), kb_dequant).bfloat16() + expected_q = torch.cat((absorbed, q[..., 128:]), -1).flatten(1) + for row in range(m): + seq, offset = divmod(row, tokens) + pos = int(lengths[seq]) - tokens + offset + page = int(pages[seq, pos // 64]) + page_offset + ckv = ag[row, 1536:2048].float() + ckv *= torch.rsqrt(ckv.square().mean() + 1.0e-6) * kv_norm.float() + expected_pool[page, pos % 64] = torch.cat((ckv, ag[row, 2048:2112].float())).bfloat16() + actual_q = _ops().mla_qkv( + ag, + qa_norm, + qb, + qb_scale, + kb, + kb_scale, + kv_norm, + pool.flatten(), + pages, + lengths, + 576, + page_offset, + tokens, + 1.0e-6, + 1.0e-6, + ) + torch.testing.assert_close(actual_q, expected_q, atol=0.002, rtol=0.02) + torch.testing.assert_close(pool, expected_pool, atol=0.002, rtol=0.02) + vb_dequant = vb.float() * vb_scale.repeat_interleave(128, 1).repeat_interleave(128, 2) + expected_output = [] + for row in range(m): + seq, offset = divmod(row, tokens) + end = int(lengths[seq]) - tokens + offset + 1 + positions = torch.arange(end, device="cuda") + kv = pool[pages[seq, positions // 64].long() + page_offset, positions % 64].float() + logits = actual_q[row].view(h, 576).float() @ kv.T / 192**0.5 + latent = (logits.softmax(-1) @ kv[:, :512]).bfloat16() + projected = torch.einsum("hk,hvk->hv", _quantize(latent), vb_dequant).bfloat16() + gate = ag[row, 2112:].float().sigmoid() + expected_output.append((projected.flatten().float() * gate).bfloat16()) + actual = _ops().mla_attn( + actual_q, + pool.flatten(), + pages, + lengths, + vb, + vb_scale, + ag, + 576, + page_offset, + tokens, + 192**-0.5, + 2112, + ) + torch.testing.assert_close(actual, torch.stack(expected_output), atol=0.002, rtol=0.02) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_q = _ops().mla_qkv( + ag, + qa_norm, + qb, + qb_scale, + kb, + kb_scale, + kv_norm, + pool.flatten(), + pages, + lengths, + 576, + page_offset, + tokens, + 1.0e-6, + 1.0e-6, + ) + graph_output = _ops().mla_attn( + graph_q, + pool.flatten(), + pages, + lengths, + vb, + vb_scale, + ag, + 576, + page_offset, + tokens, + 192**-0.5, + 2112, + ) + graph.replay() + torch.testing.assert_close(graph_q, actual_q, atol=0, rtol=0) + torch.testing.assert_close(graph_output, actual, atol=0, rtol=0) + lengths.add_(1) + graph.replay() + updated_q = _ops().mla_qkv( + ag, + qa_norm, + qb, + qb_scale, + kb, + kb_scale, + kv_norm, + pool.flatten(), + pages, + lengths, + 576, + page_offset, + tokens, + 1.0e-6, + 1.0e-6, + ) + updated_output = _ops().mla_attn( + updated_q, + pool.flatten(), + pages, + lengths, + vb, + vb_scale, + ag, + 576, + page_offset, + tokens, + 192**-0.5, + 2112, + ) + torch.testing.assert_close(graph_q, updated_q, atol=0, rtol=0) + torch.testing.assert_close(graph_output, updated_output, atol=0, rtol=0) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 9d20ff46b988..5854b5fb9e4a 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -130,7 +130,15 @@ def test_every_target_ships_both_products(): def test_targets_only_import_their_declared_family_implementation(): """K3 targets may share their family implementation, never another family's.""" - k3_shared_modules = {"modeling", "weights", "decode_comm", "decode_gemv", "spec_worker"} + k3_shared_modules = { + "modeling", + "weights", + "decode_comm", + "decode_gemv", + "spec_worker", + "kimi_k3_nvfp4.implementation", + "kimi_k3_nvfp4.weights", + } for arch in _ARCHS: for target_dir, name in _routed_targets(arch): for source_file in target_dir.glob("*.py"): @@ -193,7 +201,10 @@ def test_uncertified_generic_calls_name_every_stock_import(): for target_dir, name in _routed_targets(arch): source_file = target_dir / "modeling.py" if target_dir.parent.name == "kimi_k3_vl": - source_file = target_dir.parent / "modeling.py" + if "nvfp4" in target_dir.name: + source_file = target_dir.parent / "kimi_k3_nvfp4" / "implementation.py" + elif "sm_103" in target_dir.name: + source_file = target_dir.parent / "modeling.py" tree = ast.parse(source_file.read_text()) declared = _declared_tuple(tree, "UNCERTIFIED_GENERIC_CALLS") if declared is None: diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py index dd9f989873d0..4dbc56c7b0ac 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""K3 targets share the model lifecycle while preserving each checkpoint's layout.""" +"""NVFP4 targets share one implementation and leave the original MXFP4 targets independent.""" import importlib from types import SimpleNamespace @@ -8,19 +8,47 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4 import ( + implementation as modeling, +) from tensorrt_llm.quantization.mode import QuantAlgo pytestmark = pytest.mark.cpu_only +@pytest.mark.parametrize("moe_tp,moe_ep", [(4, 4), (16, 1)]) +def test_sm100_mxfp4_keeps_its_own_model(moe_tp: int, moe_ep: int) -> None: + module = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl." + f"kimi_k3_mxfp4__sm_100__tp16_moetp{moe_tp}ep{moe_ep}.modeling" + ) + target = getattr(module, f"ModelingV2KimiK3Mxfp4Sm100Tp16Moetp{moe_tp}ep{moe_ep}") + assert target.forward.__module__ == module.__name__ + assert target.forward is not modeling.ModelingV2KimiK3.forward + + +def test_fp8_mla_preserves_raw_query_weights_before_backend_transformation() -> None: + attention = modeling.K3DecodeMLA.__new__(modeling.K3DecodeMLA) + torch.nn.Module.__init__(attention) + codes = torch.ones(128, 128).to(torch.float8_e4m3fn) + scales = torch.full((1, 1), 0.125) + attention.q_b_proj = SimpleNamespace( + weight=codes, + weight_scale=scales, + out_features=128, + in_features=128, + ) + attention.prepare_fp8_decode_weights() + codes.zero_() + scales.zero_() + assert attention._fp8_qb_weight.dtype == torch.float8_e4m3fn + torch.testing.assert_close(attention._fp8_qb_weight.float(), torch.ones(128, 128)) + torch.testing.assert_close(attention._fp8_qb_scale, torch.full((1, 1), 0.125)) + + @pytest.mark.parametrize( "sm,fmt,moe_tp,moe_ep", [ - (100, "mxfp4", 4, 4), - (100, "mxfp4", 16, 1), - (103, "mxfp4", 4, 4), - (103, "mxfp4", 16, 1), (103, "nvfp4", 4, 4), (107, "nvfp4", 4, 4), ], diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index 05fb5eb7ff69..437363f44dc6 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -548,7 +548,7 @@ def test_kimi_k3_nvfp4_declines_excluded_experts(monkeypatch: pytest.MonkeyPatch def test_kimi_k3_nvfp4_constructor_uses_per_layer_quantization(): from types import SimpleNamespace - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4.implementation import ( KimiK3MoERuntime, ) @@ -567,7 +567,7 @@ def test_kimi_k3_nvfp4_constructor_uses_per_layer_quantization(): @pytest.mark.parametrize("prefix", ["language_model.model.", "model.", ""]) def test_kimi_k3_nvfp4_experts_honor_excluded_aliases(prefix: str) -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.modeling import ( + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4.implementation import ( KimiK3MoERuntime, ) from tensorrt_llm.models.modeling_utils import QuantConfig @@ -613,7 +613,9 @@ def test_kimi_k3_nvfp4_keeps_mla_dspark_worker(): @pytest.mark.parametrize("prefix", ["language_model.model.", "model.", ""]) def test_kimi_k3_nvfp4_attention_quantization_preserves_config(prefix: str) -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4 import ( + implementation as target, + ) from tensorrt_llm.models.modeling_utils import QuantConfig projection = QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=128) @@ -633,7 +635,9 @@ def test_kimi_k3_nvfp4_attention_quantization_preserves_config(prefix: str) -> N def test_kimi_k3_nvfp4_attention_honors_exclusions() -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4 import ( + implementation as target, + ) from tensorrt_llm.models.modeling_utils import QuantConfig config = ModelConfig( @@ -657,7 +661,9 @@ def test_kimi_k3_nvfp4_attention_honors_exclusions() -> None: def test_kimi_k3_nvfp4_attention_rejects_unsupported_quantization( algo: QuantAlgo, group_size: int ) -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4 import ( + implementation as target, + ) from tensorrt_llm.models.modeling_utils import QuantConfig config = ModelConfig( @@ -671,7 +677,9 @@ def test_kimi_k3_nvfp4_attention_rejects_unsupported_quantization( def test_kimi_k3_nvfp4_attention_rejects_conflicting_block_sizes() -> None: - from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl import modeling as target + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4 import ( + implementation as target, + ) from tensorrt_llm.models.modeling_utils import QuantConfig config = ModelConfig( From 6b1c3b5d3a09ad7b0ac9ddd279f9e1bf584de08c Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:49:42 +0800 Subject: [PATCH 4/8] Support K3 KV and SSM precision combinations (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 21 +- .../kimi_k3_vl/kimi_k3_nvfp4/attention.py | 119 ++++++++- .../kimi_k3_nvfp4/implementation.py | 99 +++++++- .../backends/fp4_mla/cache_update.py | 2 +- tensorrt_llm/_torch/attention/mla.py | 4 +- .../k3_kda_verify/k3_kda_verify_kernel.py | 13 +- .../cute_dsl_kernels/k3_kda_verify/op.py | 15 +- .../test_lists/test-db/l0_b200.yml | 1 + .../test_modeling_v2_k3_cache_precisions.py | 238 ++++++++++++++++++ .../test_modeling_v2_k3_fp8_fused.py | 25 ++ .../_torch/modeling_v2/ssm/_kda_cells.py | 9 +- .../ssm/test_modeling_v2_k3_bf16_verify.py | 86 +++++++ .../test_modeling_v2_kimi_k3_drift.py | 18 ++ 13 files changed, 625 insertions(+), 25 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py create mode 100644 tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index ac03f4362c8a..e66a73d97530 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -16,9 +16,24 @@ shared by thin SM103 and SM107 target classes. NVFP4 retains checkpoint FP8 projection weights and activations: decode fuses KDA's gate, convolution, recurrence and output normalization, and MLA's query/absorption/cache writes and attention/V-B/output gate. The ordinary FP8 GEMMs remain responsible for -input and output projections. MLA's fused decode currently reads a BF16 KV -cache with 64-token pages; other layouts use the native path. The KDA fused -verify path retains FP32 states and the original per-token draft records. +input and output projections. MLA's fused decode uses a BF16 KV cache with +64-token pages, or composes the fused FP8 projections with the native FP8 +and NVFP4 cache/attention kernels. NVFP4 KV requires 128-token pages and the +V2 cache manager; its high-precision tail and packed scales retain the native +layout. The KDA fused decode and verify accept FP32 or BF16 recurrent-state +storage, with FP32 recurrence arithmetic and FP32 per-token draft records. +These KV and SSM storage choices are independent of the NVFP4 expert weights. +For NVFP4 KV and BF16 recurrent-state storage: + +```yaml +modeling_v2: require +kv_cache_config: + dtype: nvfp4 + tokens_per_block: 128 + mamba_ssm_cache_dtype: bfloat16 + use_kv_cache_manager_v2: true +``` + The NVFP4 expert kernels remain separate from the MXFP4/MXFP8 kernels. Every hardware and quantization target requires its own accuracy gates. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py index 627d026db775..3b203ec37ccf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py @@ -175,6 +175,7 @@ def _mla_qkv( T: tl.constexpr, QA_EPS: tl.constexpr, KV_EPS: tl.constexpr, + STORE_PAGED: tl.constexpr = True, ): head, chunk = tl.program_id(0), tl.program_id(1) rows = chunk * 16 + tl.arange(0, 16) @@ -216,11 +217,16 @@ def _mla_qkv( qpe = tl.gather(q, tl.broadcast_to((128 + pe)[None, :], (16, 64)), 1) tl.store(OUT + base + 512 + pe[None, :], qpe, rows[:, None] < M) if head == 0: - seq = rows // T - length = tl.load(LENGTHS + seq, rows < M, 0) - pos = length - T + rows % T - page = tl.load(PAGES + seq * PAGE_STRIDE + pos // 64, (rows < M) & (pos >= 0), 0) - cache_row = ((page.to(tl.int64) + PAGE_OFFSET) * 64 + pos % 64) * ROW_STRIDE + if STORE_PAGED: + seq = rows // T + length = tl.load(LENGTHS + seq, rows < M, 0) + pos = length - T + rows % T + page = tl.load(PAGES + seq * PAGE_STRIDE + pos // 64, (rows < M) & (pos >= 0), 0) + cache_row = ((page.to(tl.int64) + PAGE_OFFSET) * 64 + pos % 64) * ROW_STRIDE + valid = (rows < M) & (pos >= 0) + else: + cache_row = rows * 576 + valid = rows < M ckv = tl.load( AG + rows[:, None] * AG_STRIDE + 1536 + latent[None, :], rows[:, None] < M, 0 ).to(tl.float32) @@ -229,13 +235,13 @@ def _mla_qkv( tl.store( POOL + cache_row[:, None] + latent[None, :], ckv, - (rows[:, None] < M) & (pos[:, None] >= 0), + valid[:, None], ) kpe = tl.load(AG + rows[:, None] * AG_STRIDE + 2048 + pe[None, :], rows[:, None] < M, 0) tl.store( POOL + cache_row[:, None] + 512 + pe[None, :], kpe, - (rows[:, None] < M) & (pos[:, None] >= 0), + valid[:, None], ) @@ -310,6 +316,105 @@ def _qkv_fake( return ag.new_empty((ag.shape[0], kb.shape[0] * 576)) +@torch.library.custom_op("trtllm::k3_fp8_mla_project_qkv", mutates_args=()) +def mla_project_qkv( + ag: torch.Tensor, + qa_norm: torch.Tensor, + qb: torch.Tensor, + qb_scale: torch.Tensor, + kb: torch.Tensor, + kb_scale: torch.Tensor, + kv_norm: torch.Tensor, + qa_eps: float, + kv_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Produce absorbed Q and normalized latent KV for quantized cache kernels.""" + h = kb.shape[0] + query = ag.new_empty((ag.shape[0], h * 576)) + latent = ag.new_empty((ag.shape[0], 576)) + _mla_qkv[(h, triton.cdiv(ag.shape[0], 16))]( + ag, + qa_norm, + qb, + qb_scale, + kb, + kb_scale, + kv_norm, + latent, + ag, + ag, + query, + ag.shape[0], + h, + ag.stride(0), + qb_scale.stride(0), + *kb.stride(), + *kb_scale.stride(), + 576, + 0, + 0, + 1, + qa_eps, + kv_eps, + STORE_PAGED=False, + num_warps=8, + ) + return query, latent + + +@mla_project_qkv.register_fake +def _project_qkv_fake(ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, qa_eps, kv_eps): + return ag.new_empty((ag.shape[0], kb.shape[0] * 576)), ag.new_empty((ag.shape[0], 576)) + + +@triton.jit +def _mla_output( + LATENT, + VB, + VS, + GATE, + OUT, + H: tl.constexpr, + GATE_STRIDE: tl.constexpr, + GATE_COL: tl.constexpr, +): + row, head = tl.program_id(0), tl.program_id(1) + k = tl.arange(0, 128) + projected_tile = tl.full((16, 128), 0, tl.float32) + for group in tl.static_range(4): + latent = tl.load(LATENT + (row * H + head) * 512 + group * 128 + k).to(tl.float32) + scale = tl.maximum(tl.max(tl.abs(latent), 0) / 448.0, 1.0e-12) + activation = tl.broadcast_to((latent / scale)[None, :], (16, 128)).to(tl.float8e4nv) + weights = tl.load(VB + head * 128 * 512 + k[None, :] * 512 + group * 128 + k[:, None]) + projected_tile += tl.dot(activation, weights) * scale * tl.load(VS + head * 4 + group) + projected = tl.sum(tl.where(tl.arange(0, 16)[:, None] == 0, projected_tile, 0), 0) + projected = projected.to(tl.bfloat16).to(tl.float32) + gate = tl.sigmoid(tl.load(GATE + row * GATE_STRIDE + GATE_COL + head * 128 + k).to(tl.float32)) + tl.store(OUT + (row * H + head) * 128 + k, projected * gate) + + +@torch.library.custom_op("trtllm::k3_fp8_mla_output", mutates_args=()) +def mla_output( + latent: torch.Tensor, + vb: torch.Tensor, + vb_scale: torch.Tensor, + gate: torch.Tensor, + gate_col: int, +) -> torch.Tensor: + """Fuse FP8 V-B absorption and output gating after FP4/FP8 attention.""" + h = vb.shape[0] + output = latent.new_empty((latent.shape[0], h * 128)) + _mla_output[(latent.shape[0], h)]( + latent, vb, vb_scale, gate, output, h, gate.stride(0), gate_col, num_warps=4 + ) + return output + + +@mla_output.register_fake +def _output_fake(latent, vb, vb_scale, gate, gate_col): + return latent.new_empty((latent.shape[0], vb.shape[0] * 128)) + + @triton.jit def _mla_attn( Q, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py index ec82423bd7cd..689534105cc7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py @@ -18,6 +18,7 @@ ) from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType # Registers trtllm::kda_mtp_decode (REQUIRED_TRTLLM_OPS), which the built-in KDA module loads on its first verify. from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 @@ -87,6 +88,8 @@ "k3_fp8_kda_decode", "k3_fp8_mla_qkv", "k3_fp8_mla_attn", + "k3_fp8_mla_project_qkv", + "k3_fp8_mla_output", # The decode path's GEMVs, the dense MLP's activation, the LM head and the embedding (decode_gemv.py). "k3_decode_gemv", "k3_ctm_gemv_wide", @@ -1312,6 +1315,8 @@ def __init__( ) } attention_config.use_cute_dsl_blockscaling_bmm = True + if torch.cuda.get_device_capability() == (10, 7): + attention_config.use_cute_dsl_blockscaling_mm = True attention_config.quant_config = QuantConfig( kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo if model_config.quant_config is not None @@ -2123,6 +2128,13 @@ class K3DecodeKDA(KimiKDALinearAttention): """ def __init__(self, *args, **kwargs) -> None: + config = kwargs.get("model_config") + if config is not None and torch.cuda.get_device_capability() == (10, 7): + config = copy.copy(config) + config._frozen = False + config.use_cute_dsl_blockscaling_mm = True + config._frozen = kwargs["model_config"]._frozen + kwargs["model_config"] = config super().__init__(*args, **kwargs) self.fp8_decode_ready = False self.decode_gemvs: Optional[_decode_gemv.K3DecodeGemvs] = None @@ -2259,8 +2271,6 @@ def forward_verify( f"Kimi K3 KDA layer {self.layer_idx}: the cache manager keeps the per-token verify states, which only " "the decode kernels write, and this layer has no fused projection weight or buffers" ) - if ssm_pool.dtype != torch.float32: - raise ValueError("Fused KDA verification requires an FP32 recurrent-state pool") core = self._k3_verify(x2d, num_steps, layer_cache, ssm_pool, slot_indices) return self._store_core(core, output) @@ -2347,6 +2357,22 @@ def _k3_step_view(self, metadata: AttentionMetadata, step: DecodeStep) -> Option != (1536, 512, 128, 64, 128) ): return None + quantized_cache = bool( + getattr(self.mqa, "has_fp4_kv_cache", False) + or getattr(self.mqa, "has_fp8_kv_cache", False) + ) + if quantized_cache: + if ( + metadata.num_contexts != 0 + or not 0 < metadata.num_generations <= 8 + or not 0 < step.tokens_per_request <= 8 + or step.num_tokens != metadata.num_generations * step.tokens_per_request + or metadata.beam_width != 1 + or getattr(metadata, "is_spec_dec_tree", False) + or getattr(metadata, "is_spec_dec_dynamic_tree", False) + ): + return None + return {"quantized_cache": True} view = k3_mla_decode_view(self.mqa, metadata, step.num_tokens) return None if isinstance(view, str) else view @@ -2365,6 +2391,11 @@ def forward( return self.forward_generic(position_ids, hidden_states, attn_metadata) x = hidden_states[: step.num_tokens] ag = torch.cat((self.kv_a_proj_with_mqa(x), self.g_proj(x)), dim=-1) + if view.get("quantized_cache", False): + core = self._quantized_decode(ag, position_ids, attn_metadata) + if not project_output: + return core + return self._project_output([core], position_ids, attn_metadata, None) fused_q = _fp8_attention.mla_qkv( ag, self.q_a_layernorm.weight, @@ -2400,6 +2431,70 @@ def forward( return core return self._project_output([core], position_ids, attn_metadata, None) + def _quantized_decode( + self, + ag: torch.Tensor, + position_ids: torch.Tensor | None, + metadata: AttentionMetadata, + ) -> torch.Tensor: + """Keep FP8 projections around the native FP4/FP8 cache and attention kernels.""" + query, latent_cache = _fp8_attention.mla_project_qkv( + ag, + self.q_a_layernorm.weight, + self._fp8_qb_weight, + self._fp8_qb_scale, + self.k_b_proj_trans, + self.k_b_proj_trans_scale, + self.kv_a_layernorm.weight, + float(self.q_a_layernorm.variance_epsilon), + float(self.kv_a_layernorm.variance_epsilon), + ) + q_pe = query.view(-1, self.num_heads_tp, 576)[..., 512:] + cu_q = torch.empty(metadata.num_seqs + 1, dtype=torch.int32, device=ag.device) + cu_kv = torch.empty_like(cu_q) + scheduler = torch.empty(1, dtype=torch.uint32, device=ag.device) + bmm1 = torch.empty(2, dtype=torch.float32, device=ag.device) + bmm2 = torch.empty(1, dtype=torch.float32, device=ag.device) + quant_q = torch.empty_like(query, dtype=torch.uint8) + self.mqa.mla_rope_generation( + query.view(-1, self.num_heads_tp, 576), + q_pe, + latent_cache, + metadata, + cu_q, + cu_kv, + scheduler, + bmm1, + bmm2, + quant_q, + fuse_fp4_q_quant=bool(self.mqa.has_fp4_kv_cache), + ) + latent = self._attn_forward_gen( + self.mqa, + query, + None, + None, + position_ids, + metadata, + attention_input_type=AttentionInputType.generation_only, + out_scale=self.out_scale, + latent_cache=latent_cache, + q_pe=q_pe, + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=scheduler, + mla_bmm1_scale=bmm1, + mla_bmm2_scale=bmm2, + quant_q_buffer=quant_q, + ) + return _fp8_attention.mla_output( + latent, + self.v_b_proj, + self.v_b_proj_scale, + ag, + self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim, + ) + def forward_generic( self, position_ids: Optional[torch.Tensor], diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py index d2fd2f7da0ce..48b7ab905834 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py @@ -101,7 +101,7 @@ def can_fuse_fp4_mla_q_quant( num_gen = metadata.num_seqs - metadata.num_contexts return bool( _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS - and get_sm_version() == 107 + and get_sm_version() in (100, 103, 107) and num_gen > 0 and getattr(metadata, "kv_cache_manager", None) is not None and q.shape[0] > 0 diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index b20229b73cbb..dc5af9eff4e0 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -730,13 +730,13 @@ def create_weights(self): and self.quant_config.layer_quant_mode.has_fp4_kv_cache() ): if ( - mla_weight_dtype != torch.bfloat16 + mla_weight_dtype not in (torch.bfloat16, torch.float8_e4m3fn) or self.mapping.cp_size != 1 or self.apply_rotary_emb or self.kv_cache_dtype == "fp8_ds_mla" ): raise ValueError( - "FP4 MLA requires BF16 absorption weights and fused RoPE " + "FP4 MLA requires BF16 or block-scaled FP8 absorption weights and fused RoPE " "on the TRTLLM backend without context parallelism or " "fp8_ds_mla cache packing." ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py index bc7fa1ce9412..c7472b775060 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py @@ -229,13 +229,13 @@ def k3_kda_verify_kernel( cs_q: cutlass.Array, # fp32 [pool][S][H * K]: raw conv inputs, column s = position s - 2 from the golden token cs_k: cutlass.Array, cs_v: cutlass.Array, - ssm: cutlass.Array, # fp32 [pool][H][V][K] at slot stride ssm_stride: the state after the last golden token + ssm: cutlass.Array, # fp32 or bf16 [pool][H][V][K] at slot stride ssm_stride: the last golden token's state state_tok: cutlass.Array, # fp32 [pool][3][NUM_SPEC][H][K]: the records of the last round's drafts slots: cutlass.Array, # int32 [N] pending: cutlass.Array, # int32 [pool]: drafts the sampler accepted last round out: cutlass.Array, # bf16 [T][H][V] proj_words: cutlass.Int32, - ssm_stride: cutlass.Int64, # fp32 elements between the pool's slots (the manager interleaves conv states) + ssm_stride: cutlass.Int64, # elements between the pool's slots (the manager interleaves conv states) H: cutlass.Constexpr[int], NUM_SPEC: cutlass.Constexpr[int], lower_bound: cutlass.Constexpr[float], @@ -356,7 +356,8 @@ def k3_kda_verify_kernel( # The drafts' records, in each slot's region of state_tok: the row innovations vn [NUM_SPEC][H][V], then beta * k # and the decay [NUM_SPEC][H][K]. The state after the accepted drafts is the pool's (the golden token's) with each # accepted draft's update replayed, S = fma(decay, S, vn * (beta * k)): the recurrence's own arithmetic, so it is - # bit-identical to the one the drafts reached. k3_kda_attn uses the same records. + # bit-identical to the one the drafts reached with an FP32 pool. A BF16 pool rounds the committed golden state + # before replay. k3_kda_attn uses the same records. CT_VN = 0 CT_WB = NUM_SPEC * H * V CT_WD = CT_WB + NUM_SPEC * H * K @@ -367,7 +368,9 @@ def k3_kda_verify_kernel( st_base = (h * V + row_a) * K for r in cutlass.range_constexpr(ROWS): for i in cutlass.range_constexpr(VEC): - r_st.store(pool.load(idx=st_base + r * K + i * 32 + lane), idx=r * VEC + i) + r_st.store( + cutlass.Float32(pool.load(idx=st_base + r * K + i * 32 + lane)), idx=r * VEC + i + ) for t_acc in cutlass.range(pend, unroll=1): rec_acc = (t_acc * H + h) * V vns_acc = [tok.load(idx=rec_acc + CT_VN + row_a + r) for r in range(ROWS)] @@ -675,7 +678,7 @@ def k3_kda_verify_kernel( if cutlass.const_expr(tr == 0): for r in cutlass.range_constexpr(ROWS): for i in cutlass.range_constexpr(VEC): - pool.store(st[r * VEC + i], idx=st_base + r * K + i * 32 + lane) + pool.store(ssm.dtype(st[r * VEC + i]), idx=st_base + r * K + i * 32 + lane) else: if lane < ROWS: tok.store( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py index 0990b5217e4c..d2d455b5674d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py @@ -108,7 +108,7 @@ def k3_kda_verify( ``proj`` bf16 [T, 4 H K + K + H + pad]: the fused projection rows [q | k | v | onorm gate | f_a | b | pad]; ``w_fb`` bf16 [H K, K], the f_b weight (out, in); ``w_q/w_k/w_v`` fp32 [H K, 4]; ``a_log`` fp32 [H]; ``dt_bias`` fp32 [H K]; ``onorm_w`` fp32 [V]; ``cs_*`` fp32 [pool, H K, 3 + num_spec] (dim-contiguous); - ``ssm`` fp32 [pool, H, V, K] (each slot dense, slots at any stride); ``state_tok`` fp32 + ``ssm`` fp32 or bf16 [pool, H, V, K] (each slot dense, slots at any stride); ``state_tok`` fp32 [pool, 3, num_spec, H, K], the drafts' records (vn, beta * k, decay); ``slots`` int32 [N]; ``pending`` int32 [pool]. With ``g_ext`` (bf16 [T, H K], the unfused f_b output) the gate is read instead of computed from f_a.""" import cuda.bindings.driver as cuda_driver @@ -133,7 +133,7 @@ def k3_kda_verify( or tuple(ssm.shape[1:]) != (num_heads, K, K) or tuple(state_tok.shape[1:]) != (3, num_spec, num_heads, K) or ssm.stride()[1:] != (K * K, K, 1) - or ssm.dtype != torch.float32 + or ssm.dtype not in (torch.float32, torch.bfloat16) or not state_tok.is_contiguous() or state_tok.dtype != torch.float32 or cs_q.shape[-1] != CONV_W - 1 + num_spec @@ -177,7 +177,16 @@ def k3_kda_verify( runtime = (proj.shape[1] // 2, n_req, ssm.stride(0)) stream = cuda_driver.CUstream(torch.cuda.current_stream(proj.device).cuda_stream) pdl = use_pdl() - key = (num_heads, num_spec, float(lower_bound), float(scale), float(eps), fold_fb, pdl) + key = ( + num_heads, + num_spec, + ssm.dtype, + float(lower_bound), + float(scale), + float(eps), + fold_fb, + pdl, + ) fn = _compiled.get(key) if fn is None: if torch.cuda.is_current_stream_capturing(): diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 783e5e0354ed..bf8bb8fcc3ba 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -504,6 +504,7 @@ l0_b200: - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_qkv.py - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_mla_attn_vb_out.py - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py + - unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py new file mode 100644 index 000000000000..7eae2fb6b86a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py @@ -0,0 +1,238 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""K3 fused FP8 projections around BF16, FP8 and NVFP4 paged MLA caches.""" + +import weakref + +import pytest +import torch +from _torch.modeling.kimi_k3_mla_reference import KimiMLAAttention +from _torch.modeling.test_kimi_k3_checkpoint import load_mla_checkpoint +from _torch.modeling.test_kimi_k3_mla import _checkpoint_weights + +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4.implementation import ( + DecodeStep, + K3DecodeMLA, +) +from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.interface import KVCacheParams +from tensorrt_llm._torch.attention.backends.utils import get_attention_backend +from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType +from tensorrt_llm._torch.utils import model_extra_attrs +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)), + reason="requires an SM100-family GPU", +) + + +@pytest.mark.parametrize("cache_dtype", ["bf16", "fp8", "nvfp4"]) +@pytest.mark.parametrize("batch", [1, 2, 4]) +@torch.inference_mode() +def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int) -> None: + import tensorrt_llm._torch.modules.kimi_k3_mla as mla_package + + torch.manual_seed(451) + monkeypatch.setattr(mla_package, "KimiK3MLAAttention", K3DecodeMLA) + config = KimiLinearConfig( + vocab_size=128, + hidden_size=256, + intermediate_size=256, + num_hidden_layers=1, + attn_res_block_size=1, + torch_dtype=torch.bfloat16, + linear_attn_config={"kda_layers": [], "full_attn_layers": [1], "num_heads": 6}, + num_attention_heads=6, + num_key_value_heads=6, + q_lora_rank=1536, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rms_norm_eps=1e-6, + mla_use_nope=True, + mla_use_output_gate=True, + max_position_embeddings=512, + ) + config._attn_implementation = "eager" + reference = KimiMLAAttention(config, layer_idx=0).bfloat16().eval() + weights = _checkpoint_weights(reference, fp8=True) + quant_algo = {"bf16": None, "fp8": QuantAlgo.FP8, "nvfp4": QuantAlgo.NVFP4}[cache_dtype] + mapping = Mapping(world_size=1, rank=0, tp_size=1) + model_config = ModelConfig( + pretrained_config=config, + mapping=mapping, + quant_config=QuantConfig(kv_cache_quant_algo=quant_algo), + quant_config_dict={ + f"model.layers.0.self_attn.{name}": QuantConfig( + quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + group_size=128, + kv_cache_quant_algo=quant_algo, + ) + for name in weights + }, + use_cute_dsl_blockscaling_mm=True, + use_cute_dsl_blockscaling_bmm=True, + ) + with torch.device("cuda"): + model = KimiLinearForCausalLM(model_config).eval() + projections = { + f"{name}.{suffix}": value + for name, pair in weights.items() + for suffix, value in pair.items() + } + projections.update( + {name: value for name, value in reference.named_parameters() if value.ndim == 1} + ) + load_mla_checkpoint(model, projections, "") + runtime = model.model.layers[0].self_attn.mixer + runtime._fp8_qb_weight = weights["q_b_proj"]["weight"].cuda().contiguous() + runtime._fp8_qb_scale = weights["q_b_proj"]["weight_scale"].cuda().contiguous() + for module in runtime.modules(): + if callable(getattr(module, "post_load_weights", None)): + module.post_load_weights() + context_length = 128 + page_size = 128 if cache_dtype == "nvfp4" else 64 + manager_class = Fp4MlaKVCacheManagerV2 if cache_dtype == "nvfp4" else KVCacheManagerV2 + manager = manager_class( + KvCacheConfig( + max_tokens=256 * batch, + dtype="auto" if cache_dtype == "bf16" else cache_dtype, + enable_block_reuse=False, + host_cache_size=0, + ), + CacheTypeCpp.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=576, + tokens_per_block=page_size, + max_seq_len=256, + max_batch_size=batch, + mapping=mapping, + dtype={"bf16": DataType.BF16, "fp8": DataType.FP8, "nvfp4": DataType.NVFP4}[cache_dtype], + max_num_tokens=batch * context_length, + pretrained_config=config, + ) + try: + requests = manager.add_dummy_requests(list(range(batch)), [context_length + 1] * batch) + assert requests is not None + for request in requests: + request.state = LlmRequestState.GENERATION_IN_PROGRESS + request.prompt_len = context_length + request.py_prompt_len = context_length + + def metadata(context: bool): + meta = get_attention_backend("TRTLLM").Metadata( + seq_lens=torch.full((batch,), context_length if context else 1, dtype=torch.int), + request_ids=list(range(batch)), + max_num_requests=batch, + num_contexts=batch if context else 0, + prompt_lens=[context_length] * batch, + max_num_tokens=batch * context_length, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0 if context else context_length] * batch, + ), + mapping=mapping, + enable_context_mla_with_cached_kv=True, + ) + meta.prepare() + return meta + + hidden = ( + torch.randn(batch, context_length + 1, 256, device="cuda", dtype=torch.bfloat16) * 0.2 + ) + ctx = metadata(True) + attrs = { + "attention_metadata": weakref.ref(ctx), + "mla_layers": {runtime.layer_idx_str: weakref.ref(runtime)}, + } + with model_extra_attrs(attrs): + runtime.forward_generic( + None, hidden[:, :context_length].flatten(0, 1).contiguous(), ctx + ) + gen = metadata(False) + attrs["attention_metadata"] = weakref.ref(gen) + x = hidden[:, -1].contiguous() + step = DecodeStep(batch, batch, 1) + assert runtime.will_run_decode_branch(gen, step) + with model_extra_attrs(attrs): + expected = runtime.forward_generic(None, x, gen) + actual = runtime(None, x, gen, step) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = runtime(None, x, gen, step) + graph.replay() + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.04) + relative_error = ( + actual.float() - expected.float() + ).norm() / expected.float().norm().clamp_min(1e-8) + assert relative_error < 0.04, relative_error + torch.testing.assert_close(captured, actual, atol=0, rtol=0) + assert runtime.k_b_proj_trans.dtype == torch.float8_e4m3fn + assert runtime.v_b_proj.dtype == torch.float8_e4m3fn + assert runtime.k_b_proj_trans_dequant is None and runtime.v_b_proj_dequant is None + finally: + manager.shutdown() + + +@pytest.mark.parametrize("batch,tokens", [(1, 1), (2, 1), (4, 1), (1, 4), (2, 4), (4, 4)]) +@torch.inference_mode() +def test_fp4_core_matches_dequantized_reference(monkeypatch, batch: int, tokens: int) -> None: + from _torch.attention.test_fp4_mla import ( + _build_fp4_mla_attention_decode_case, + _fp4_mla_attention_decode_reference, + ) + + from tensorrt_llm._torch.attention.backends.fp4_mla import run_fp4_mla_attention_decode + + monkeypatch.setenv("TRTLLM_FP4_MLA_ATTENTION_BACKEND", "triton") + manager, meta, query, kv_rank, rope_dim = _build_fp4_mla_attention_decode_case( + seq_lens=[131 + 2 * i for i in range(batch)], + num_heads=6, + seed=452, + query_len_per_seq=tokens, + ) + try: + output = torch.empty_like(query[..., :kv_rank]) + state = meta.fp4_mla_state + run_fp4_mla_attention_decode( + meta, + layer_idx=0, + local_layer=0, + q=query, + output=output, + sm_scale=0.1, + kv_lora_rank=kv_rank, + qk_rope_head_dim=rope_dim, + prequantized_q=state.prequantized_q, + prequantized_q_sf=state.prequantized_q_sf, + q_batch_capacity=state.q_batch_capacity, + ) + expected, _ = _fp4_mla_attention_decode_reference( + meta, + query[..., :kv_rank], + query[..., kv_rank:], + sm_scale=0.1, + kv_lora_rank=kv_rank, + qk_rope_head_dim=rope_dim, + ) + actual = output.float().flatten(1) + expected = expected.flatten(1) + relative = (actual - expected).norm(dim=1) / expected.norm(dim=1) + cosine = torch.nn.functional.cosine_similarity(actual, expected, dim=1) + assert bool((relative < 0.15).all()), relative + assert bool((cosine > 0.99).all()), cosine + torch.testing.assert_close(actual, expected, atol=0.15, rtol=0.15) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py index d42ced111f03..2ff3e5cb8fc9 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py @@ -182,6 +182,16 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): ) torch.testing.assert_close(actual_q, expected_q, atol=0.002, rtol=0.02) torch.testing.assert_close(pool, expected_pool, atol=0.002, rtol=0.02) + projected_q, latent_cache = _ops().mla_project_qkv( + ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, 1.0e-6, 1.0e-6 + ) + torch.testing.assert_close(projected_q, actual_q, atol=0, rtol=0) + latent = ag[:, 1536:2048].float() + latent = ( + latent * torch.rsqrt(latent.square().mean(-1, keepdim=True) + 1.0e-6) * kv_norm.float() + ).bfloat16() + torch.testing.assert_close(latent_cache[:, :512], latent, atol=0.002, rtol=0.02) + torch.testing.assert_close(latent_cache[:, 512:], ag[:, 2048:2112], atol=0, rtol=0) vb_dequant = vb.float() * vb_scale.repeat_interleave(128, 1).repeat_interleave(128, 2) expected_output = [] for row in range(m): @@ -209,8 +219,20 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): 2112, ) torch.testing.assert_close(actual, torch.stack(expected_output), atol=0.002, rtol=0.02) + latent_output = _rand(m, h * 512, scale=0.1) + projected = torch.einsum( + "mhk,hvk->mhv", _quantize(latent_output.view(m, h, 512)), vb_dequant + ).bfloat16() + expected = (projected.flatten(1).float() * ag[:, 2112:].float().sigmoid()).bfloat16() + torch.testing.assert_close( + _ops().mla_output(latent_output, vb, vb_scale, ag, 2112), expected, atol=0.002, rtol=0.02 + ) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): + graph_projected_q, graph_latent = _ops().mla_project_qkv( + ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, 1.0e-6, 1.0e-6 + ) + graph_projected_output = _ops().mla_output(latent_output, vb, vb_scale, ag, 2112) graph_q = _ops().mla_qkv( ag, qa_norm, @@ -245,6 +267,9 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): graph.replay() torch.testing.assert_close(graph_q, actual_q, atol=0, rtol=0) torch.testing.assert_close(graph_output, actual, atol=0, rtol=0) + torch.testing.assert_close(graph_projected_q, actual_q, atol=0, rtol=0) + torch.testing.assert_close(graph_latent, latent_cache, atol=0, rtol=0) + torch.testing.assert_close(graph_projected_output, expected, atol=0.002, rtol=0.02) lengths.add_(1) graph.replay() updated_q = _ops().mla_qkv( diff --git a/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py index 3abd20fbcc5f..b3d18f3f3eb6 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py +++ b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py @@ -48,7 +48,12 @@ def load_ops(): # --------------------------------------------------------------------------------------------------------------- # The cache objects # --------------------------------------------------------------------------------------------------------------- -def build_manager(num_layers: int, num_spec: int | None = None, max_batch_size: int = 8): +def build_manager( + num_layers: int, + num_spec: int | None = None, + max_batch_size: int = 8, + state_dtype: torch.dtype = torch.float32, +): """A real MambaHybridCacheManagerV2 with ``num_layers`` KDA layers (and one attention layer). With ``num_spec``: MTP-style speculation of ``num_spec`` drafts, the KDA replay caches and the drafts' records (``kda_draft_records``) that k3_kda_verify / k3_kda_attn read and write.""" @@ -69,7 +74,7 @@ def build_manager(num_layers: int, num_spec: int | None = None, max_batch_size: mamba_num_layers=num_layers, mamba_layer_mask=[True] * num_layers + [False], mamba_cache_dtype=torch.bfloat16, - mamba_ssm_cache_dtype=torch.float32, + mamba_ssm_cache_dtype=state_dtype, kv_cache_config=KvCacheConfig(max_tokens=512, enable_block_reuse=False), kv_cache_type=CacheTypeCpp.SELF, num_layers=1, diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py new file mode 100644 index 000000000000..02b4d98a6a77 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 state storage with FP32 recurrence, draft replay and CUDA graphs.""" + +import _kda_cells as kc +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)), + reason="requires an SM100-family GPU", +) + + +def _run(weights: dict, pools: dict, projection: torch.Tensor, slots: torch.Tensor, drafts: int): + return k3_kda_verify( + projection, + weights["w_fb"], + weights["w_q"], + weights["w_k"], + weights["w_v"], + weights["a_log"], + weights["dt_bias"], + weights["onorm_w"], + pools["cs_q"], + pools["cs_k"], + pools["cs_v"], + pools["ssm"], + pools["state_tok"], + slots, + pools["pending"], + drafts, + kc.LOWER_BOUND, + kc.SCALE, + kc.EPS, + ) + + +@pytest.mark.parametrize("batch,steps", [(1, 2), (2, 4), (4, 8)]) +@torch.inference_mode() +def test_bf16_verify_replay_and_graph(batch: int, steps: int) -> None: + manager = kc.build_manager(2, num_spec=steps - 1, state_dtype=torch.bfloat16) + try: + slots = kc.request_slots(manager, batch, first_id=700) + pools = kc.layer_pools(manager, 1) + kc.fill_pools(pools, seed=431) + reference = kc.clone_pools(pools) + reference["ssm"] = pools["ssm"].float() + weights = kc.make_weights(432) + initial = kc.snapshot(pools) + generator = torch.Generator(device="cuda").manual_seed(433) + projection = torch.randn( + batch * steps, kc.PROJ, generator=generator, device="cuda" + ).bfloat16() + # Compile both storage dtypes before capture, then restore every live buffer. + _run(weights, reference, projection, slots, steps - 1) + _run(weights, pools, projection, slots, steps - 1) + for name, value in initial.items(): + pools[name].copy_(value) + reference[name].copy_(value) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = _run(weights, pools, projection, slots, steps - 1) + for rnd in range(5): + projection.copy_(torch.randn(projection.shape, generator=generator, device="cuda")) + graph.replay() + expected = _run(weights, reference, projection, slots, steps - 1) + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.01) + torch.testing.assert_close(pools["ssm"], reference["ssm"].bfloat16(), atol=0, rtol=0) + torch.testing.assert_close( + pools["state_tok"], reference["state_tok"], atol=1e-5, rtol=1e-5 + ) + # Both paths start the next round from a BF16 committed golden state. + reference["ssm"].copy_(reference["ssm"].bfloat16()) + accepted = torch.tensor( + kc.pending_schedule(batch, steps - 1, rnd), device="cuda", dtype=torch.int32 + ) + for p in (pools, reference): + p["pending"][slots.long()] = accepted + untouched = kc.other_slots(pools["ssm"].shape[0], slots.tolist()) + assert kc.same_rows(pools, initial, untouched, ("ssm", "state_tok", "cs_q", "cs_k", "cs_v")) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py index 4dbc56c7b0ac..2be90a8d7255 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py @@ -46,6 +46,24 @@ def test_fp8_mla_preserves_raw_query_weights_before_backend_transformation() -> torch.testing.assert_close(attention._fp8_qb_scale, torch.full((1, 1), 0.125)) +@pytest.mark.parametrize("sm", [(10, 3), (10, 7)]) +def test_nvfp4_kda_selects_rubin_projection_backend(monkeypatch, sm) -> None: + from tensorrt_llm._torch.model_config import ModelConfig + + captured = [] + + def initialize(module, *args, model_config, **kwargs): + torch.nn.Module.__init__(module) + captured.append(model_config) + + monkeypatch.setattr(modeling.KimiKDALinearAttention, "__init__", initialize) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: sm) + config = ModelConfig(use_cute_dsl_blockscaling_mm=False) + modeling.K3DecodeKDA(model_config=config) + assert captured[0].use_cute_dsl_blockscaling_mm == (sm == (10, 7)) + assert not config.use_cute_dsl_blockscaling_mm + + @pytest.mark.parametrize( "sm,fmt,moe_tp,moe_ep", [ From 186975c533aef66fb18b9b217316cd9352f7c82b Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:10:32 +0800 Subject: [PATCH 5/8] Cover K3 quantized-cache verification layouts (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../kimi_k3_nvfp4/implementation.py | 5 ++- .../test_modeling_v2_k3_cache_precisions.py | 31 +++++++++++++------ 2 files changed, 25 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py index 689534105cc7..60b59a40f718 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py @@ -2368,6 +2368,7 @@ def _k3_step_view(self, metadata: AttentionMetadata, step: DecodeStep) -> Option or not 0 < step.tokens_per_request <= 8 or step.num_tokens != metadata.num_generations * step.tokens_per_request or metadata.beam_width != 1 + or getattr(metadata, "is_ragged_verify", False) or getattr(metadata, "is_spec_dec_tree", False) or getattr(metadata, "is_spec_dec_dynamic_tree", False) ): @@ -2455,7 +2456,9 @@ def _quantized_decode( scheduler = torch.empty(1, dtype=torch.uint32, device=ag.device) bmm1 = torch.empty(2, dtype=torch.float32, device=ag.device) bmm2 = torch.empty(1, dtype=torch.float32, device=ag.device) - quant_q = torch.empty_like(query, dtype=torch.uint8) + quant_q = torch.empty( + (ag.shape[0], self.num_heads_tp, 576), dtype=torch.uint8, device=ag.device + ) self.mqa.mla_rope_generation( query.view(-1, self.num_heads_tp, 576), q_pe, diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py index 7eae2fb6b86a..d0424aa15c09 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """K3 fused FP8 projections around BF16, FP8 and NVFP4 paged MLA caches.""" +import copy import weakref import pytest @@ -15,7 +16,7 @@ K3DecodeMLA, ) from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 -from tensorrt_llm._torch.attention.backends.interface import KVCacheParams +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures, KVCacheParams from tensorrt_llm._torch.attention.backends.utils import get_attention_backend from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig from tensorrt_llm._torch.model_config import ModelConfig @@ -24,7 +25,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType from tensorrt_llm._torch.utils import model_extra_attrs -from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -36,9 +37,9 @@ @pytest.mark.parametrize("cache_dtype", ["bf16", "fp8", "nvfp4"]) -@pytest.mark.parametrize("batch", [1, 2, 4]) +@pytest.mark.parametrize("batch,tokens", [(1, 1), (2, 1), (4, 1), (1, 4), (2, 4), (4, 4)]) @torch.inference_mode() -def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int) -> None: +def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int, tokens: int) -> None: import tensorrt_llm._torch.modules.kimi_k3_mla as mla_package torch.manual_seed(451) @@ -121,9 +122,10 @@ def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int) -> dtype={"bf16": DataType.BF16, "fp8": DataType.FP8, "nvfp4": DataType.NVFP4}[cache_dtype], max_num_tokens=batch * context_length, pretrained_config=config, + spec_config=MTPDecodingConfig(max_draft_len=tokens - 1) if tokens > 1 else None, ) try: - requests = manager.add_dummy_requests(list(range(batch)), [context_length + 1] * batch) + requests = manager.add_dummy_requests(list(range(batch)), [context_length + tokens] * batch) assert requests is not None for request in requests: request.state = LlmRequestState.GENERATION_IN_PROGRESS @@ -132,7 +134,9 @@ def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int) -> def metadata(context: bool): meta = get_attention_backend("TRTLLM").Metadata( - seq_lens=torch.full((batch,), context_length if context else 1, dtype=torch.int), + seq_lens=torch.full( + (batch,), context_length if context else tokens, dtype=torch.int + ), request_ids=list(range(batch)), max_num_requests=batch, num_contexts=batch if context else 0, @@ -145,12 +149,15 @@ def metadata(context: bool): ), mapping=mapping, enable_context_mla_with_cached_kv=True, + max_total_draft_tokens=tokens - 1, + runtime_features=AttentionRuntimeFeatures(has_speculative_draft_tokens=tokens > 1), ) meta.prepare() return meta hidden = ( - torch.randn(batch, context_length + 1, 256, device="cuda", dtype=torch.bfloat16) * 0.2 + torch.randn(batch, context_length + tokens, 256, device="cuda", dtype=torch.bfloat16) + * 0.2 ) ctx = metadata(True) attrs = { @@ -163,8 +170,8 @@ def metadata(context: bool): ) gen = metadata(False) attrs["attention_metadata"] = weakref.ref(gen) - x = hidden[:, -1].contiguous() - step = DecodeStep(batch, batch, 1) + x = hidden[:, context_length:].flatten(0, 1).contiguous() + step = DecodeStep(batch * tokens, batch, tokens) assert runtime.will_run_decode_branch(gen, step) with model_extra_attrs(attrs): expected = runtime.forward_generic(None, x, gen) @@ -219,8 +226,12 @@ def test_fp4_core_matches_dequantized_reference(monkeypatch, batch: int, tokens: prequantized_q_sf=state.prequantized_q_sf, q_batch_capacity=state.q_batch_capacity, ) + reference_meta = copy.copy(meta) + reference_meta.fp4_mla_state = copy.copy(state) + reference_meta.fp4_mla_state.workspaces = dict(state.workspaces) + reference_meta.fp4_mla_state.workspaces.pop("_fp4_mla_attention_p_buf", None) expected, _ = _fp4_mla_attention_decode_reference( - meta, + reference_meta, query[..., :kv_rank], query[..., kv_rank:], sm_scale=0.1, From 8ecc69f3772e8421d119cd11c4f08d0ea5e64722 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:21:03 +0800 Subject: [PATCH 6/8] Isolate NVFP4 state precision dispatch from MXFP4 (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../kimi_k3_nvfp4/implementation.py | 6 +- .../models/kimi_k3_vl/kimi_k3_nvfp4/verify.py | 215 ++++++++++++++++++ .../k3_kda_verify/k3_kda_verify_kernel.py | 13 +- .../cute_dsl_kernels/k3_kda_verify/op.py | 15 +- .../test_modeling_v2_k3_cache_precisions.py | 2 +- .../ssm/test_modeling_v2_k3_bf16_verify.py | 6 +- 6 files changed, 231 insertions(+), 26 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py index 60b59a40f718..2444f19694bd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py @@ -16,7 +16,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.k3_drafter_attn_qknorm import ( k3_drafter_attn_qknorm, ) -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify from tensorrt_llm._torch.attention.backends import AttentionMetadata from tensorrt_llm._torch.attention.backends.interface import AttentionInputType @@ -63,6 +62,7 @@ from . import decode_moe as _decode_moe from . import spec_worker as _spec_worker from . import weights as _weights +from .verify import k3_fp8_kda_verify if TYPE_CHECKING: from transformers import PretrainedConfig @@ -84,7 +84,7 @@ "kda_mtp_decode", "dsv3_router_gemm_op", "dsv3_fused_a_gemm_op", - "k3_kda_verify", + "k3_fp8_kda_verify", "k3_fp8_kda_decode", "k3_fp8_mla_qkv", "k3_fp8_mla_attn", @@ -2286,7 +2286,7 @@ def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch rows = torch.cat( (self.qkvg_proj(x), torch.nn.functional.linear(x, self._bfa_proj_weight)), dim=-1 ) - out = k3_kda_verify( + out = k3_fp8_kda_verify( rows, self.f_b_proj.weight, w_q, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py new file mode 100644 index 000000000000..c6254a25e8b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``trtllm::k3_fp8_kda_verify``: Kimi K3's KDA speculative verify for 1 + num_spec tokens per request, from the fused +projection rows to the gated-norm core output, with the drafts' records for the next round. + +State contract (differs from ``trtllm::kda_mtp_decode``'s replay caches): after a call the pool ``ssm`` holds the +state after each request's golden token and ``state_tok[slot]`` the records of its drafts: the row innovations vn +[num_spec][H][V], then beta * k and the decay [num_spec][H][K]. The conv caches hold the raw inputs at positions +-2..num_spec around the golden token. The next call starts from the pool state with the drafts the sampler accepted +(``pending[slot]``) replayed from their records, S = fma(decay, S, vn (beta k)), the recurrence's own arithmetic. +The kernel compiles on the first call for each configuration, which must happen outside CUDA-graph capture. +""" + +from __future__ import annotations + +import os +import threading +from typing import Dict, Optional + +import torch + +K = 128 +CONV_W = 4 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + # detach(): DLPack refuses tensors that require grad (weights are parameters). + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def _index_arg(t: torch.Tensor): + """An int32 index tensor the kernel reads element by element (``slots``, ``pending``), declared at its element's + alignment: the mixer passes slices such as ``state_indices[num_prefills:]``, which start on any 4-byte boundary.""" + return _arg(t, t.element_size()) + + +def _flat(t: torch.Tensor) -> torch.Tensor: + """A 1-D view of a dense tensor's storage (a view with its last two dims transposed is fine), never a copy.""" + if t.is_contiguous(): + return t.view(-1) + swapped = t.transpose(-1, -2) + if swapped.is_contiguous(): + return swapped.reshape(-1) + raise ValueError(f"expected a dense tensor, got shape {tuple(t.shape)} strides {t.stride()}") + + +def _slots_view(t: torch.Tensor) -> torch.Tensor: + """A [slots, slot elements] view of a pool whose slots are dense but may be strided (the Mamba cache manager + coalesces each slot's per-layer states), never a copy. The kernels address it flat from the first slot at 64-bit + offsets; the view keeps every extent within the DSL's 32-bit sizes however far the last slot lies.""" + return t.view(t.shape[0], -1) + + +def _words(t: torch.Tensor) -> torch.Tensor: + return _flat(t).view(torch.int32) + + +def use_pdl() -> bool: + return os.environ.get("TRTLLM_ENABLE_PDL", "1") == "1" + + +@torch.library.custom_op( + "trtllm::k3_fp8_kda_verify", + mutates_args=("cs_q", "cs_k", "cs_v", "ssm", "state_tok"), + device_types="cuda", +) +def k3_fp8_kda_verify( + proj: torch.Tensor, + w_fb: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + onorm_w: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + ssm: torch.Tensor, + state_tok: torch.Tensor, + slots: torch.Tensor, + pending: torch.Tensor, + num_spec: int, + lower_bound: float, + scale: float, + eps: float, + g_ext: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Gated-norm KDA core output [T, H, V] bf16 for T = N * (1 + num_spec) verify tokens. + + ``proj`` bf16 [T, 4 H K + K + H + pad]: the fused projection rows [q | k | v | onorm gate | f_a | b | pad]; + ``w_fb`` bf16 [H K, K], the f_b weight (out, in); ``w_q/w_k/w_v`` fp32 [H K, 4]; ``a_log`` fp32 [H]; + ``dt_bias`` fp32 [H K]; ``onorm_w`` fp32 [V]; ``cs_*`` fp32 [pool, H K, 3 + num_spec] (dim-contiguous); + ``ssm`` fp32 or bf16 [pool, H, V, K] (each slot dense, slots at any stride); ``state_tok`` fp32 + [pool, 3, num_spec, H, K], the drafts' records (vn, beta * k, decay); ``slots`` int32 [N]; ``pending`` int32 + [pool]. With ``g_ext`` (bf16 [T, H K], the unfused f_b output) the gate is read instead of computed from f_a.""" + import cuda.bindings.driver as cuda_driver + + from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import k3_kda_verify_kernel as kernel + + n_req = slots.shape[0] + num_heads = w_fb.shape[0] // K + tokens = n_req * (1 + num_spec) + v_dim = onorm_w.shape[0] + if ( + proj.dtype != torch.bfloat16 + or proj.dim() != 2 + or proj.shape[0] != tokens + or proj.shape[1] % 8 != 0 + or proj.shape[1] < 4 * num_heads * K + K + num_heads + or not proj.is_contiguous() + or tuple(w_fb.shape) != (num_heads * K, K) + or w_fb.dtype != torch.bfloat16 + or not w_fb.is_contiguous() + or v_dim != K + or tuple(ssm.shape[1:]) != (num_heads, K, K) + or tuple(state_tok.shape[1:]) != (3, num_spec, num_heads, K) + or ssm.stride()[1:] != (K * K, K, 1) + or ssm.dtype not in (torch.float32, torch.bfloat16) + or not state_tok.is_contiguous() + or state_tok.dtype != torch.float32 + or cs_q.shape[-1] != CONV_W - 1 + num_spec + or slots.dtype != torch.int32 + or pending.dtype != torch.int32 + or any( + t.dtype != torch.float32 + for t in (w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v) + ) + ): + raise ValueError( + f"k3_fp8_kda_verify: unsupported call (w_q/w_k/w_v/a_log/dt_bias/onorm_w/cs_* must be fp32: " + f"{[str(t.dtype) for t in (w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v)]}) " + f"proj {tuple(proj.shape)} {proj.dtype} contiguous {proj.is_contiguous()}, " + f"w_fb {tuple(w_fb.shape)} {w_fb.dtype} contiguous {w_fb.is_contiguous()}, onorm_w {tuple(onorm_w.shape)}, " + f"ssm {tuple(ssm.shape)} {ssm.dtype} strides {ssm.stride()}, state_tok {tuple(state_tok.shape)} " + f"{state_tok.dtype} contiguous {state_tok.is_contiguous()}, cs_q {tuple(cs_q.shape)}, num_spec {num_spec}, " + f"slots {tuple(slots.shape)} {slots.dtype}, pending {pending.dtype} (tokens expected {tokens})" + ) + fold_fb = g_ext is None + out = torch.empty(tokens, num_heads, v_dim, dtype=torch.bfloat16, device=proj.device) + args = ( + _arg(w_fb), + _arg(_words(proj)), + _arg(_words(proj if fold_fb else g_ext)), + _arg(_flat(w_q)), + _arg(_flat(w_k)), + _arg(_flat(w_v)), + _arg(_flat(a_log)), + _arg(_flat(dt_bias)), + _arg(_flat(onorm_w)), + _arg(_flat(cs_q)), + _arg(_flat(cs_k)), + _arg(_flat(cs_v)), + _arg(_slots_view(ssm)), + _arg(_slots_view(state_tok)), + _index_arg(slots), + _index_arg(pending), + _arg(_flat(out)), + ) + runtime = (proj.shape[1] // 2, n_req, ssm.stride(0)) + stream = cuda_driver.CUstream(torch.cuda.current_stream(proj.device).cuda_stream) + pdl = use_pdl() + key = ( + num_heads, + num_spec, + ssm.dtype, + float(lower_bound), + float(scale), + float(eps), + fold_fb, + pdl, + ) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "trtllm::k3_fp8_kda_verify must run once per configuration outside CUDA-graph capture first " + "(it compiles its kernel on the first call)." + ) + import cutlass.cute as cute + + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.k3_kda_verify, *args, *runtime, num_heads, num_spec, float(lower_bound), float(scale), + float(eps), fold_fb, pdl, stream, + ) # fmt: skip + fn(*args, *runtime, stream) + return out + + +@k3_fp8_kda_verify.register_fake +def _(proj, w_fb, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, state_tok, slots, pending, + num_spec, lower_bound, scale, eps, g_ext=None): # fmt: skip + return proj.new_empty( + (proj.shape[0], w_fb.shape[0] // K, onorm_w.shape[0]), dtype=torch.bfloat16 + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py index c7472b775060..bc7fa1ce9412 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/k3_kda_verify_kernel.py @@ -229,13 +229,13 @@ def k3_kda_verify_kernel( cs_q: cutlass.Array, # fp32 [pool][S][H * K]: raw conv inputs, column s = position s - 2 from the golden token cs_k: cutlass.Array, cs_v: cutlass.Array, - ssm: cutlass.Array, # fp32 or bf16 [pool][H][V][K] at slot stride ssm_stride: the last golden token's state + ssm: cutlass.Array, # fp32 [pool][H][V][K] at slot stride ssm_stride: the state after the last golden token state_tok: cutlass.Array, # fp32 [pool][3][NUM_SPEC][H][K]: the records of the last round's drafts slots: cutlass.Array, # int32 [N] pending: cutlass.Array, # int32 [pool]: drafts the sampler accepted last round out: cutlass.Array, # bf16 [T][H][V] proj_words: cutlass.Int32, - ssm_stride: cutlass.Int64, # elements between the pool's slots (the manager interleaves conv states) + ssm_stride: cutlass.Int64, # fp32 elements between the pool's slots (the manager interleaves conv states) H: cutlass.Constexpr[int], NUM_SPEC: cutlass.Constexpr[int], lower_bound: cutlass.Constexpr[float], @@ -356,8 +356,7 @@ def k3_kda_verify_kernel( # The drafts' records, in each slot's region of state_tok: the row innovations vn [NUM_SPEC][H][V], then beta * k # and the decay [NUM_SPEC][H][K]. The state after the accepted drafts is the pool's (the golden token's) with each # accepted draft's update replayed, S = fma(decay, S, vn * (beta * k)): the recurrence's own arithmetic, so it is - # bit-identical to the one the drafts reached with an FP32 pool. A BF16 pool rounds the committed golden state - # before replay. k3_kda_attn uses the same records. + # bit-identical to the one the drafts reached. k3_kda_attn uses the same records. CT_VN = 0 CT_WB = NUM_SPEC * H * V CT_WD = CT_WB + NUM_SPEC * H * K @@ -368,9 +367,7 @@ def k3_kda_verify_kernel( st_base = (h * V + row_a) * K for r in cutlass.range_constexpr(ROWS): for i in cutlass.range_constexpr(VEC): - r_st.store( - cutlass.Float32(pool.load(idx=st_base + r * K + i * 32 + lane)), idx=r * VEC + i - ) + r_st.store(pool.load(idx=st_base + r * K + i * 32 + lane), idx=r * VEC + i) for t_acc in cutlass.range(pend, unroll=1): rec_acc = (t_acc * H + h) * V vns_acc = [tok.load(idx=rec_acc + CT_VN + row_a + r) for r in range(ROWS)] @@ -678,7 +675,7 @@ def k3_kda_verify_kernel( if cutlass.const_expr(tr == 0): for r in cutlass.range_constexpr(ROWS): for i in cutlass.range_constexpr(VEC): - pool.store(ssm.dtype(st[r * VEC + i]), idx=st_base + r * K + i * 32 + lane) + pool.store(st[r * VEC + i], idx=st_base + r * K + i * 32 + lane) else: if lane < ROWS: tok.store( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py index d2d455b5674d..0990b5217e4c 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/k3_kda_verify/op.py @@ -108,7 +108,7 @@ def k3_kda_verify( ``proj`` bf16 [T, 4 H K + K + H + pad]: the fused projection rows [q | k | v | onorm gate | f_a | b | pad]; ``w_fb`` bf16 [H K, K], the f_b weight (out, in); ``w_q/w_k/w_v`` fp32 [H K, 4]; ``a_log`` fp32 [H]; ``dt_bias`` fp32 [H K]; ``onorm_w`` fp32 [V]; ``cs_*`` fp32 [pool, H K, 3 + num_spec] (dim-contiguous); - ``ssm`` fp32 or bf16 [pool, H, V, K] (each slot dense, slots at any stride); ``state_tok`` fp32 + ``ssm`` fp32 [pool, H, V, K] (each slot dense, slots at any stride); ``state_tok`` fp32 [pool, 3, num_spec, H, K], the drafts' records (vn, beta * k, decay); ``slots`` int32 [N]; ``pending`` int32 [pool]. With ``g_ext`` (bf16 [T, H K], the unfused f_b output) the gate is read instead of computed from f_a.""" import cuda.bindings.driver as cuda_driver @@ -133,7 +133,7 @@ def k3_kda_verify( or tuple(ssm.shape[1:]) != (num_heads, K, K) or tuple(state_tok.shape[1:]) != (3, num_spec, num_heads, K) or ssm.stride()[1:] != (K * K, K, 1) - or ssm.dtype not in (torch.float32, torch.bfloat16) + or ssm.dtype != torch.float32 or not state_tok.is_contiguous() or state_tok.dtype != torch.float32 or cs_q.shape[-1] != CONV_W - 1 + num_spec @@ -177,16 +177,7 @@ def k3_kda_verify( runtime = (proj.shape[1] // 2, n_req, ssm.stride(0)) stream = cuda_driver.CUstream(torch.cuda.current_stream(proj.device).cuda_stream) pdl = use_pdl() - key = ( - num_heads, - num_spec, - ssm.dtype, - float(lower_bound), - float(scale), - float(eps), - fold_fb, - pdl, - ) + key = (num_heads, num_spec, float(lower_bound), float(scale), float(eps), fold_fb, pdl) fn = _compiled.get(key) if fn is None: if torch.cuda.is_current_stream_capturing(): diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py index d0424aa15c09..76fae98fcdbf 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py @@ -81,7 +81,7 @@ def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int, tok ) for name in weights }, - use_cute_dsl_blockscaling_mm=True, + use_cute_dsl_blockscaling_mm=torch.cuda.get_device_capability() == (10, 7), use_cute_dsl_blockscaling_bmm=True, ) with torch.device("cuda"): diff --git a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py index 02b4d98a6a77..5c707480343b 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py @@ -6,7 +6,9 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog.ssm.k3_kda_verify import k3_kda_verify +from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4.verify import ( + k3_fp8_kda_verify, +) pytestmark = pytest.mark.skipif( not torch.cuda.is_available() @@ -16,7 +18,7 @@ def _run(weights: dict, pools: dict, projection: torch.Tensor, slots: torch.Tensor, drafts: int): - return k3_kda_verify( + return k3_fp8_kda_verify( projection, weights["w_fb"], weights["w_q"], From e6c431b92aecac4c35ce33645bae0e5a9647a3d4 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:32:38 +0800 Subject: [PATCH 7/8] Convert BF16 verify states explicitly in NVFP4 kernel (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../models/kimi_k3_vl/kimi_k3_nvfp4/verify.py | 4 +- .../kimi_k3_vl/kimi_k3_nvfp4/verify_kernel.py | 830 ++++++++++++++++++ 2 files changed, 832 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py index c6254a25e8b1..cc103be29262 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py @@ -113,7 +113,7 @@ def k3_fp8_kda_verify( [pool]. With ``g_ext`` (bf16 [T, H K], the unfused f_b output) the gate is read instead of computed from f_a.""" import cuda.bindings.driver as cuda_driver - from tensorrt_llm._torch.cute_dsl_kernels.k3_kda_verify import k3_kda_verify_kernel as kernel + from . import verify_kernel as kernel n_req = slots.shape[0] num_heads = w_fb.shape[0] // K @@ -200,7 +200,7 @@ def k3_fp8_kda_verify( fn = _compiled.get(key) if fn is None: fn = _compiled[key] = cute.compile( - kernel.k3_kda_verify, *args, *runtime, num_heads, num_spec, float(lower_bound), float(scale), + kernel.k3_fp8_kda_verify, *args, *runtime, num_heads, num_spec, float(lower_bound), float(scale), float(eps), fold_fb, pdl, stream, ) # fmt: skip fn(*args, *runtime, stream) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify_kernel.py new file mode 100644 index 000000000000..cf947f5f4f11 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify_kernel.py @@ -0,0 +1,830 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi K3 KDA speculative verify as one kernel: the forget-gate up-projection (f_b), the per-token pre-compute and +the delta-rule recurrence over the 1 + NUM_SPEC verify tokens. It commits the golden token's state and each draft's +records (vn, beta * k, decay), from which the next round replays the accepted drafts. + +Grid (H, N, 8), cluster (1, 1, 8), 256 threads. CTA (h, n, s) owns V rows [16 s, 16 s + 16) of local head h for +request n; warp w owns rows 16 s + 2 w and 16 s + 2 w + 1, and lane l keys k = 32 i + l (i < 4) of each. + +Before ``griddepcontrol.wait`` (written by the previous verify step, or constant): + the slot and P, the drafts the sampler accepted last round; the starting state into registers (the pool state, + committed after the last golden token, with P drafts replayed from their records); the raw conv inputs at positions + -3..-1 before the first new token (conv-cache columns P..P+2) for q and k (128 channels) and v (this CTA's 16); + the conv weights, dt_bias, A_log and the output-norm weight; the head's W_fb^T block by TMA. +After the wait (this step's fused projection rows [q | k | v | onorm gate | f_a | b]): the new tokens' raw q, k, v, + f_a, b and output gate. +f_b: g[t, c] = bf16(sum_i f_a[t, i] W_fb[128 h + c, i]), rounded as the unfused bf16 GEMV output is. +Pre-compute: warp t = verify token t: q, k = l2norm(silu(conv4(u))) (q also scaled), the lower-bound gate, beta and + this CTA's 16 v channels. +Recurrence: the arithmetic of ``kda_mtp_decode``'s V-split path over the verify tokens, unrolled; the state after + token 0 (the golden token) is committed to the pool; token t >= 1 leaves its draft's records (vn, + beta * k, decay) in ``state_tok[slot]``. +Epilogue: the gated RMSNorm of the outputs over V (per-token sums of squares stored into every peer's shared + memory, one cluster barrier), bf16 output rows; CTA 0 rewrites the q/k conv cache (the window at the + golden token and the raw inputs of the drafts), every CTA the same for its v channels. + +With FOLD_FB False the gate comes from ``g_ext`` (the unfused f_b output); the kernel is then bit-exact against +``kda_mtp_decode`` fed the same gate and replaying the same accepted drafts. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +K = 128 # key head dim; the value head dim V is the same +V = 128 +CONV_W = 4 +THREADS = 256 +WARPS = THREADS // 32 +V_SPLIT = 8 +V_CTA = V // V_SPLIT # V rows per CTA +ROWS = V_CTA // WARPS # V rows per warp +VEC = K // 32 # keys per lane +BF16_BYTES = 2 +CHUNK = 8 # bf16 per 16-byte load +SMEM = cutlass.AddressSpace.smem +RMEM = cutlass.AddressSpace.rmem + +assert ROWS == 2, "the recurrence below processes one pair of V rows per warp" + +# f_b on the tensor cores: D[c, t] = sum_i W_fb[128 h + c, i] f_a[t, i], M = 128 channels, N = 8 tokens, K = 128. +CTA_M = 128 +MMA_N = 8 +MMA_K = 16 +TMA_K_BOX = 64 +TMA_COPY_ITERS = K // TMA_K_BOX +K_BLOCKS_PER_HALF = TMA_K_BOX // MMA_K +LEADING = 16 +STRIDE = 8 * TMA_K_BOX * BF16_BYTES +A_HALF_ELEMS = CTA_M * TMA_K_BOX +B_HALF_ELEMS = MMA_N * TMA_K_BOX +STEP = (MMA_K * BF16_BYTES) >> 4 +A_BOX = A_HALF_ELEMS >> 3 +B_BOX = B_HALF_ELEMS >> 3 +TMEM_COLS = 32 + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): whether + phase ``parity`` has completed, acquiring at cluster scope. For barriers that other CTAs' st.async complete (their + complete_tx releases at cluster scope; a CTA-scope acquire does not synchronize with it).""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _try_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.try_wait.parity.acquire.cluster on a barrier of this CTA (``mbar``, its shared-memory pointer): as + ``_test_wait_cluster``, with try_wait's bounded suspend.""" + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _st_async_f32(dst, value, mbar, *, loc=None, ip=None): + """st.async of one fp32 to a shared::cluster address, completing ``mbar`` (a shared::cluster address) by 4 bytes.""" + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Float32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$0], $1, [$2];", "r,f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +def _lo(word): + return (word << cutlass.Int32(16)).bitcast(cutlass.Float32) + + +def _hi(word): + return (word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + + +def _bf16(x): + return cutlass.Float32(cutlass.BFloat16(x)) + + +def _butterfly(value): + for offset in [16, 8, 4, 2, 1]: + value += cute.arch.shuffle_sync_bfly(value, offset=offset, mask=-1, mask_and_clamp=31) + return value + + +def _store8(dst, words, idx): + """The 8 bf16 of four int32 words, as fp32, into dst[idx : idx + 8].""" + dst.store( + (_lo(cutlass.Int32(words[0])), _hi(cutlass.Int32(words[0])), _lo(cutlass.Int32(words[1])), + _hi(cutlass.Int32(words[1]))), + idx=idx, alignment=16, + ) # fmt: skip + dst.store( + (_lo(cutlass.Int32(words[2])), _hi(cutlass.Int32(words[2])), _lo(cutlass.Int32(words[3])), + _hi(cutlass.Int32(words[3]))), + idx=idx + 4, alignment=16, + ) # fmt: skip + + +def _qk_pre(tk, lane, s_uq, s_wq, s_uk, s_wk, s_q, s_k, scale): + """q and k of verify token ``tk`` from the conv history in shared memory: conv4 + SiLU, then the L2 norm (q also + scaled), lane l keys 32 i + l (the arithmetic of kda_mtp_decode's pre-compute warps). Traced inline: its Python + loops unroll.""" + pq = [cutlass.Float32(0.0)] * VEC + for i in range(VEC): + c = i * 32 + lane + conv = cutlass.Float32(0.0) + for w in range(CONV_W - 1): + conv += s_uq.load(idx=(tk + w) * K + c) * s_wq.load(idx=w * K + c) + conv += s_uq.load(idx=(tk + CONV_W - 1) * K + c) * s_wq.load(idx=(CONV_W - 1) * K + c) + e = cute.math.exp(-conv, fastmath=True) + pq[i] = conv * cute.arch.rcp_approx(cutlass.Float32(1.0) + e) + sum_q = cutlass.Float32(0.0) + for i in range(VEC): + sum_q += pq[i] * pq[i] + sum_q = _butterfly(sum_q) + rnorm_q = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + for i in range(VEC): + s_q.store(pq[i] * rnorm_q, idx=tk * K + i * 32 + lane) + pk = [cutlass.Float32(0.0)] * VEC + for i in range(VEC): + c = i * 32 + lane + conv = s_uk.load(idx=tk * K + c) * s_wk.load(idx=c) + for w in range(1, CONV_W - 1): + conv += s_uk.load(idx=(tk + w) * K + c) * s_wk.load(idx=w * K + c) + conv += s_uk.load(idx=(tk + CONV_W - 1) * K + c) * s_wk.load(idx=(CONV_W - 1) * K + c) + pk[i] = conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-conv, fastmath=True) + ) + sum_k = cutlass.Float32(0.0) + for i in range(VEC): + sum_k += pk[i] * pk[i] + sum_k = _butterfly(sum_k) + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in range(VEC): + s_k.store(pk[i] * rnorm_k, idx=tk * K + i * 32 + lane) + + +@cute.kernel +def k3_fp8_kda_verify_kernel( + tma_wfb: cutlass.GridConstant[ + cuda.TensorMap + ], # W_fb [H * K, K] bf16, the head's 128 rows by one call + tma_fa: cutlass.GridConstant[ + cuda.TensorMap + ], # the projection rows' f_a columns, box 8 tokens x 64 + proj: cutlass.Array, # int32 words of the fused projection rows, bf16 [T, 2 * proj_words] + g_ext: cutlass.Array, # int32 words of bf16 [T, H * K], the unfused f_b output (FOLD_FB False only) + w_q: cutlass.Array, # fp32 [H * K, CONV_W] + w_k: cutlass.Array, + w_v: cutlass.Array, + a_log: cutlass.Array, # fp32 [H] + dt_bias: cutlass.Array, # fp32 [H * K] + onorm_w: cutlass.Array, # fp32 [V] + cs_q: cutlass.Array, # fp32 [pool][S][H * K]: raw conv inputs, column s = position s - 2 from the golden token + cs_k: cutlass.Array, + cs_v: cutlass.Array, + ssm: cutlass.Array, # fp32 or bf16 [pool][H][V][K] at slot stride ssm_stride: the last golden token's state + state_tok: cutlass.Array, # fp32 [pool][3][NUM_SPEC][H][K]: the records of the last round's drafts + slots: cutlass.Array, # int32 [N] + pending: cutlass.Array, # int32 [pool]: drafts the sampler accepted last round + out: cutlass.Array, # bf16 [T][H][V] + proj_words: cutlass.Int32, + ssm_stride: cutlass.Int64, # elements between the pool's slots (the manager interleaves conv states) + H: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + FOLD_FB: cutlass.Constexpr[bool], + USE_PDL: cutlass.Constexpr[bool], +): + NT = NUM_SPEC + 1 # verify tokens per request + S = CONV_W - 1 + NUM_SPEC # conv-cache columns + HK = H * K + ROWS_U = ( + CONV_W - 1 + NT + ) # raw conv inputs by position: rows 0..2 before token 0, then the tokens + # The fused projection row [q | k | v | onorm gate | f_a | b | pad], in int32 words (bf16 pairs). + Q_W = 0 + K_W = HK // 2 + V_W = HK + OG_W = 3 * HK // 2 + FA_W = 2 * HK + B_COL = 4 * HK + K # bf16 column + assert NT <= WARPS and NT % 2 == 0 + + tidx, _, _ = cute.arch.thread_idx() + lane = tidx % 32 + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + h, n, vs = cute.arch.block_idx() + v0 = vs * V_CTA + ch0 = h * K + row0 = n * NT + + smem_a = cutlass.Array( + cutlass.BFloat16, CTA_M * K, space=SMEM, alignment=1024 + ) # W_fb rows, 128B-swizzled + smem_b = cutlass.Array( + cutlass.BFloat16, MMA_N * K, space=SMEM, alignment=1024 + ) # f_a, 128B-swizzled + w_full = cutlass.Array(cutlass.Int64, 1, space=SMEM, alignment=8) + x_full = cutlass.Array(cutlass.Int64, 1, space=SMEM, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=SMEM, alignment=8) + ss_ready = cutlass.Array( + cutlass.Int64, 1, space=SMEM, alignment=8 + ) # the peers' norm partials have landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=SMEM) + s_uq = cutlass.Array(cutlass.Float32, ROWS_U * K, space=SMEM, alignment=16) + s_uk = cutlass.Array(cutlass.Float32, ROWS_U * K, space=SMEM, alignment=16) + s_uv = cutlass.Array(cutlass.Float32, ROWS_U * V_CTA, space=SMEM, alignment=16) + s_wq = cutlass.Array(cutlass.Float32, CONV_W * K, space=SMEM, alignment=16) # [tap][channel] + s_wk = cutlass.Array(cutlass.Float32, CONV_W * K, space=SMEM, alignment=16) + s_wv = cutlass.Array(cutlass.Float32, CONV_W * V_CTA, space=SMEM, alignment=16) + s_dtb = cutlass.Array(cutlass.Float32, K, space=SMEM, alignment=16) + s_onw = cutlass.Array(cutlass.Float32, V_CTA, space=SMEM, alignment=16) + s_gr = cutlass.Array( + cutlass.Float32, NT * K, space=SMEM, alignment=16 + ) # f_b output, bf16-rounded + s_braw = cutlass.Array(cutlass.Float32, NT, space=SMEM, alignment=16) + s_og = cutlass.Array(cutlass.Float32, NT * V_CTA, space=SMEM, alignment=16) + s_q = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) + s_k = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) + s_dec = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) # exp(gate) + s_kd = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) # exp(gate) * k + s_bk = cutlass.Array(cutlass.Float32, NT * K, space=SMEM, alignment=16) # beta * k + s_beta = cutlass.Array(cutlass.Float32, NT, space=SMEM, alignment=16) + s_v = cutlass.Array(cutlass.Float32, NT * V_CTA, space=SMEM, alignment=16) + s_o = cutlass.Array( + cutlass.Float32, NT * V_CTA, space=SMEM, alignment=16 + ) # bf16-rounded raw outputs + s_ss = cutlass.Array( + cutlass.Float32, V_SPLIT * NT, space=SMEM, alignment=16 + ) # [source CTA][token] + s_rs = cutlass.Array(cutlass.Float32, NT, space=SMEM, alignment=16) + r_st = cutlass.Array(cutlass.Float32, ROWS * VEC, space=RMEM) + + # ---- Prologue: nothing here depends on the predecessor grid. + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + tma_ptr_w = tma_wfb.get_ptr() + tma_ptr_x = tma_fa.get_ptr() + if cutlass.const_expr(FOLD_FB): + if warp == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_x) + if prims.elect_sync(): + prims.mbarrier_init(w_full, 1) + prims.mbarrier_init(x_full, 1) + prims.mbarrier_init(acc_done, 1) + if warp == 2: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == 3: + if prims.elect_sync(): + prims.mbarrier_init(ss_ready, 1) + prims.mbarrier_arrive_expect_tx(ss_ready, V_SPLIT * NT * 4) + prims.fence_mbarrier_init() + # Cluster formation: the peers' shared memory is addressable from here on. + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if cutlass.const_expr(FOLD_FB): + if warp == 0: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(w_full, CTA_M * K * BF16_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a, + tma_ptr_w, + (cutlass.Int32(0), ch0, cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(0)), + w_full, + ) + + slot = slots.load(idx=n) + pend = pending.load(idx=slot) + pend = cutlass.Int32( + cutlass.select_(pend > cutlass.Int32(NUM_SPEC), cutlass.Int32(NUM_SPEC), pend) + ) + exp_a = cute.math.exp(a_log.load(idx=h), fastmath=True) + row_a = v0 + warp * ROWS + # The warp owns both rows and lane l keys 32 i + l (the arithmetic of kda_mtp_decode). + # The drafts' records, in each slot's region of state_tok: the row innovations vn [NUM_SPEC][H][V], then beta * k + # and the decay [NUM_SPEC][H][K]. The state after the accepted drafts is the pool's (the golden token's) with each + # accepted draft's update replayed, S = fma(decay, S, vn * (beta * k)): the recurrence's own arithmetic, so it is + # bit-identical to the one the drafts reached with an FP32 pool. A BF16 pool rounds the committed golden state + # before replay. k3_kda_attn uses the same records. + CT_VN = 0 + CT_WB = NUM_SPEC * H * V + CT_WD = CT_WB + NUM_SPEC * H * K + CT_SLOT = CT_WD + NUM_SPEC * H * K # a slot's records + # The slot's pool state and records from their first element: the slot offset in 64 bits, once. + pool = ssm.subview(cutlass.Int64(slot) * ssm_stride) + tok = state_tok.subview(cutlass.Int64(slot) * CT_SLOT) + st_base = (h * V + row_a) * K + for r in cutlass.range_constexpr(ROWS): + for i in cutlass.range_constexpr(VEC): + r_st.store( + cutlass.Float32(pool.load(idx=st_base + r * K + i * 32 + lane)), idx=r * VEC + i + ) + for t_acc in cutlass.range(pend, unroll=1): + rec_acc = (t_acc * H + h) * V + vns_acc = [tok.load(idx=rec_acc + CT_VN + row_a + r) for r in range(ROWS)] + wbs_acc = [tok.load(idx=CT_WB + (t_acc * H + h) * K + i * 32 + lane) for i in range(VEC)] + wds_acc = [tok.load(idx=CT_WD + (t_acc * H + h) * K + i * 32 + lane) for i in range(VEC)] + sts_acc = [r_st.load(idx=j) for j in range(ROWS * VEC)] + for r in cutlass.range_constexpr(ROWS): + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + vb0_acc, vb1_acc = cute.arch.mul_packed_f32x2( + (vns_acc[r], vns_acc[r]), (wbs_acc[_p], wbs_acc[_p + 1]) + ) + sts_acc[r * VEC + _p], sts_acc[r * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(wds_acc[_p], wds_acc[_p + 1]), src_b=(sts_acc[r * VEC + _p], sts_acc[r * VEC + _p + 1]), + src_c=(vb0_acc, vb1_acc), + ) # fmt: skip + for j in cutlass.range_constexpr(ROWS * VEC): + r_st.store(sts_acc[j], idx=j) + + # Conv weights, transposed to [tap][channel]: q by threads 0..127, k by 128..255. + if tidx < K: + wq4 = w_q.load(idx=(ch0 + tidx) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wq.store(cutlass.Float32(wq4[w]), idx=w * K + tidx) + else: + ck_w = tidx - K + wk4 = w_k.load(idx=(ch0 + ck_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wk.store(cutlass.Float32(wk4[w]), idx=w * K + ck_w) + # Conv history (positions -3..-1 = columns P..P+2), the v weights, the output-norm weight, dt_bias. + QH = 3 * K // 4 # float4 loads per q (or k) history + VH = 3 * V_CTA // 4 + if tidx < QH: + m_q = tidx // (K // 4) + c4_q = (tidx % (K // 4)) * 4 + s_uq.store(cs_q.load(idx=(slot * S + pend + m_q) * HK + ch0 + c4_q, vector_size=4, alignment=16), + idx=m_q * K + c4_q, alignment=16) # fmt: skip + elif tidx < 2 * QH: + m_k = (tidx - QH) // (K // 4) + c4_k = ((tidx - QH) % (K // 4)) * 4 + s_uk.store(cs_k.load(idx=(slot * S + pend + m_k) * HK + ch0 + c4_k, vector_size=4, alignment=16), + idx=m_k * K + c4_k, alignment=16) # fmt: skip + elif tidx < 2 * QH + VH: + m_v = (tidx - 2 * QH) // (V_CTA // 4) + c4_v = ((tidx - 2 * QH) % (V_CTA // 4)) * 4 + s_uv.store(cs_v.load(idx=(slot * S + pend + m_v) * HK + ch0 + v0 + c4_v, vector_size=4, alignment=16), + idx=m_v * V_CTA + c4_v, alignment=16) # fmt: skip + elif tidx < 2 * QH + VH + V_CTA: + cv_w = tidx - (2 * QH + VH) + wv4 = w_v.load(idx=(ch0 + v0 + cv_w) * CONV_W, vector_size=CONV_W, alignment=16) + for w in cutlass.range_constexpr(CONV_W): + s_wv.store(cutlass.Float32(wv4[w]), idx=w * V_CTA + cv_w) + elif tidx < 2 * QH + VH + 2 * V_CTA: + cv_n = tidx - (2 * QH + VH + V_CTA) + s_onw.store(onorm_w.load(idx=v0 + cv_n), idx=cv_n) + if tidx < K // 4: + s_dtb.store( + dt_bias.load(idx=ch0 + tidx * 4, vector_size=4, alignment=16), + idx=tidx * 4, + alignment=16, + ) + + # ---- This step's projection rows. + if cutlass.const_expr(USE_PDL): + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(FOLD_FB): + if warp == 1: + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(x_full, MMA_N * K * BF16_BYTES) + for half in cutlass.range_constexpr(TMA_COPY_ITERS): + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(half * B_HALF_ELEMS), tma_ptr_x, + (cutlass.Int32(FA_W * 2 + half * TMA_K_BOX), row0), x_full, + ) # fmt: skip + CH = K // CHUNK # 16-byte chunks per 128-channel row + for it in cutlass.range_constexpr((2 * NT * CH + THREADS - 1) // THREADS): + item = tidx + it * THREADS + if item < 2 * NT * CH: + stream_qk = item // (NT * CH) + t_qk = (item // CH) % NT + j_qk = item % CH + words_qk = proj.load( + idx=(row0 + t_qk) * proj_words + Q_W + ch0 // 2 + stream_qk * (K_W - Q_W) + j_qk * (CHUNK // 2), + vector_size=4, alignment=16, + ) # fmt: skip + if stream_qk == 0: + _store8(s_uq, words_qk, (CONV_W - 1 + t_qk) * K + j_qk * CHUNK) + else: + _store8(s_uk, words_qk, (CONV_W - 1 + t_qk) * K + j_qk * CHUNK) + VCH = V_CTA // CHUNK # 16-byte chunks per 16-channel v (or gate) slice + n_v = NT * VCH + n_g = 0 if FOLD_FB else NT * CH + for it in cutlass.range_constexpr((2 * n_v + NT + n_g + THREADS - 1) // THREADS): + item = tidx + it * THREADS + if item < n_v: + t_v = item // VCH + j_v = item % VCH + words_v = proj.load(idx=(row0 + t_v) * proj_words + V_W + (ch0 + v0) // 2 + j_v * (CHUNK // 2), + vector_size=4, alignment=16) # fmt: skip + _store8(s_uv, words_v, (CONV_W - 1 + t_v) * V_CTA + j_v * CHUNK) + elif item < 2 * n_v: + t_og = (item - n_v) // VCH + j_og = (item - n_v) % VCH + words_og = proj.load(idx=(row0 + t_og) * proj_words + OG_W + (ch0 + v0) // 2 + j_og * (CHUNK // 2), + vector_size=4, alignment=16) # fmt: skip + _store8(s_og, words_og, t_og * V_CTA + j_og * CHUNK) + elif item < 2 * n_v + NT: + t_b = item - 2 * n_v + word_b = proj.load(idx=(row0 + t_b) * proj_words + (B_COL + h) // 2) + b_in = _lo(word_b) + if (B_COL + h) % 2 == 1: + b_in = _hi(word_b) + s_braw.store(b_in, idx=t_b) + elif item < 2 * n_v + NT + n_g: + gi = item - (2 * n_v + NT) + t_g = gi // CH + j_g = gi % CH + words_g = g_ext.load(idx=(row0 + t_g) * (HK // 2) + ch0 // 2 + j_g * (CHUNK // 2), vector_size=4, + alignment=16) # fmt: skip + _store8(s_gr, words_g, t_g * K + j_g * CHUNK) + prims.barrier_cta_sync(0) + + # ---- f_b on the tensor cores (warp 2 issues, warps 4-7 read TMEM) beside the q/k pre-compute (warps 0-3). + if cutlass.const_expr(FOLD_FB): + if warp == 2: + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cutlass.BFloat16, + b_dtype=cutlass.BFloat16, + n_dim=MMA_N, + m_dim=CTA_M, + ) + desc_a_base = prims.Tcgen05SmemDesc.build( + start_address=smem_a, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + desc_b_base = prims.Tcgen05SmemDesc.build( + start_address=smem_b, leading_byte_offset=LEADING, stride_byte_offset=STRIDE, + layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B, + ) # fmt: skip + tmem_acc = cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32) + while not cute.arch.mbarrier_try_wait(w_full.data_ptr(), 0): + pass + while not cute.arch.mbarrier_try_wait(x_full.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for kb in cutlass.range_constexpr(TMA_COPY_ITERS * K_BLOCKS_PER_HALF): + box = kb // K_BLOCKS_PER_HALF + within = kb % K_BLOCKS_PER_HALF + if prims.elect_sync(): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, prims.CTAGroup.CTA_1, tmem_acc, + desc_a_base + (box * A_BOX + within * STEP), desc_b_base + (box * B_BOX + within * STEP), + idesc, kb != 0, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + if warp < 4: + # q and k of tokens 2w and 2w + 1 (the arithmetic of kda_mtp_decode's pre-compute warps); NT not a multiple of 4 + # rounds the tokens per warp up and the warps past NT idle. + qk_per_warp = (NT + 3) // 4 + for sub in cutlass.range_constexpr(qk_per_warp): + tk = warp * qk_per_warp + sub + if cutlass.const_expr(NT % 4 == 0): + _qk_pre(tk, lane, s_uq, s_wq, s_uk, s_wk, s_q, s_k, scale) + else: + if tk < NT: + _qk_pre(tk, lane, s_uq, s_wq, s_uk, s_wk, s_q, s_k, scale) + else: + if cutlass.const_expr(FOLD_FB): + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), num=MMA_N + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + c_fb = (warp - 4) * 32 + lane + for t_fb in cutlass.range_constexpr(NT): + s_gr.store(_bf16(cutlass.Float32(acc[t_fb])), idx=t_fb * K + c_fb) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.barrier_cta_sync(1, thread_count=128) + if warp == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), TMEM_COLS + ) + # v (this CTA's 16 channels) and beta, one (token, channel) per thread of warps 4-7. + item_v = tidx - 128 + if item_v < NT * V_CTA: + t_cv = item_v // V_CTA + c_cv = item_v % V_CTA + conv_v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(CONV_W): + conv_v += s_uv.load(idx=(t_cv + w) * V_CTA + c_cv) * s_wv.load(idx=w * V_CTA + c_cv) + conv_v = conv_v * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-conv_v, fastmath=True) + ) + s_v.store(conv_v, idx=t_cv * V_CTA + c_cv) + if item_v < NT: + b_pre = s_braw.load(idx=item_v) + s_beta.store( + cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-b_pre, fastmath=True)), + idx=item_v, + ) + prims.barrier_cta_sync(0) + + # ---- The recurrence's per-token operands: warp t, token t: decay = exp(gate), decay * k, beta * k. + if warp < NT: + tg = warp + r_beta_g = s_beta.load(idx=tg) + for i_pair in cutlass.range_constexpr(VEC // 2): + dks = [] + for i in (i_pair * 2, i_pair * 2 + 1): + c = i * 32 + lane + g_raw = s_gr.load(idx=tg * K + c) + s_dtb.load(idx=c) + x = exp_a * g_raw + sig = cute.arch.rcp_approx(cutlass.Float32(1.0) + cute.math.exp(-x, fastmath=True)) + dks.append( + (cute.math.exp(lower_bound * sig, fastmath=True), s_k.load(idx=tg * K + c), c) + ) + (d0, k0v, c0), (d1, k1v, c1) = dks + bk0, bk1 = cute.arch.mul_packed_f32x2((r_beta_g, r_beta_g), (k0v, k1v)) + kd0, kd1 = cute.arch.mul_packed_f32x2((d0, d1), (k0v, k1v)) + s_dec.store(d0, idx=tg * K + c0) + s_dec.store(d1, idx=tg * K + c1) + s_kd.store(kd0, idx=tg * K + c0) + s_kd.store(kd1, idx=tg * K + c1) + s_bk.store(bk0, idx=tg * K + c0) + s_bk.store(bk1, idx=tg * K + c1) + prims.barrier_cta_sync(0) + + # ---- Recurrence over the verify tokens (register-resident state). + st = [r_st.load(idx=j) for j in range(ROWS * VEC)] + outs = [] + for tr in cutlass.range_constexpr(NT): + r_v_val = s_v.load(idx=tr * V_CTA + warp * ROWS + lane % ROWS) + r_q = [cutlass.Float32(0.0)] * VEC + r_k = [cutlass.Float32(0.0)] * VEC + r_decay = [cutlass.Float32(0.0)] * VEC + r_bk = [cutlass.Float32(0.0)] * VEC + for i in cutlass.range_constexpr(VEC): + ki = i * 32 + lane + r_q[i] = s_q.load(idx=tr * K + ki) + r_k[i] = s_kd.load(idx=tr * K + ki) + r_bk[i] = s_bk.load(idx=tr * K + ki) + r_decay[i] = s_dec.load(idx=tr * K + ki) + ra = 0 + rb = 1 + r_va = cute.arch.shuffle_sync(r_v_val, ra) + r_vb = cute.arch.shuffle_sync(r_v_val, rb) + shk_a1 = cutlass.Float32(0.0) + shk_a2 = cutlass.Float32(0.0) + shk_b1 = cutlass.Float32(0.0) + shk_b2 = cutlass.Float32(0.0) + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + shk_a1, shk_a2 = cute.arch.fma_packed_f32x2( + src_a=(st[ra * VEC + _p], st[ra * VEC + _p + 1]), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_a1, shk_a2), + ) + shk_b1, shk_b2 = cute.arch.fma_packed_f32x2( + src_a=(st[rb * VEC + _p], st[rb * VEC + _p + 1]), + src_b=(r_k[_p], r_k[_p + 1]), + src_c=(shk_b1, shk_b2), + ) + shk_a = shk_a1 + shk_a2 + shk_b = shk_b1 + shk_b2 + for offset in [16, 8, 4, 2, 1]: + shk_a += cute.arch.shuffle_sync_bfly(shk_a, offset=offset, mask=-1, mask_and_clamp=31) + shk_b += cute.arch.shuffle_sync_bfly(shk_b, offset=offset, mask=-1, mask_and_clamp=31) + vn_a = r_va - shk_a + vn_b = r_vb - shk_b + shq_a1 = cutlass.Float32(0.0) + shq_a2 = cutlass.Float32(0.0) + shq_b1 = cutlass.Float32(0.0) + shq_b2 = cutlass.Float32(0.0) + for _pi in cutlass.range_constexpr(VEC // 2): + _p = _pi * 2 + vnbk_a0, vnbk_a1 = cute.arch.mul_packed_f32x2((vn_a, vn_a), (r_bk[_p], r_bk[_p + 1])) + vnbk_b0, vnbk_b1 = cute.arch.mul_packed_f32x2((vn_b, vn_b), (r_bk[_p], r_bk[_p + 1])) + st[ra * VEC + _p], st[ra * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_decay[_p], r_decay[_p + 1]), + src_b=(st[ra * VEC + _p], st[ra * VEC + _p + 1]), + src_c=(vnbk_a0, vnbk_a1), + ) + st[rb * VEC + _p], st[rb * VEC + _p + 1] = cute.arch.fma_packed_f32x2( + src_a=(r_decay[_p], r_decay[_p + 1]), + src_b=(st[rb * VEC + _p], st[rb * VEC + _p + 1]), + src_c=(vnbk_b0, vnbk_b1), + ) + shq_a1, shq_a2 = cute.arch.fma_packed_f32x2( + src_a=(st[ra * VEC + _p], st[ra * VEC + _p + 1]), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_a1, shq_a2), + ) + shq_b1, shq_b2 = cute.arch.fma_packed_f32x2( + src_a=(st[rb * VEC + _p], st[rb * VEC + _p + 1]), + src_b=(r_q[_p], r_q[_p + 1]), + src_c=(shq_b1, shq_b2), + ) + # The q-dot feeds only the output, not the next token: its reduction waits until after the loop. + outs.append([shq_a1 + shq_a2, shq_b1 + shq_b2]) + # The golden token's state to the pool; a draft's record: the warp's row innovations (lane r, row r). + if cutlass.const_expr(tr == 0): + for r in cutlass.range_constexpr(ROWS): + for i in cutlass.range_constexpr(VEC): + pool.store(ssm.dtype(st[r * VEC + i]), idx=st_base + r * K + i * 32 + lane) + else: + if lane < ROWS: + tok.store( + cutlass.Float32(cutlass.select_(lane == 0, vn_a, vn_b)), + idx=CT_VN + ((tr - 1) * H + h) * V + row_a + lane, + ) + + for offset in [16, 8, 4, 2, 1]: + for tr in cutlass.range_constexpr(NT): + for r in cutlass.range_constexpr(ROWS): + outs[tr][r] += cute.arch.shuffle_sync_bfly( + outs[tr][r], offset=offset, mask=-1, mask_and_clamp=31 + ) + if lane == 0: + for tr in cutlass.range_constexpr(NT): + s_o.store(_bf16(outs[tr][0]), idx=tr * V_CTA + warp * ROWS + 0) + s_o.store(_bf16(outs[tr][1]), idx=tr * V_CTA + warp * ROWS + 1) + + # ---- Output gated RMSNorm over V: per-token partial sums of squares to every CTA of the cluster. + prims.barrier_cta_sync(0) + if tidx < NT * V_CTA: + t_out = tidx // V_CTA + x_o = s_o.load(idx=tidx) + ss = x_o * x_o + for off_ss in [V_CTA >> d for d in range(1, V_CTA.bit_length()) if V_CTA >> d > 0]: + ss = ss + cute.arch.shuffle_sync_bfly(ss, offset=off_ss, mask=-1, mask_and_clamp=31) + if tidx % V_CTA == 0: + # st.async into every CTA of the cluster (this one too), completing on its mailbox barrier: no cluster + # barrier, so nothing waits for this CTA's state stores to drain. + for r in cutlass.range_constexpr(V_SPLIT): + _st_async_f32( + _mapa_u32(s_ss.subview(vs * NT + t_out).data_ptr(), r), + ss, + _mapa_u32(ss_ready.data_ptr(), r), + ) + # Every peer's partials are in this CTA's shared memory once its mailbox completes. Each peer sends only after + # its recurrence, i.e. after it consumed its reads of the q/k conv history, so CTA 0's rewrite below cannot race + # them; and no CTA exits before its own mailbox has received every peer's partials. + while not _try_wait_cluster(ss_ready.data_ptr(), 0): + pass + if tidx < NT: + total = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(V_SPLIT): + total = total + s_ss.load(idx=r * NT + tidx) + s_rs.store(cute.math.rsqrt(total / V + eps), idx=tidx) + prims.barrier_cta_sync(0) + if tidx < NT * V_CTA: + t_y = tidx // V_CTA + v_y = tidx % V_CTA + z = s_og.load(idx=tidx) + gate = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.math.exp(-z, fastmath=True)) + y = s_o.load(idx=tidx) * s_rs.load(idx=t_y) * s_onw.load(idx=v_y) * gate + out.store(cutlass.BFloat16(y), idx=((row0 + t_y) * H + h) * V + v0 + v_y) + + # ---- The drafts' per-key records (beta * k and the decay), one CTA per head, after every peer's norm partials + # have arrived: each peer sends them after its recurrence, i.e. after its prologue read the last round's records. + if vs == 0: + for it_r in cutlass.range_constexpr((NUM_SPEC * K + THREADS - 1) // THREADS): + item_r = tidx + it_r * THREADS + if item_r < NUM_SPEC * K: + t_r = item_r // K + key_r = item_r % K + tok.store( + s_bk.load(idx=(t_r + 1) * K + key_r), + idx=CT_WB + (t_r * H + h) * K + key_r, + ) + tok.store( + s_dec.load(idx=(t_r + 1) * K + key_r), + idx=CT_WD + (t_r * H + h) * K + key_r, + ) + + # ---- Conv caches for the next round: column s = position s - 2 = row s + 1 of the position-indexed inputs. + if vs == 0: + if tidx < K: + for s in cutlass.range_constexpr(S): + cs_q.store(s_uq.load(idx=(s + 1) * K + tidx), idx=(slot * S + s) * HK + ch0 + tidx) + else: + ck_c = tidx - K + for s in cutlass.range_constexpr(S): + cs_k.store(s_uk.load(idx=(s + 1) * K + ck_c), idx=(slot * S + s) * HK + ch0 + ck_c) + for it in cutlass.range_constexpr((S * V_CTA + THREADS - 1) // THREADS): + item = tidx + it * THREADS + if item < S * V_CTA: + s_col = item // V_CTA + v_col = item % V_CTA + cs_v.store( + s_uv.load(idx=(s_col + 1) * V_CTA + v_col), + idx=(slot * S + s_col) * HK + ch0 + v0 + v_col, + ) + + +@cute.jit +def k3_fp8_kda_verify( + w_fb: cute.Tensor, # bf16 [H * K, K]: the f_b weight (out, in) + proj: cute.Tensor, # int32 words of the fused projection [T, proj_words] + g_ext: cute.Tensor, # int32 words of the unfused f_b output [T, H * K / 2] (FOLD_FB False) + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + a_log: cute.Tensor, + dt_bias: cute.Tensor, + onorm_w: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + ssm: cute.Tensor, + state_tok: cute.Tensor, + slots: cute.Tensor, + pending: cute.Tensor, + out: cute.Tensor, + proj_words: cutlass.Int32, + n_req: cutlass.Int32, + ssm_stride: cutlass.Int64, + H: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + FOLD_FB: cutlass.Constexpr[bool], + USE_PDL: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + """One launch of the verify kernel for ``n_req`` requests of 1 + NUM_SPEC tokens each.""" + # W_fb [H K, K] as five TMA dimensions (64-element column chunk, row, chunk index, 1, 1): one call lands the + # head's 128 rows, both 128-byte-swizzled halves; strides in 16-byte units. + tma_wfb = cuda.create_tensor_map_tiled( + global_address=w_fb.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[TMA_K_BOX, H * K, K // TMA_K_BOX, 1, 1], + global_strides=[(K * BF16_BYTES) // 16, (TMA_K_BOX * BF16_BYTES) // 16, (H * K * K * BF16_BYTES) // 16, + (H * K * K * BF16_BYTES) // 16], + box_dims=[TMA_K_BOX, CTA_M, TMA_COPY_ITERS, 1, 1], + swizzle=cuda.TensorMapSwizzle.s128b, + ) # fmt: skip + # The projection rows as bf16 [T, 2 proj_words]: an 8-token x 64-column box of f_a per call. + tma_fa = cuda.create_tensor_map_tiled( + global_address=proj.iterator.toint(), + dtype=cutlass.BFloat16, + global_dims=[proj_words * 2, n_req * (NUM_SPEC + 1)], + global_strides=[(proj_words * 4) // 16], + box_dims=[TMA_K_BOX, MMA_N], + swizzle=cuda.TensorMapSwizzle.s128b, + ) + k3_fp8_kda_verify_kernel( + tma_wfb, tma_fa, proj, g_ext, w_q, w_k, w_v, a_log, dt_bias, onorm_w, cs_q, cs_k, cs_v, ssm, state_tok, slots, + pending, out, proj_words, ssm_stride, H, NUM_SPEC, lower_bound, scale, eps, FOLD_FB, USE_PDL, + ).launch( + grid=(H, n_req, V_SPLIT), block=(THREADS, 1, 1), cluster=(1, 1, V_SPLIT), stream=stream, use_pdl=USE_PDL, + ) # fmt: skip From fbde4c91b39c13c54a8736e400751428bb790508 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:05:19 +0800 Subject: [PATCH 8/8] Enable K3 hybrid FP4 KV cache allocation (By Agent) Compose the native FP4 MLA layout with the hybrid V2 cache manager so quantized KV and recurrent-state pools share the existing lifecycle. Match fused Q-B quantization to the selected FP8 GEMM recipe and cover the cache factory, real hybrid allocation and precision combinations. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_experimental/modeling_v2/README.md | 10 ++ .../kimi_k3_vl/kimi_k3_nvfp4/attention.py | 12 +- .../kimi_k3_nvfp4/implementation.py | 23 +++- tensorrt_llm/_torch/pyexecutor/_util.py | 13 ++- .../_torch/pyexecutor/config_utils.py | 3 +- .../kv_cache/mamba_cache_manager.py | 49 ++++++++- .../_torch/pyexecutor/model_loader.py | 3 +- .../test_modeling_v2_k3_cache_precisions.py | 39 ++++++- .../test_modeling_v2_k3_fp8_fused.py | 18 ++- .../_torch/modeling_v2/ssm/_kda_cells.py | 8 +- .../modeling_v2/test_modeling_v2_claims.py | 1 + .../test_modeling_v2_kimi_k3_drift.py | 103 ++++++++++++++++++ 12 files changed, 252 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index e66a73d97530..ab2a8829f03d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -23,10 +23,20 @@ V2 cache manager; its high-precision tail and packed scales retain the native layout. The KDA fused decode and verify accept FP32 or BF16 recurrent-state storage, with FP32 recurrence arithmetic and FP32 per-token draft records. These KV and SSM storage choices are independent of the NVFP4 expert weights. +The fused Q-B projection follows the selected GEMM backend's FP8 recipe, +including weight resmoothing and UE8M0 activation scales where required. + +| MLA KV storage | KDA SSM storage | KDA recurrence / draft records | +|---|---|---| +| BF16 | FP32 or BF16 | FP32 | +| FP8 E4M3 | FP32 or BF16 | FP32 | +| NVFP4 | FP32 or BF16 | FP32 | + For NVFP4 KV and BF16 recurrent-state storage: ```yaml modeling_v2: require +enable_chunked_prefill: false kv_cache_config: dtype: nvfp4 tokens_per_block: 128 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py index 3b203ec37ccf..b6e8b19a3226 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py @@ -176,6 +176,7 @@ def _mla_qkv( QA_EPS: tl.constexpr, KV_EPS: tl.constexpr, STORE_PAGED: tl.constexpr = True, + QB_UE8M0: tl.constexpr = False, ): head, chunk = tl.program_id(0), tl.program_id(1) rows = chunk * 16 + tl.arange(0, 16) @@ -198,6 +199,8 @@ def _mla_qkv( a *= rms[:, None] * tl.load(QA_NORM + cols)[None, :].to(tl.float32) a = a.to(tl.bfloat16).to(tl.float32) a_scale = tl.maximum(tl.max(tl.abs(a), 1) / 448.0, 1.0e-12) + if QB_UE8M0: + a_scale = tl.exp2(tl.ceil(tl.log2(a_scale))) a8 = (a / a_scale[:, None]).to(tl.float8e4nv) b = tl.load(QB + qb_rows[None, :] * 1536 + cols[:, None], n[None, :] < 192, 0.0) bs = tl.load(QB_SCALE + (qb_rows // 128) * QB_SCALE_STRIDE + block, n < 192, 0) @@ -262,6 +265,7 @@ def mla_qkv( tokens_per_request: int, qa_eps: float, kv_eps: float, + qb_ue8m0: bool = False, ) -> torch.Tensor: """Fuse Q-A norm, FP8 Q-B/absorption and the paged latent-cache update.""" h = kb.shape[0] @@ -290,6 +294,7 @@ def mla_qkv( tokens_per_request, qa_eps, kv_eps, + QB_UE8M0=qb_ue8m0, num_warps=8, ) return out @@ -312,6 +317,7 @@ def _qkv_fake( tokens_per_request, qa_eps, kv_eps, + qb_ue8m0=False, ): return ag.new_empty((ag.shape[0], kb.shape[0] * 576)) @@ -327,6 +333,7 @@ def mla_project_qkv( kv_norm: torch.Tensor, qa_eps: float, kv_eps: float, + qb_ue8m0: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Produce absorbed Q and normalized latent KV for quantized cache kernels.""" h = kb.shape[0] @@ -357,13 +364,16 @@ def mla_project_qkv( qa_eps, kv_eps, STORE_PAGED=False, + QB_UE8M0=qb_ue8m0, num_warps=8, ) return query, latent @mla_project_qkv.register_fake -def _project_qkv_fake(ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, qa_eps, kv_eps): +def _project_qkv_fake( + ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, qa_eps, kv_eps, qb_ue8m0=False +): return ag.new_empty((ag.shape[0], kb.shape[0] * 576)), ag.new_empty((ag.shape[0], 576)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py index 2444f19694bd..328734c2ed70 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py @@ -55,6 +55,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig +from tensorrt_llm.quantization.utils.fp8_utils import resmooth_to_fp8_e8m0 from . import attention as _fp8_attention from . import decode_comm as _decode_comm @@ -161,6 +162,7 @@ "tensorrt_llm._torch.models.modeling_dspark.draft_is_embedded_in_target", "tensorrt_llm._torch.models.modeling_dflash.DFlashForCausalLM", "tensorrt_llm._torch.pyexecutor.config_utils.is_mla", + "tensorrt_llm.quantization.utils.fp8_utils.resmooth_to_fp8_e8m0", ) # The K3 decode kernels' bounds: the token-count kernels take one tile of DECODE_MAX_TOKENS rows; the request-aware @@ -2275,10 +2277,12 @@ def forward_verify( return self._store_core(core, output) def _k3_verify(self, x, num_steps, layer_cache, ssm_pool, slot_indices) -> torch.Tensor: - """The core output ``[N T, H, 128]`` of N requests of T verify tokens: ``ssm/k3_kda_attn`` for one request of - 8 tokens, else ``ssm/k3_kda_verify`` on the projection's rows (the ``kda_proj`` decode GEMV site where it - takes them). Each slot's state after the golden token, its drafts' states and its conv window are written in - place.""" + """Verify N requests of T tokens from FP8 QKVG and BF16 gate projections. + + Return ``[N T, H, 128]`` and update each slot's committed state, FP32 + draft records and convolution window. The recurrence uses FP32 with + either FP32 or BF16 committed-state storage. + """ num_requests = x.shape[0] // num_steps slots = slot_indices[:num_requests] w_q, w_k, w_v = self._get_mtp_conv_weights() @@ -2314,10 +2318,11 @@ class K3DecodeMLA(KimiK3MLAAttention): def __init__(self, **kwargs) -> None: super().__init__(**kwargs) self._fp8_qb_weight = self._fp8_qb_scale = None + self._fp8_qb_ue8m0 = False self._fp8_decode_ready = False def prepare_fp8_decode_weights(self) -> None: - """Retain checkpoint Q-B codes before the GEMM backend transforms them.""" + """Prepare FP8 Q-B codes and scales with the selected GEMM's quantization recipe.""" linear = self.q_b_proj if linear.weight.dtype != torch.float8_e4m3fn: raise ValueError("NVFP4 fused MLA requires checkpoint FP8 Q-B projections") @@ -2329,6 +2334,12 @@ def prepare_fp8_decode_weights(self) -> None: ) self._fp8_qb_weight = linear.weight.detach().clone().contiguous() self._fp8_qb_scale = scale.detach().clone().contiguous() + self._fp8_qb_ue8m0 = not ( + getattr(linear, "use_cute_dsl_blockscaling_mm", False) + or getattr(linear, "disable_deep_gemm", False) + ) or (linear.weight.is_cuda and torch.cuda.get_device_capability() == (10, 7)) + if self._fp8_qb_ue8m0: + resmooth_to_fp8_e8m0(self._fp8_qb_weight, self._fp8_qb_scale) def post_load_weights(self) -> None: super().post_load_weights() @@ -2413,6 +2424,7 @@ def forward( view["tokens_per_request"], float(self.q_a_layernorm.variance_epsilon), float(self.kv_a_layernorm.variance_epsilon), + self._fp8_qb_ue8m0, ) core = _fp8_attention.mla_attn( fused_q, @@ -2449,6 +2461,7 @@ def _quantized_decode( self.kv_a_layernorm.weight, float(self.q_a_layernorm.variance_epsilon), float(self.kv_a_layernorm.variance_epsilon), + self._fp8_qb_ue8m0, ) q_pe = query.view(-1, self.num_heads_tp, 576)[..., 512:] cu_q = torch.empty(metadata.num_seqs + 1, dtype=torch.int32, device=ag.device) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9351a49a0fc5..b98b3de2da0a 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -209,9 +209,12 @@ def get_kv_cache_manager_cls( "FP4 MLA disaggregated serving requires the follow-up " "Python NIXL integration.") if is_hybrid_linear(config): - raise NotImplementedError( - "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not " - "support hybrid linear-attention models.") + if use_py_mamba_cache_manager() or os.environ.get( + "TLLM_MAMBA_MANAGER_PREFERENCE"): + raise ValueError( + "Hybrid FP4 MLA requires MambaHybridCacheManagerV2; unset " + "TRTLLM_USE_PY_MAMBA / TLLM_MAMBA_MANAGER_PREFERENCE.") + return MambaHybridCacheManagerV2 from ..attention.backends.fp4_mla.cache_manager import \ Fp4MlaKVCacheManagerV2 @@ -3353,6 +3356,10 @@ def _create_kv_cache_manager( # states in place, replacing the intermediate-buffer + promotion # flow for KDA layers. kda_extra_kwargs = {} + if (quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache() and issubclass( + kv_cache_manager_cls, MambaHybridCacheManagerV2)): + kda_extra_kwargs["pretrained_config"] = config kda_replay_manager_types = (MixedMambaHybridCacheManager, MambaHybridCacheManagerV2) if (spec_config is not None diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index a9b2b16bfbe3..8abfd3d2fbf2 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -440,8 +440,7 @@ def supports_fp4_mla_attention(model_config: "ModelConfig") -> bool: """Whether the model uses the dedicated dense TRTLLM FP4 MLA path.""" return (is_mla(model_config.pretrained_config) and model_config.attn_backend == "TRTLLM" - and model_config.sparse_attention_config is None - and not is_hybrid_linear(model_config.pretrained_config)) + and model_config.sparse_attention_config is None) def uses_fp4_mla_attention(model_config: "ModelConfig") -> bool: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 8ecddb4b29d6..f93f36a11110 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -2198,6 +2198,21 @@ def _estimate_mamba_hybrid_cache_cost( num_layers=local_attention_layers, **kwargs, ) if local_attention_layers > 0 else 0) + attention_fixed = 0 + quant_config = getattr(model_config, "quant_config", None) + if (local_attention_layers > 0 and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache() + and getattr(model_config.pretrained_config, "kv_lora_rank", None)): + from ...attention.backends.fp4_mla.cache_manager import \ + Fp4MlaV2CacheLayoutPolicy + attention_slope, attention_fixed = Fp4MlaV2CacheLayoutPolicy.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + tokens_per_block=tokens_per_block, + spec_config=spec_config, + max_batch_size=max_batch_size + num_reserved_dummy_slots, + ) state_bytes_per_rank = (local_mamba_layers * params.get_states_bytes_per_layer(mapping)) if not is_draft: @@ -2233,7 +2248,7 @@ def _estimate_mamba_hybrid_cache_cost( else: partial_attention_slots = (max_resident_sequences * unaligned_fixed_rules) - intercept = (fixed_state_slots * state_bytes_per_rank + + intercept = (attention_fixed + fixed_state_slots * state_bytes_per_rank + partial_attention_slots * attention_block_bytes) if interval is None: @@ -3252,6 +3267,22 @@ def __init__( self._setup_mtp_intermediate_states(spec_config, max_batch_size) + self._attention_cache_layout_policy = None + if dtype == DataType.NVFP4 and kv_cache_type == CacheTypeCpp.SELFKONLY: + from ...attention.backends.fp4_mla.cache_manager import \ + Fp4MlaV2CacheLayoutPolicy + kv_cache_config = Fp4MlaV2CacheLayoutPolicy.configure_manager( + self, + kv_cache_config, + kv_cache_type, + dtype=dtype, + tokens_per_block=tokens_per_block, + head_dim=head_dim, + pretrained_config=kwargs.get("pretrained_config"), + spec_config=spec_config, + is_disagg=kwargs.get("is_disagg", False), + ) + kv_cache_config = kv_cache_config.model_copy(deep=True) if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: block_reuse_config = kv_cache_config.block_reuse_config @@ -3280,6 +3311,8 @@ def __init__( num_reserved_index_slots=self._num_reserved_dummy_slots, **kwargs, ) + if self._attention_cache_layout_policy is not None: + self._attention_cache_layout_policy.finalize_manager(self) self.mamba_layer_offsets = { layer_id: offset @@ -3866,11 +3899,16 @@ def get_cache_size_per_token(model_config, ) def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: + if local_layer_idx >= len(self.pp_layers): + return False return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: layer_id = int(self.impl.layer_grouping[pool_id][0]) + if (getattr(self, "_attention_cache_layout_policy", None) is not None + and layer_id >= self.num_local_layers): + return Role.MLA_HP_TAIL, None if self._is_local_mamba_layer(layer_id): return MambaRole.SSM_STATE, None return super()._get_pool_roles(pool_id) @@ -4223,7 +4261,7 @@ def _build_cache_config( for _ in range(ssm_floor_slots) ]), ] - return replace( + result = replace( config, cache_tiers=cache_tiers, layers=layers, @@ -4234,6 +4272,10 @@ def _build_cache_config( # attempted, while the runtime config still needs the invariant. commit_min_snapshot=True, ) + if getattr(self, "_attention_cache_layout_policy", None) is not None: + result = self._attention_cache_layout_policy._build_cache_config( + self, result) + return result def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, state_shape: List[int]) -> torch.Tensor: @@ -4377,6 +4419,9 @@ def _setup_replay_buffers(self, spec_config) -> None: def _attention_cache_bytes_per_token(self) -> int: # Mamba layers have zero KV heads, so the generic calculation naturally # returns only bytes owned by local attention layers. + if getattr(self, "_attention_cache_layout_policy", None) is not None: + return self._attention_cache_layout_policy.get_cache_bytes_per_token( + self) return super().get_cache_bytes_per_token() def get_cache_bytes_per_token(self) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f2f4e12a855..d3e30922596e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -219,8 +219,7 @@ def validate_fp4_mla_config(model_config: ModelConfig, and model_config.quant_config.quant_mode.has_fp4_kv_cache()): return if not supports_fp4_mla_attention(model_config): - raise ValueError("Dense FP4 MLA requires the TRTLLM attention backend; " - "hybrid linear attention is not supported.") + raise ValueError("Dense FP4 MLA requires the TRTLLM attention backend.") if model_config.mapping.cp_size != 1: raise ValueError("FP4 MLA does not support context parallelism.") if llm_args.kv_cache_config.use_kv_cache_manager_v2 is False: diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py index 76fae98fcdbf..4937c64cd26a 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py @@ -22,6 +22,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType from tensorrt_llm._torch.utils import model_extra_attrs @@ -36,12 +37,18 @@ ) -@pytest.mark.parametrize("cache_dtype", ["bf16", "fp8", "nvfp4"]) +@pytest.mark.parametrize( + "cache_dtype", ["bf16", "fp8", "nvfp4", "nvfp4_ssm_fp32", "nvfp4_ssm_bf16"] +) @pytest.mark.parametrize("batch,tokens", [(1, 1), (2, 1), (4, 1), (1, 4), (2, 4), (4, 4)]) -@torch.inference_mode() +@torch.no_grad() def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int, tokens: int) -> None: import tensorrt_llm._torch.modules.kimi_k3_mla as mla_package + hybrid = cache_dtype.startswith("nvfp4_ssm_") + state_dtype = torch.bfloat16 if cache_dtype.endswith("bf16") else torch.float32 + if hybrid: + cache_dtype = "nvfp4" torch.manual_seed(451) monkeypatch.setattr(mla_package, "KimiK3MLAAttention", K3DecodeMLA) config = KimiLinearConfig( @@ -96,22 +103,37 @@ def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int, tok ) load_mla_checkpoint(model, projections, "") runtime = model.model.layers[0].self_attn.mixer - runtime._fp8_qb_weight = weights["q_b_proj"]["weight"].cuda().contiguous() - runtime._fp8_qb_scale = weights["q_b_proj"]["weight_scale"].cuda().contiguous() + runtime.prepare_fp8_decode_weights() for module in runtime.modules(): if callable(getattr(module, "post_load_weights", None)): module.post_load_weights() context_length = 128 page_size = 128 if cache_dtype == "nvfp4" else 64 manager_class = Fp4MlaKVCacheManagerV2 if cache_dtype == "nvfp4" else KVCacheManagerV2 + manager_kwargs = {} + if hybrid: + manager_class = MambaHybridCacheManagerV2 + manager_kwargs = dict( + mamba_d_state=128, + mamba_d_conv=4, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_head_dim=128, + mamba_num_layers=1, + mamba_layer_mask=[False, True], + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=state_dtype, + layer_mask=[True, False], + conv_state_layout="q_k_v", + ) manager = manager_class( - KvCacheConfig( + kv_cache_config=KvCacheConfig( max_tokens=256 * batch, dtype="auto" if cache_dtype == "bf16" else cache_dtype, enable_block_reuse=False, host_cache_size=0, ), - CacheTypeCpp.SELFKONLY, + kv_cache_type=CacheTypeCpp.SELFKONLY, num_layers=1, num_kv_heads=1, head_dim=576, @@ -123,10 +145,15 @@ def test_fp8_mla_cache_precisions(monkeypatch, cache_dtype: str, batch: int, tok max_num_tokens=batch * context_length, pretrained_config=config, spec_config=MTPDecodingConfig(max_draft_len=tokens - 1) if tokens > 1 else None, + **manager_kwargs, ) try: requests = manager.add_dummy_requests(list(range(batch)), [context_length + tokens] * batch) assert requests is not None + if hybrid: + assert manager.get_ssm_states(1).dtype == state_dtype + assert manager.get_buffers(1) is None + assert manager.get_fp4_mla_page_table_spec(0).hp_is_paged for request in requests: request.state = LlmRequestState.GENERATION_IN_PROGRESS request.prompt_len = context_length diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py index 2ff3e5cb8fc9..be5a511dab17 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py @@ -33,9 +33,11 @@ def _rand(*shape, scale=1.0): return (torch.randn(shape, device="cuda") * scale).bfloat16() -def _quantize(x): +def _quantize(x, ue8m0=False): groups = x.float().unflatten(-1, (-1, 128)) scale = groups.abs().amax(-1, keepdim=True).clamp_min(4.48e-10) / 448 + if ue8m0: + scale = torch.exp2(torch.ceil(torch.log2(scale))) codes = (groups / scale).to(torch.float8_e4m3fn) return (codes.float() * scale).flatten(-2) @@ -122,7 +124,8 @@ def test_fp8_kda_fused_decode(batch, state_dtype): @pytest.mark.parametrize("batch,tokens", [(1, 1), (2, 1), (4, 1), (1, 4), (4, 8)]) -def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): +@pytest.mark.parametrize("qb_ue8m0", [False, True]) +def test_fp8_mla_fused_query_cache_and_attention(batch, tokens, qb_ue8m0): torch.manual_seed(422) torch.backends.cuda.matmul.allow_tf32 = False h, m = 6, batch * tokens @@ -130,6 +133,8 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): qa_norm, kv_norm = _rand(1536, scale=0.2), _rand(512, scale=0.2) qb = _rand(h * 192, 1536, scale=0.05).to(torch.float8_e4m3fn) qb_scale = torch.rand(9, 12, device="cuda") * 0.2 + 0.1 + if qb_ue8m0: + qb_scale = torch.exp2(torch.ceil(torch.log2(qb_scale))) kb = _rand(h, 512, 128, scale=0.05).to(torch.float8_e4m3fn) kb_scale = torch.rand(h, 4, 1, device="cuda") * 0.2 + 0.1 vb = _rand(h, 128, 512, scale=0.05).to(torch.float8_e4m3fn) @@ -152,7 +157,7 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): qa * torch.rsqrt(qa.square().mean(-1, keepdim=True) + 1.0e-6) * qa_norm.float() ).bfloat16() qb_dequant = qb.float() * qb_scale.repeat_interleave(128, 0).repeat_interleave(128, 1) - q = (_quantize(qa) @ qb_dequant.T).bfloat16().view(m, h, 192) + q = (_quantize(qa, qb_ue8m0) @ qb_dequant.T).bfloat16().view(m, h, 192) kb_dequant = kb.float() * kb_scale.repeat_interleave(128, 1).repeat_interleave(128, 2) absorbed = torch.einsum("mhk,hvk->mhv", _quantize(q[..., :128]), kb_dequant).bfloat16() expected_q = torch.cat((absorbed, q[..., 128:]), -1).flatten(1) @@ -179,11 +184,12 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): tokens, 1.0e-6, 1.0e-6, + qb_ue8m0, ) torch.testing.assert_close(actual_q, expected_q, atol=0.002, rtol=0.02) torch.testing.assert_close(pool, expected_pool, atol=0.002, rtol=0.02) projected_q, latent_cache = _ops().mla_project_qkv( - ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, 1.0e-6, 1.0e-6 + ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, 1.0e-6, 1.0e-6, qb_ue8m0 ) torch.testing.assert_close(projected_q, actual_q, atol=0, rtol=0) latent = ag[:, 1536:2048].float() @@ -230,7 +236,7 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): graph_projected_q, graph_latent = _ops().mla_project_qkv( - ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, 1.0e-6, 1.0e-6 + ag, qa_norm, qb, qb_scale, kb, kb_scale, kv_norm, 1.0e-6, 1.0e-6, qb_ue8m0 ) graph_projected_output = _ops().mla_output(latent_output, vb, vb_scale, ag, 2112) graph_q = _ops().mla_qkv( @@ -249,6 +255,7 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): tokens, 1.0e-6, 1.0e-6, + qb_ue8m0, ) graph_output = _ops().mla_attn( graph_q, @@ -288,6 +295,7 @@ def test_fp8_mla_fused_query_cache_and_attention(batch, tokens): tokens, 1.0e-6, 1.0e-6, + qb_ue8m0, ) updated_output = _ops().mla_attn( updated_q, diff --git a/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py index b3d18f3f3eb6..a4a2ec0c8fcd 100644 --- a/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py +++ b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py @@ -7,8 +7,8 @@ Not a test file (no ``test_`` prefix): the four entry tests import it by name, as ``comm/`` imports ``_rank_job``. Shapes: 6 heads, K = V = 128, conv width 4 (the TP16 rank slice); conv states bf16 in the ``[q | k | v]`` layout, -SSM states fp32. The manager coalesces each slot's per-layer states, so every per-layer view is a strided slot view, -the layout the model hands these ops. +SSM states fp32 by default or bf16 when requested. The manager coalesces each slot's per-layer states, so every +per-layer view is a strided slot view, the layout the model hands these ops. """ from __future__ import annotations @@ -105,8 +105,8 @@ def request_slots(mgr, count: int, first_id: int) -> torch.Tensor: def layer_pools(mgr, layer: int) -> dict: """The manager's views of one KDA layer's pools: ``conv`` bf16 [slots, 3 HK, W - 1] (q | k | v channels) and - ``ssm`` fp32 [slots, H, V, K] (plain decode); with the replay caches also ``cs_q`` / ``cs_k`` / ``cs_v`` fp32 - [slots, HK, W - 1 + num_spec] (dim-contiguous), ``state_tok`` fp32 [slots, 3, num_spec, H, K] and ``pending``, + ``ssm`` fp32 or bf16 [slots, H, V, K] (plain decode); with the replay caches also ``cs_q`` / ``cs_k`` / ``cs_v`` + fp32 [slots, HK, W - 1 + num_spec] (dim-contiguous), ``state_tok`` fp32 [slots, 3, num_spec, H, K] and ``pending``, the accepted-draft record every layer shares (``prev_num_accepted_tokens``).""" pools = {"conv": mgr.get_conv_states(layer), "ssm": mgr.get_ssm_states(layer)} if mgr.use_kda_replay_update: diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 5854b5fb9e4a..d272c9ecdbaf 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -166,6 +166,7 @@ def test_targets_only_import_their_declared_family_implementation(): _NON_COMPUTE_IMPORTS = frozenset( { "AllReduceStrategy", + "AttentionInputType", "AttentionMetadata", "AuxStreamType", "MambaHybridCacheManagerV2", diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py index 2be90a8d7255..f0d25788d748 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_kimi_k3_drift.py @@ -16,6 +16,80 @@ pytestmark = pytest.mark.cpu_only +def test_nvfp4_k3_selects_hybrid_v2_cache_manager(monkeypatch): + from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.pyexecutor._util import ( + _create_kv_cache_manager, + get_kv_cache_manager_cls, + ) + from tensorrt_llm._torch.pyexecutor.config_utils import uses_fp4_mla_attention + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + ) + from tensorrt_llm._torch.pyexecutor.model_loader import validate_fp4_mla_config + from tensorrt_llm.llmapi import KvCacheConfig + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantConfig + + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + config = ModelConfig( + pretrained_config=KimiLinearConfig( + num_hidden_layers=2, + linear_attn_config={ + "kda_layers": [1], + "full_attn_layers": [2], + "head_dim": 128, + "num_heads": 6, + "short_conv_kernel_size": 4, + }, + q_lora_rank=1536, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + ), + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4), + ) + cache = KvCacheConfig(dtype="nvfp4", use_kv_cache_manager_v2=True) + assert uses_fp4_mla_attention(config) + validate_fp4_mla_config( + config, SimpleNamespace(kv_cache_config=cache, enable_chunked_prefill=False) + ) + assert get_kv_cache_manager_cls(config, cache) is MambaHybridCacheManagerV2 + captured = {} + + class RecordingManager(MambaHybridCacheManagerV2): + def __init__(self, *args, **kwargs): + captured.update(kwargs) + + _create_kv_cache_manager( + model_engine=None, + kv_cache_manager_cls=RecordingManager, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + kv_cache_config=cache, + tokens_per_block=128, + max_seq_len=256, + max_batch_size=4, + spec_config=None, + sparse_attention_config=None, + max_num_tokens=256, + max_beam_width=1, + kv_connector_manager=None, + model_config=config, + dtype=torch.bfloat16, + ) + assert captured["pretrained_config"] is config.pretrained_config + with pytest.raises(ValueError, match="use_kv_cache_manager_v2"): + get_kv_cache_manager_cls( + config, cache.model_copy(update={"use_kv_cache_manager_v2": False}) + ) + monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1") + with pytest.raises(ValueError, match="MambaHybridCacheManagerV2"): + get_kv_cache_manager_cls(config, cache) + + @pytest.mark.parametrize("moe_tp,moe_ep", [(4, 4), (16, 1)]) def test_sm100_mxfp4_keeps_its_own_model(moe_tp: int, moe_ep: int) -> None: module = importlib.import_module( @@ -37,6 +111,8 @@ def test_fp8_mla_preserves_raw_query_weights_before_backend_transformation() -> weight_scale=scales, out_features=128, in_features=128, + use_cute_dsl_blockscaling_mm=True, + disable_deep_gemm=False, ) attention.prepare_fp8_decode_weights() codes.zero_() @@ -46,6 +122,33 @@ def test_fp8_mla_preserves_raw_query_weights_before_backend_transformation() -> torch.testing.assert_close(attention._fp8_qb_scale, torch.full((1, 1), 0.125)) +@pytest.mark.parametrize( + "backend,ue8m0", + [ + ({}, True), + ({"use_cute_dsl_blockscaling_mm": False, "disable_deep_gemm": False}, True), + ({"use_cute_dsl_blockscaling_mm": True, "disable_deep_gemm": False}, False), + ({"use_cute_dsl_blockscaling_mm": False, "disable_deep_gemm": True}, False), + ], +) +def test_fp8_mla_query_recipe_handles_weight_read_specialization(monkeypatch, backend, ue8m0): + attention = modeling.K3DecodeMLA.__new__(modeling.K3DecodeMLA) + torch.nn.Module.__init__(attention) + codes = torch.ones(128, 128).to(torch.float8_e4m3fn) + scales = torch.full((1, 1), 0.125) + attention.q_b_proj = SimpleNamespace( + weight=codes, weight_scale=scales, out_features=128, in_features=128, **backend + ) + calls = [] + monkeypatch.setattr(modeling, "resmooth_to_fp8_e8m0", lambda w, s: calls.append((w, s))) + attention.prepare_fp8_decode_weights() + assert attention._fp8_qb_ue8m0 == ue8m0 + assert len(calls) == int(ue8m0) + if calls: + assert calls[0][0].data_ptr() != codes.data_ptr() + assert calls[0][1].data_ptr() != scales.data_ptr() + + @pytest.mark.parametrize("sm", [(10, 3), (10, 7)]) def test_nvfp4_kda_selects_rubin_projection_backend(monkeypatch, sm) -> None: from tensorrt_llm._torch.model_config import ModelConfig