diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 7dc176ff06da..ab2a8829f03d 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,43 @@ 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 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 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. +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 + 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. + ## Using it ```python diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_comm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_comm.py new file mode 100644 index 000000000000..f5327daf9474 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/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/decode_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_gemv.py new file mode 100644 index 000000000000..301b709643b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/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/decode_moe.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe.py new file mode 100644 index 000000000000..43f643d49f40 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/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/decode_moe_tp16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe_tp16.py new file mode 100644 index 000000000000..b797c4fe6108 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/decode_moe_tp16.py @@ -0,0 +1,477 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# 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. +""" + +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, +) + +# 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, +) +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 + +# 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) + + +@dataclass(eq=False) +class K3DecodeMoe: + # 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`.""" + + head: K3MoeHeadWorkspace + small: K3MoeState + wide: K3MoeWideState + mnnvl: MnnvlWorkspace + exchange: Optional[K3LatentExchange] = None + # The one- and two-token engines' workspaces (``create`` builds both). + m1: Optional[K3MoeM1State] = None + m2: Optional[K3MoeM2State] = None + + # 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), + ) + + +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"] + # 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}"] + 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 + # The one- and two-token engines' handles. + m1: K3MoeM1Layer + m2: K3MoeM2Layer + + @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, + # The engines' handles; no wide build. + small=state.small.layer(*weights), + wide=None, + m1=state.m1.layer(*weights), + m2=state.m2.layer(*weights), + ) + + def warm_up(self, moe: nn.Module) -> None: + # 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.""" + 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 + # 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) + 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: + # 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``.""" + rows = hidden_states.shape[0] + if hidden_states.dtype != torch.bfloat16: + return False + # No wide step (k3_moe's wide build does not fit 896 local experts). + return 0 < rows <= 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 + # 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)) + 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() + + # 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) + + 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_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/__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..b6e8b19a3226 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/attention.py @@ -0,0 +1,548 @@ +# 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, + 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) + 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) + 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) + 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: + 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) + 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, + 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, + valid[:, None], + ) + + +@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, + qb_ue8m0: bool = False, +) -> 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, + QB_UE8M0=qb_ue8m0, + 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, + qb_ue8m0=False, +): + 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, + qb_ue8m0: bool = False, +) -> 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, + 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, qb_ue8m0=False +): + 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, + 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..328734c2ed70 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/implementation.py @@ -0,0 +1,3514 @@ +# 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.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 +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 tensorrt_llm.quantization.utils.fp8_utils import resmooth_to_fp8_e8m0 + +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 +from .verify import k3_fp8_kda_verify + +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_fp8_kda_verify", + "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", + "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", + "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 +# 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 + 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 + 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: + 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 + + @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" + ) + 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: + """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() + 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_fp8_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_qb_ue8m0 = False + self._fp8_decode_ready = False + + def prepare_fp8_decode_weights(self) -> None: + """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") + 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() + 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() + 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 + 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_ragged_verify", False) + 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 + + 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) + 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, + 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), + self._fp8_qb_ue8m0, + ) + 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 _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), + 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) + 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( + (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, + 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], + 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/verify.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4/verify.py new file mode 100644 index 000000000000..cc103be29262 --- /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 . import 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_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) + 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/_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 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/__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..d98271c207ac --- /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 ..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) + + +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..723db3326e18 --- /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 ..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/__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..f48ed5103338 --- /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_moetp4ep4.""" 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..d1b9d9269d56 --- /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,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 NVFP4, SM107, attention TP16 and MoE TP4 x EP4.""" + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +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) + + +def _check_construction(model_config: ModelConfig) -> None: + _check_target_construction(model_config, _TARGET) + + +class ModelingV2KimiK3Nvfp4Sm107Tp16Moetp4ep4(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_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..723db3326e18 --- /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,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 ..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/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 854616aebd9f..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 one GPU architecture these 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. # @@ -64,6 +59,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. @@ -116,18 +140,58 @@ 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 mxfp4_on_sm100 and parallel == "tp16_moetp4ep4": + 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 + 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_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, ) @@ -141,6 +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 " - "experts split moe_tp 4 x moe_ep 4 or, set explicitly, moe_tp 16 x moe_ep 1" + "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 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/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/spec_worker.py new file mode 100644 index 000000000000..abc1a1388213 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/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/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/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_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/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/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/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index c794e45d5903..bf8bb8fcc3ba 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -503,6 +503,8 @@ 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 + - 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/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/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_cache_precisions.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py new file mode 100644 index 000000000000..4937c64cd26a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_cache_precisions.py @@ -0,0 +1,276 @@ +# 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 copy +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 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 +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 +from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig +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", "nvfp4_ssm_fp32", "nvfp4_ssm_bf16"] +) +@pytest.mark.parametrize("batch,tokens", [(1, 1), (2, 1), (4, 1), (1, 4), (2, 4), (4, 4)]) +@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( + 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=torch.cuda.get_device_capability() == (10, 7), + 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.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( + 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, + ), + kv_cache_type=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, + 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 + 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 tokens, 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, + 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 + tokens, 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[:, 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) + 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, + ) + 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( + 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_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/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..be5a511dab17 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_k3_fp8_fused.py @@ -0,0 +1,315 @@ +# 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, 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) + + +@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)]) +@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 + 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 + 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) + 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_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) + 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, + 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, qb_ue8m0 + ) + 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): + 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) + 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, qb_ue8m0 + ) + graph_projected_output = _ops().mla_output(latent_output, vb, vb_scale, ag, 2112) + 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, + qb_ue8m0, + ) + 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) + 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( + 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, + qb_ue8m0, + ) + 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/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/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/ssm/_kda_cells.py b/tests/unittest/_torch/modeling_v2/ssm/_kda_cells.py index 3abd20fbcc5f..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 @@ -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, @@ -100,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/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..5c707480343b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/ssm/test_modeling_v2_k3_bf16_verify.py @@ -0,0 +1,88 @@ +# 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.models.kimi_k3_vl.kimi_k3_nvfp4.verify import ( + k3_fp8_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_fp8_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_claims.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py index 59d0f48a8559..d272c9ecdbaf 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,17 @@ 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", + "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"): @@ -143,6 +147,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 " @@ -156,6 +166,7 @@ def test_targets_do_not_share_files(): _NON_COMPUTE_IMPORTS = frozenset( { "AllReduceStrategy", + "AttentionInputType", "AttentionMetadata", "AuxStreamType", "MambaHybridCacheManagerV2", @@ -189,7 +200,13 @@ 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": + 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: 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..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 @@ -1,135 +1,239 @@ # 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. +"""NVFP4 targets share one implementation and leave the original MXFP4 targets independent.""" -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: +import importlib +from types import SimpleNamespace - # >>> 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. -""" +import pytest +import torch -from __future__ import annotations +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 -import difflib -from pathlib import Path -from typing import List, Tuple +pytestmark = pytest.mark.cpu_only -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" +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, ) - 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) + 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}) ) - return drift - + monkeypatch.setenv("TRTLLM_USE_PY_MAMBA", "1") + with pytest.raises(ValueError, match="MambaHybridCacheManagerV2"): + get_kv_cache_manager_cls(config, cache) -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) +@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" ) - - -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"]) + 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, + use_cute_dsl_blockscaling_mm=True, + disable_deep_gemm=False, + ) + 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( + "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 + + 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", + [ + (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" + ) + 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, + ) + ) + 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 ae2731dbe63e..437363f44dc6 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,14 +464,237 @@ 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")) + + +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"]) +@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))) == ( + 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"]) +@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" + 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(): + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kimi_k3_nvfp4.implementation import ( + 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) + 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) + + +@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.kimi_k3_nvfp4.implementation 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 + 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 import ( + implementation 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 import ( + implementation 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 import ( + implementation 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 import ( + implementation 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..85d205860ad1 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,91 @@ 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 import modeling as target + from tensorrt_llm.models.modeling_utils import QuantConfig + + 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, + 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. 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),