diff --git a/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp b/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp index 081f0eed9945..9a2b4271ab00 100644 --- a/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp +++ b/cpp/tensorrt_llm/thop/kdaDecodeOp.cpp @@ -205,14 +205,21 @@ at::Tensor kda_decode_fusion_forward(at::Tensor x_q, at::Tensor x_k, at::Tensor at::Tensor cs_v, at::Tensor a_log, at::Tensor g, at::Tensor dt_bias, at::Tensor beta, at::Tensor onorm_g, at::Tensor onorm_weight, std::optional ssm_state_indices, at::Tensor cu_seqlens, at::Tensor state, bool apply_onorm, bool update_conv_cache, bool use_lower_bound, bool apply_beta_sigmoid, double lower_bound, - double scale, double onorm_eps) + double scale, double onorm_eps, std::optional output) { validate_kda_decode_fusion_inputs(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, apply_onorm, update_conv_cache); int const B = static_cast(x_q.size(1)); int const HV = static_cast(x_v.size(2)); - auto out = at::empty({B, 1, HV, kDimV}, x_q.options()); + auto out = output.has_value() ? *output : at::empty({B, 1, HV, kDimV}, x_q.options()); + if (output.has_value()) + { + TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::kBFloat16, "out must be a CUDA bfloat16 tensor"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + TORCH_CHECK(out.dim() == 4 && out.size(0) == B && out.size(1) == 1 && out.size(2) == HV && out.size(3) == kDimV, + "out must have shape [B, 1, HV, 128]"); + } launch_selected_kernel(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, cs_q, cs_k, cs_v, a_log, g, dt_bias, beta, onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, out, apply_onorm, update_conv_cache, use_lower_bound, apply_beta_sigmoid, lower_bound, scale, onorm_eps); @@ -236,7 +243,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor? ssm_state_indices, Tensor cu_seqlens, Tensor(d!) state, " "bool apply_onorm, bool update_conv_cache, bool use_lower_bound, " "bool apply_beta_sigmoid, float lower_bound, float scale, " - "float onorm_eps) -> Tensor"); + "float onorm_eps, Tensor(e!)? output=None) -> Tensor(e!)"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index c5893c21bfae..4152f14ccb07 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -23,6 +23,7 @@ Gemma4UnifiedTextConfig, Gemma4UnifiedVisionConfig, ) +from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig from tensorrt_llm._torch.configs.laguna import LagunaConfig from tensorrt_llm._torch.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig @@ -54,6 +55,12 @@ def _register_custom_configs_with_transformers() -> None: "kimi_k2": DeepseekV3Config, "deepseek_v4": DeepseekV4Config, "gemma4_assistant": Gemma4AssistantConfig, + # Kimi K3 text config ("kimi_linear"). The composite "kimi_k3" + # model_type is flattened to the text config by + # pyexecutor.config_utils.load_pretrained_config; registering the + # text config here lets AutoConfig / AutoTokenizer resolve + # "kimi_linear" without trust_remote_code. + "kimi_linear": KimiLinearConfig, "laguna": LagunaConfig, # minicpmv4_6 is only registered in transformers>=5.7.0; register our # own composite config so AutoTokenizer.from_pretrained works on older @@ -86,6 +93,7 @@ def _register_custom_configs_with_transformers() -> None: "Gemma4UnifiedConfig", "Gemma4UnifiedTextConfig", "Gemma4UnifiedVisionConfig", + "KimiLinearConfig", "LagunaConfig", "MiniCPMV4_6Config", "MiniCPMV4_6VisionConfig", diff --git a/tensorrt_llm/_torch/configs/kimi_linear.py b/tensorrt_llm/_torch/configs/kimi_linear.py new file mode 100644 index 000000000000..b00fa6627cd0 --- /dev/null +++ b/tensorrt_llm/_torch/configs/kimi_linear.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-tree config for Kimi K3 ("kimi_linear") text checkpoints. + +Mirrors the checkpoint-shipped ``configuration_kimi_k3.KimiLinearConfig`` so +TRT-LLM can parse Kimi K3 checkpoints without ``trust_remote_code`` for the +config. The top-level Kimi K3 checkpoints use a composite VLM config +(``model_type: kimi_k3``) whose ``text_config`` is this class; +``load_pretrained_config`` flattens the composite config to this text config +(TRT-LLM runs the text model only). +""" + +from typing import Optional + +from transformers.configuration_utils import PretrainedConfig + + +class KimiLinearConfig(PretrainedConfig): + model_type = "kimi_linear" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=163840, + hidden_size=4096, + head_dim=None, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=False, + moe_intermediate_size: Optional[int] = None, + moe_renormalize: bool = True, + moe_router_activation_func: str = "sigmoid", + num_experts: Optional[int] = None, + num_experts_per_token: Optional[int] = None, + num_shared_experts: int = 0, + routed_scaling_factor: float = 1.0, + first_k_dense_replace: int = 0, + moe_layer_freq: int = 1, + use_grouped_topk: bool = True, + num_expert_group: int = 1, + topk_group: int = 1, + q_lora_rank: Optional[int] = None, + kv_lora_rank: Optional[int] = None, + qk_nope_head_dim: Optional[int] = None, + qk_rope_head_dim: Optional[int] = None, + v_head_dim: Optional[int] = None, + mla_use_nope: Optional[bool] = False, + mla_use_output_gate: Optional[bool] = False, + num_nextn_predict_layers: int = 0, + linear_attn_config: Optional[dict] = None, + attn_res_block_size: Optional[int] = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: Optional[float] = None, + activation_situ_linear_beta: Optional[float] = None, + max_position_embeddings: int = 4096, + routed_expert_hidden_size: Optional[int] = None, + topk_method: str = "noaux_tc", + **kwargs, + ): + # NOTE: unlike the checkpoint-shipped config class, do not accept a + # ``model_type`` kwarg that shadows the class attribute; transformers + # keys registry lookups off the class attribute. + kwargs.pop("model_type", None) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate + # moe config + self.num_experts = num_experts + self.num_experts_per_token = num_experts_per_token + self.moe_renormalize = moe_renormalize + self.num_shared_experts = num_shared_experts + self.routed_scaling_factor = routed_scaling_factor + self.moe_router_activation_func = moe_router_activation_func + assert self.moe_router_activation_func in ("softmax", "sigmoid") + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + self.use_grouped_topk = use_grouped_topk + self.num_expert_group = num_expert_group + self.topk_group = topk_group + self.num_nextn_predict_layers = num_nextn_predict_layers + + self.attn_res_block_size = attn_res_block_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.max_position_embeddings = max_position_embeddings + self.routed_expert_hidden_size = routed_expert_hidden_size + self.topk_method = topk_method + + if linear_attn_config is not None: + assert linear_attn_config["kda_layers"] is not None + assert linear_attn_config["full_attn_layers"] is not None + self.linear_attn_config = linear_attn_config + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def is_mla(self): + return ( + self.q_lora_rank is not None + or self.kv_lora_rank is not None + or self.qk_nope_head_dim is not None + or self.qk_rope_head_dim is not None + or self.v_head_dim is not None + or self.mla_use_nope is True + ) + + @property + def is_moe(self): + return self.num_experts is not None + + @property + def is_linear_attn(self) -> bool: + return not ( + self.linear_attn_config is None + or ( + isinstance(self.linear_attn_config, dict) + and self.linear_attn_config["kda_layers"] is not None + and len(self.linear_attn_config["kda_layers"]) == 0 + ) + ) + + def is_kda_layer(self, layer_idx: int) -> bool: + """0-indexed layer check; ``kda_layers`` in the config is 1-indexed.""" + return ( + self.linear_attn_config is not None + and (layer_idx + 1) in self.linear_attn_config["kda_layers"] + ) + + def is_full_attn_layer(self, layer_idx: int) -> bool: + """0-indexed layer check; ``full_attn_layers`` is 1-indexed.""" + return ( + self.linear_attn_config is not None + and (layer_idx + 1) in self.linear_attn_config["full_attn_layers"] + ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 7a5408b77b79..9ebbd8df355e 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -294,11 +294,39 @@ def _(logits, pass @torch.library.register_fake("trtllm::kda_decode") - def _(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v, - conv_state_q, conv_state_k, conv_state_v, a_log, g, dt_bias, beta, - onorm_g, onorm_weight, ssm_state_indices, cu_seqlens, state, - apply_onorm, update_conv_cache, use_lower_bound, apply_beta_sigmoid, - lower_bound, scale, onorm_eps): + def _(x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + bias_q: torch.Tensor, + bias_k: torch.Tensor, + bias_v: torch.Tensor, + conv_state_q: torch.Tensor, + conv_state_k: torch.Tensor, + conv_state_v: torch.Tensor, + a_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + onorm_g: torch.Tensor, + onorm_weight: torch.Tensor, + ssm_state_indices: Optional[torch.Tensor], + cu_seqlens: torch.Tensor, + state: torch.Tensor, + apply_onorm: bool, + update_conv_cache: bool, + use_lower_bound: bool, + apply_beta_sigmoid: bool, + lower_bound: float, + scale: float, + onorm_eps: float, + output: Optional[torch.Tensor] = None) -> torch.Tensor: + # Mirror the CUDA impl: write into the caller-provided output when + # given (schema returns Tensor(e!)), else allocate. + if output is not None: + return output # x_q is [1, tokens, H, 128]; the kernel emits one row per token. return x_q.new_empty((x_q.size(1), 1, x_v.size(2), x_v.size(3))) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 08f6f325329c..fe37d7a3484e 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -29,8 +29,8 @@ from transformers.utils import HF_MODULES_CACHE from tensorrt_llm._torch.pyexecutor.config_utils import ( - get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid, - load_pretrained_config) + get_kimi_linear_num_attention_layers, get_qwen3_hybrid_num_attention_layers, + is_kimi_linear, is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config) from tensorrt_llm._utils import (get_sm_version, is_sm_100f, torch_dtype_to_binding) from tensorrt_llm.bindings import LayerType as LayerTypeCpp @@ -72,11 +72,31 @@ def _is_lock_infra_error(exc: BaseException) -> bool: if isinstance(exc, PermissionError): return True if isinstance(exc, OSError): + # EEXIST: filelock's ensure_directory_exists() can lose the + # mkdir(exist_ok=True) race on NFS when many ranks start at once + # (the post-EEXIST is_dir() recheck sees a stale attribute cache). + # An un-creatable lock dir is broken infra, not contention. return exc.errno in (errno.EACCES, errno.EPERM, errno.ENOLCK, - errno.ESTALE) + errno.ESTALE, errno.EEXIST) return False +def _release_lock_ignoring_infra_errors(lock: "filelock.BaseFileLock") -> None: + """Release ``lock``, downgrading broken-lock-infra errors to a warning. + + NFS can return ENOLCK/ESTALE from the unlock ``flock`` call itself (e.g. + lock-daemon exhaustion when many ranks start simultaneously). The config + load the lock protected has already completed at release time, so + crashing the process here would fail an otherwise healthy executor. + """ + try: + lock.release() + except (PermissionError, OSError) as e: + if not _is_lock_infra_error(e): + raise + logger.warning(f"config lock release failed ({e}), continuing") + + @contextlib.contextmanager def config_file_lock(timeout: int = 10): """ @@ -120,12 +140,12 @@ def config_file_lock(timeout: int = 10): try: yield finally: - tmp_lock.release() + _release_lock_ignoring_infra_errors(tmp_lock) else: try: yield finally: - lock.release() + _release_lock_ignoring_infra_errors(lock) @dataclass(kw_only=True) @@ -1391,6 +1411,8 @@ def get_num_attention_layers(self) -> int: return cfg.hybrid_override_pattern.count("*") if is_qwen3_hybrid(cfg): return get_qwen3_hybrid_num_attention_layers(cfg) + if is_kimi_linear(cfg): + return get_kimi_linear_num_attention_layers(cfg) return cfg.num_hidden_layers def get_num_mamba_layers(self) -> int: @@ -1401,6 +1423,9 @@ def get_num_mamba_layers(self) -> int: if is_qwen3_hybrid(cfg): return cfg.num_hidden_layers - get_qwen3_hybrid_num_attention_layers( cfg) + if is_kimi_linear(cfg): + return cfg.num_hidden_layers - get_kimi_linear_num_attention_layers( + cfg) return 0 diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 4f4905d83086..d48fc280c157 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -30,6 +30,7 @@ from .modeling_hunyuan_moe import HunYuanMoEV1ForCausalLM from .modeling_hyperclovax import HCXVisionForCausalLM from .modeling_kimi_k25 import KimiK25ForConditionalGeneration +from .modeling_kimi_linear import KimiLinearForCausalLM from .modeling_laguna import LagunaForCausalLM from .modeling_llama import LlamaForCausalLM from .modeling_llava_next import LlavaNextModel @@ -89,6 +90,7 @@ "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", "KimiK25ForConditionalGeneration", + "KimiLinearForCausalLM", "LlamaForCausalLM", "LlavaNextModel", "MiniCPMV4_6Model", diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index a4e8f898548b..411c47b830a9 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import glob +import json import multiprocessing import os import threading @@ -221,11 +222,71 @@ def _with_weight_cache(self, weight_files: List[str], self._cache_loaded_weights(cache_key, weights) return weights + def cleanup(self) -> None: + # Drop lazy safetensors handles (if any) so the mmaps are released. + self._lazy_handles = [] + super().cleanup() + + @staticmethod + def _is_kimi_k3_checkpoint(checkpoint_dir: str) -> bool: + """Kimi K3 checkpoints (~1.5 TB) must not be materialized in host RAM.""" + config_path = os.path.join(checkpoint_dir, "config.json") + if not os.path.isfile(config_path): + return False + # Do not swallow read/parse failures: every rank must take the same + # branch here (the non-Kimi path enqueues collectives), so a + # rank-local transient error routing one rank differently would + # deadlock the job. Propagating fails fast on all ranks instead. + with open(config_path) as f: + model_type = json.load(f).get("model_type") + return model_type in ("kimi_k3", "kimi_linear") + + def _load_lazy_safetensors( + self, + checkpoint_dir: str, + use_consolidated: bool = False) -> dict[str, Any]: + """Return a dict of name -> lazy safetensors slices. + + Values are ``safetensors`` PySafeSlice objects: ``v[:]`` (or any + indexing) materializes only the requested bytes from the mmapped + file. This lets a model's ``load_weights`` stream a huge checkpoint + and read only its rank-local shard (e.g. Kimi K3 expert-parallel + expert slices) without ever holding the full checkpoint in RAM. + """ + weight_files = sorted(glob.glob(f"{checkpoint_dir}/*.safetensors")) + if not weight_files: + raise RuntimeError(f"No safetensors files in {checkpoint_dir}.") + # Same sharded-vs-consolidated selection as the eager path below: + # when both flavors are present, keep only the requested one. + filtered_weight_files = [ + x for x in weight_files + if ("consolidated" in os.path.split(x)[1]) == use_consolidated + ] + if len(filtered_weight_files) > 0: + weight_files = filtered_weight_files + weights: dict[str, Any] = {} + handles = [] + for file_name in weight_files: + handle = safetensors.safe_open(file_name, + framework="pt", + device="cpu") + handles.append(handle) + for name in handle.keys(): + weights[name] = handle.get_slice(name) + # Keep the file handles alive for as long as the loader lives; the + # slices reference them. Released in cleanup(). + self._lazy_handles = handles + logger.info(f"Lazily opened {len(weight_files)} safetensors files " + f"({len(weights)} tensors) from {checkpoint_dir}") + return ConsumableWeightsDict(weights) + def load_weights(self, checkpoint_dir: str, mapping: Mapping, use_consolidated: bool = False, **kwargs) -> dict[str, Any]: + if self._is_kimi_k3_checkpoint(checkpoint_dir): + return self._load_lazy_safetensors(checkpoint_dir, use_consolidated) weight_files = glob.glob(f"{checkpoint_dir}/*.safetensors") # Some model checkpoint directories contain not only the sharded safetensors, but one # consolidated tensor. In the presence of both, we favor the former unless specified explicitly, as there really is no need diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py new file mode 100644 index 000000000000..60b2b7c157a5 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -0,0 +1,2640 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KimiLinearForCausalLM — Kimi K3 text model, PyTorch backend. + +Runtime integration of the Kimi K3 hybrid architecture for the standard +TRT-LLM PyTorch-backend flow (``LLM(model=) -> generate``): + +* 93 decoder layers: 69 KDA (Kimi Delta Attention, linear attention) layers + and 24 MLA (absorbed-MQA, NoPE) layers, per the 1-indexed + ``linear_attn_config.kda_layers`` / ``full_attn_layers`` schedule. +* Layer 0 uses a dense SiTU MLP; layers 1..92 use the 896-expert latent MoE + (top-16, sigmoid + e_score_correction_bias routing, MXFP4 routed experts, + 2 shared experts, latent down/norm/up projections). +* The attention-residual ("attn_res") scheme from the HF reference + ``modeling_kimi.py`` is applied per token: snapshot mixing before + ``input_layernorm`` / ``post_attention_layernorm`` and at the model output, + with a new snapshot appended whenever ``layer_idx % attn_res_block_size == 0``. + +Caching +------- +KDA states live on the mamba side of a ``MixedMambaHybridCacheManager`` +(wired in ``pyexecutor/_util.py``): per layer, a short-conv slot of +``[3 * num_heads * head_dim, W]`` bf16 (the full FLA ``ShortConvolution`` +cache window, sections ``[q | k | v]``) and a delta-rule recurrent slot of +``[num_heads, head_dim, head_dim]`` fp32 (``[H, V, K]``, the +``state_v_first`` FLA layout). MLA layers use the paged-KV side with +``num_kv_heads=1`` and ``head_dim = kv_lora_rank + qk_rope_head_dim`` (576), +SELFKONLY, exactly like DeepSeek MLA. + +MLA prefill routing +------------------- +The in-tree ``KimiK3MLAAttention`` routes prefill through the normal +unabsorbed MLA context FMHA. It consumes the executor's original mixed-batch +metadata, letting the shared MLA implementation dispatch context and cached +generation work in one forward call. + +Parallelism +----------- +The routed-expert bank supports a configurable MoE TP x EP split +(``moe_tp_size * moe_ep_size == mapping.tp_size``); the default is EP-only +(``moe_ep_size == mapping.tp_size``), the historical K3 layout. Under EP each +MoE layer holds a contiguous ``num_experts / moe_ep_size`` slice of the MXFP4 +expert bank (whole experts); under MoE TP each rank holds ALL experts, with +w1/w3 column-sharded and w2 row-sharded along the intermediate dim +(``intermediate / moe_tp_size`` per rank; group-32 MXFP4 packed bytes and +scales sliced consistently by the stock TRTLLM-Gen quant-method loaders). +The split is EP-only unless the user sets ``moe_tensor_parallel_size`` / +``moe_expert_parallel_size`` explicitly (or the ``TLLM_K3_MOE_TP_SIZE`` / +``TLLM_K3_MOE_EP_SIZE`` env overrides). Routing is computed replicated; the +routed partial sums — EP partials of whole experts, or TP partials over the +intermediate shards — are all-reduced in the latent space (before +``routed_expert_norm`` / ``routed_expert_up_proj``, which are +nonlinear/linear layers applied to the full sum). ``lm_head`` uses the stock +``LMHead`` (vocab-sharded + gather), so logits are identical on all ranks. + +Speculative decoding: SA (suffix automaton, one-engine, draft-weight-free); +the KDA/MLA runtimes implement multi-token verification with deferred +state promotion. + +Chunked prefill is supported: continuation chunks feed the previous KDA +conv/recurrent state back into the FLA kernels (``use_initial_states``) +and the MLA prefill path natively attends over the cached latent prefix +(``kv_len = cached + q_len``). KV-cache block reuse is supported as an +opt-in via ``kv_cache_config.enable_block_reuse=true``, which routes to +the unified-pool ``CppMambaHybridCacheManager`` (per-block KDA state +snapshots every ``mamba_state_cache_interval`` tokens, FORCE_CHUNK +context chunking). + +Not supported: pipeline parallelism, draft-head spec-dec modes +(MTP/Eagle — no draft-head checkpoint exists). SA speculative decoding +is validated only without block reuse (Mixed cache manager). +""" + +from __future__ import annotations + +import copy +import os +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import torch +from torch import nn + +from ..._utils import is_sm_100f +from ...logger import logger +from ...mapping import Mapping +from ...models.modeling_utils import QuantAlgo, QuantConfig +from ..attention_backend import AttentionMetadata +from ..distributed import AllReduce, AllReduceStrategy +from ..model_config import ModelConfig +from ..modules.fused_moe import ConfigurableMoE, create_moe +from ..modules.kimi_k3_moe._mlp import KimiK3MLP, KimiK3RMSNorm +from ..modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate +from ..modules.linear import Linear as TrtllmLinear +from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.rms_norm import RMSNorm +from ..utils import ActType_TrtllmGen +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, register_auto_model + +# 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" +) + +_KDA_INDEXED_STATE_POOL_ENABLED = os.environ.get("TLLM_KDA_ENABLE_INDEXED_STATE_POOL", "1") == "1" + +# Routed-expert MoE TP/EP split overrides (read per model init, not import). +# Highest precedence; either one may be set alone, the other is derived from +# tp_size. Without them, an explicit moe_tensor_parallel_size / +# moe_expert_parallel_size pair from the user config is honored, and the +# default stays EP-only (moe_ep == tp_size). +_K3_MOE_TP_ENV = "TLLM_K3_MOE_TP_SIZE" +_K3_MOE_EP_ENV = "TLLM_K3_MOE_EP_SIZE" + +if TYPE_CHECKING: + from transformers import PretrainedConfig + +# 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" +_KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES = ( + ".self_attn.mixer.k_b_proj_trans", + ".self_attn.mixer.v_b_proj", +) + +# Serve the replicated MoE-layer MLP projections (shared-expert gate/up/down +# and the latent up/down projection) from an FP8 copy of their weights instead +# of BF16. Under attention data-parallelism every rank re-reads these dense +# weights in full on every decode step, so decode is bound by that HBM read; +# an FP8 (e4m3) weight with 128x128 block scales roughly halves those bytes. +# The MLA projections and the routed MXFP4 experts are left untouched (the KDA +# q/k/v/g/o projections have their own switch below). The FP8 weight read is +# lossy relative to BF16; set this to "0" to keep BF16. +_KIMI_K3_FP8_WEIGHT_READ_ENV = "KIMI_K3_FP8_WEIGHT_READ" + +# Also read the KDA linear-attention q/k/v/g/o projections at FP8 block-scale. +# These are the largest single replicated weight read (~61 GB/rank of the +# ~109 GB BF16 read per decode step). They use the same FP8 path as the MLP +# projections above but are gated separately: the recurrent linear-attention +# core is more accuracy-sensitive than the feed-forward MLPs, so set this to +# "0" to keep the KDA projections in BF16 while still reading the MLPs at FP8. +# The master KIMI_K3_FP8_WEIGHT_READ switch and the SM100 gate still apply. +_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV = "KIMI_K3_FP8_WEIGHT_READ_KDA" + +# Also read the MLA (full-attention) q_a/q_b/o and output-gate projections at +# FP8 block-scale. These are the replicated attention weights the MLP pass and +# the KDA pass above leave in BF16, and they are re-read in full by every rank +# each decode step under attention data-parallelism. Two MLA projections are +# deliberately kept in BF16: kv_a_proj_with_mqa outputs kv_lora_rank + +# qk_rope_head_dim (576, not a multiple of 128, so no exact 128x128 block +# scale), and kv_b_proj's weight is consumed directly (not through its forward) +# by the absorbed-decode _kv_b_absorb_split to build the k/v absorb matrices, +# which has no FP8 dequant path. The master KIMI_K3_FP8_WEIGHT_READ switch and +# the SM100 gate still apply. +_KIMI_K3_FP8_WEIGHT_READ_MLA_ENV = "KIMI_K3_FP8_WEIGHT_READ_MLA" + +# Opt-in (prototype): keep the KimiKDARuntime decode fast path — fused +# in-projection + persistent conv staging + precomputed kernel-layout +# constants (``_forward_decode``) — when the KDA projections are read at FP8 +# block-scale. By default the FP8 KDA read routes decode through the +# reference path, which re-does ~70 us/layer of glue per decode step around +# the 5 us kernel (see ``_forward_decode``'s docstring). With this set to +# "1", the fast path issues the loader's fused FP8 ``qkvg_proj`` GEMM for +# q/k/v/g plus one small BF16 GEMV for [f_a | b] +# (``finalize_decode_weights_fp8``), so FP8 weight storage and the decode +# glue savings coexist. Requires the FP8 KDA read to be active; no effect +# otherwise. Default on ("0" disables). +_KIMI_K3_KDA_GLUE_FP8_ENV = "KIMI_K3_KDA_GLUE_FP8" + +# FP8 read for the fused shared-expert gate_up_proj. Default follows the +# parallel layout (on under attention DP, off under TP — see the conversion +# helper's comment); set 0/1 to force either. +_KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV = "KIMI_K3_FP8_WEIGHT_READ_GATE_UP" + + +# --------------------------------------------------------------------------- +# Config helpers. +# --------------------------------------------------------------------------- + + +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"] + + +# --------------------------------------------------------------------------- +# attn_res: per-token snapshot mixing (HF `_apply_attn_res`). +# --------------------------------------------------------------------------- + + +KIMI_K3_FUSED_ATTN_RES_ENV = "KIMI_K3_FUSED_ATTN_RES" +"""Set to ``0`` to disable the in-tree fused Torch op +``trtllm::attn_res_fwd`` (Blackwell only). Default: fused with fallback.""" + +_FUSED_ATTN_RES_ENABLED = os.environ.get(KIMI_K3_FUSED_ATTN_RES_ENV, "1") == "1" + + +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: + 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 _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) + + +# --------------------------------------------------------------------------- +# Dense / shared-expert MLP: fused [gate | up] layout (``KimiK3MLP``). +# +# The HF checkpoint stores separate ``gate_proj`` / ``up_proj`` tensors; +# ``load_weights`` row-concatenates them into ``gate_up_proj`` (see +# ``_gate_up_ckpt_keys``), replacing two GEMMs + torch.cat with one GEMM. +# --------------------------------------------------------------------------- + + +_GATE_UP_FUSED_SUFFIX = ".gate_up_proj.weight" + + +def _gate_up_ckpt_keys(fused_key: str) -> Tuple[str, str]: + """Checkpoint ``(gate_proj, up_proj)`` keys whose row-concat loads the + fused ``gate_up_proj`` parameter named by ``fused_key``.""" + return ( + fused_key.replace(_GATE_UP_FUSED_SUFFIX, ".gate_proj.weight"), + fused_key.replace(_GATE_UP_FUSED_SUFFIX, ".up_proj.weight"), + ) + + +# --------------------------------------------------------------------------- +# FP8 block-scale weight read for the replicated MoE-layer MLP projections. +# --------------------------------------------------------------------------- + + +class _Fp8BlockScaleWeightReadLinear(nn.Module): + """Bias-free ``nn.Linear`` replacement that reads its weight at FP8. + + The BF16 weight ``[out, in]`` is quantized once (at load) to + ``float8_e4m3fn`` with 128x128 block scales, then served through the + DeepGEMM ``fp8_swap_ab_gemm`` kernel — the same FP8 block-scale GEMM the + quantized DeepSeek block-scale path uses. The activation stays BF16 and is + quantized inside the kernel, so only the weight's storage/read precision + changes. Halving the weight bytes cuts the dominant HBM read that bounds + K3's memory-bound decode step. Both ``out`` and ``in`` are multiples of + 128 for every projection this is applied to, so the block scales cover the + weight exactly. + """ + + def __init__( + self, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, out_features: int + ) -> None: + super().__init__() + self.in_features = weight_fp8.shape[1] + self.out_features = out_features + # Buffers (not parameters): these are the module's weights post-load; + # there is nothing further to load into them and they must not be + # touched by any later autocast/dtype move. + self.register_buffer("weight", weight_fp8, persistent=False) + self.register_buffer("weight_scale", weight_scale, persistent=False) + + @staticmethod + def quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """BF16 ``[out, in]`` weight -> (FP8 weight, deep_gemm-ready scale). + + Both dims must be multiples of 128. Because the 128x128 block scale is + computed per block, concatenating several such weights along ``out`` + and quantizing the result is per-block identical to quantizing each + separately (no block crosses a 128-aligned boundary), so a fused + weight's row slices equal the individually quantized weights. + """ + # Lazy imports: only pulled in on the FP8 path. + from ...deep_gemm.utils.math import per_block_cast_to_fp8 + from ...quantization.utils.fp8_utils import ( + resmooth_to_fp8_e8m0, + transform_sf_into_required_layout, + ) + + # 128x128 block-scale FP8 weight, then the exact SM100 deep_gemm scale + # preparation the shipping FP8-block-scale Linear uses: resmooth to + # UE8M0 and pack the scale into deep_gemm's TMA-aligned MN-major + # layout. fp8_swap_ab_gemm runs with disable_ue8m0_cast=True, so it + # consumes this pre-formatted scale directly (a plain FP32 block scale + # would be misread and produce garbage). + weight_fp8, weight_scale = per_block_cast_to_fp8(weight, use_ue8m0=False) + weight_fp8, weight_scale = resmooth_to_fp8_e8m0( + weight_fp8.contiguous(), weight_scale.contiguous().float() + ) + weight_scale = transform_sf_into_required_layout( + weight_scale, + mn=weight_fp8.shape[0], + k=weight_fp8.shape[1], + recipe=(1, 128, 128), + is_sfa=False, + ) + return weight_fp8, weight_scale + + @classmethod + def from_linear(cls, linear: nn.Linear) -> "_Fp8BlockScaleWeightReadLinear": + assert linear.bias is None, "FP8 weight read expects a bias-free Linear" + weight_fp8, weight_scale = cls.quantize_weight(linear.weight.data) + return cls(weight_fp8, weight_scale, linear.out_features) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out_shape = x.shape[:-1] + (self.out_features,) + out = torch.ops.trtllm.fp8_swap_ab_gemm( + x.reshape(-1, x.shape[-1]), + self.weight, + self.weight_scale, + output_dtype=x.dtype, + disable_ue8m0_cast=True, + ) + return out.reshape(out_shape) + + +def _convert_moe_mlps_to_fp8_weight_read( + model: nn.Module, include_fused_gate_up: bool = True +) -> int: + """Swap the replicated MoE-layer MLP projections to an FP8 weight read. + + Targets the shared-expert MLP (gate/up/down) and the latent up/down + projection on every MoE layer — the bias-free BF16 projections that + attention data-parallelism re-reads in full each decode step. Attention + (MLA/KDA), the routed MXFP4 experts and the dense layer-0 MLP are left in + BF16. Returns the number of projections converted. + """ + import gc + + count = 0 + + def _swap(parent: nn.Module, attr: str) -> None: + nonlocal count + child = getattr(parent, attr, None) + if isinstance(child, nn.Linear): + setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child)) + # Release the original BF16 weight storage now. The loader holds a + # transient name->Parameter map that keeps it alive until load + # returns, so without this the FP8 copy is purely additive and + # fragments the pool the FP8 GEMM autotuner and KV-cache init need. + child.weight.data = child.weight.data.new_empty(0) + count += 1 + + for layer in model.layers: + moe = getattr(layer, "block_sparse_moe", None) + if moe is None: + continue + shared = getattr(moe, "shared_experts", None) + if shared is not None: + # KimiK3MLP fuses gate and up into gate_up_proj; keep the split + # names too so either MLP layout converts. The fused gate_up read + # only pays off when attention DP re-reads it per rank per step; + # under TP the bf16 GEMM overlaps on the aux stream and the FP8 + # quantize+GEMM would serialize onto the critical path. + shared_attrs = ( + ("gate_proj", "up_proj", "gate_up_proj", "down_proj") + if include_fused_gate_up + else ("gate_proj", "up_proj", "down_proj") + ) + for attr in shared_attrs: + _swap(shared, attr) + for attr in ("routed_expert_down_proj", "routed_expert_up_proj"): + _swap(moe, attr) + + # Return the freed BF16 blocks to the driver so the raw (non-caching- + # allocator) allocations made during executor creation succeed on the + # tight DEP16 memory headroom. + if count: + gc.collect() + torch.cuda.empty_cache() + return count + + +def _convert_kda_projections_to_fp8_weight_read(model: nn.Module) -> int: + """Swap the KDA linear-attention q/k/v/g/o projections to an FP8 weight read. + + Targets the large bias-free BF16 projections of every KDA linear-attention + layer (``q_proj``/``k_proj``/``v_proj``/``g_proj``/``o_proj``, each + ``[out, in]`` with both dims a multiple of 128) — the single largest + replicated weight read, re-read in full by every rank each decode step + under attention data-parallelism. The smaller state-path projections are + left in BF16 on purpose: ``b_proj`` outputs ``num_heads`` (not a multiple + of 128, so no exact 128x128 block scale), and the forget gate + ``f_a``/``f_b``, the low-rank ``g_a``/``g_b`` gate, the short convolutions + and ``dt`` are small and feed the accuracy-sensitive recurrent decay. + + ``q_proj``/``k_proj``/``v_proj`` and the full-rank ``g_proj`` all read the + same normed hidden, so their weights are additionally concatenated into one + fused ``qkvg_proj`` FP8 GEMM used by the decode path + (``KimiKDALinearAttention._decode_via_optimized``): the decode step is + launch-bound at the small generation batch, and one GEMM (with one shared + activation quant) replaces four. The fused weight is the only storage — the + individual ``q_proj``/``k_proj``/``v_proj``/``g_proj`` modules are rebuilt to + read a **view** of their slice of it (with their own block scale), so the + prefill/verify paths keep calling them per projection with no extra memory. + ``o_proj`` reads the decode-kernel output (not the shared hidden) and is + converted on its own. Returns the number of projections converted. + """ + import gc + + count = 0 + + for layer in model.layers: + if not getattr(layer, "is_kda", False): + continue + mixer = getattr(getattr(layer, "self_attn", None), "mixer", None) + if mixer is None: + continue + + # Projections that read the same normed hidden (g_proj only in the + # full-rank-gate config; the low-rank g_a/g_b gate stays BF16). + group = [(a, getattr(mixer, a, None)) for a in ("q_proj", "k_proj", "v_proj", "g_proj")] + group = [(a, c) for a, c in group if isinstance(c, nn.Linear)] + + if group: + # One fused FP8 weight [sum(out), in]; row slices equal the + # individually quantized weights (see quantize_weight). + fused_bf16 = torch.cat([c.weight.data for _, c in group], dim=0) + fused_fp8, fused_scale = _Fp8BlockScaleWeightReadLinear.quantize_weight(fused_bf16) + fused = _Fp8BlockScaleWeightReadLinear(fused_fp8, fused_scale, fused_bf16.shape[0]) + mixer.qkvg_proj = fused + mixer.qkvg_split_sizes = [c.out_features for _, c in group] + del fused_bf16 + + # Rebuild each projection to read a view of its slice of the fused + # weight (own block scale); the fused weight is the sole storage. + offset = 0 + for attr, child in group: + n = child.out_features + _, own_scale = _Fp8BlockScaleWeightReadLinear.quantize_weight(child.weight.data) + setattr( + mixer, + attr, + _Fp8BlockScaleWeightReadLinear(fused.weight[offset : offset + n], own_scale, n), + ) + # Free the original BF16 storage (the loader's transient + # name->Parameter map keeps it alive until load returns, so + # without this the FP8 copy is purely additive on the tight + # DEP16 pool). + child.weight.data = child.weight.data.new_empty(0) + offset += n + count += 1 + + # o_proj reads the decode-kernel output, so it is not part of the fused + # hidden-reading group; convert it on its own. + o_proj = getattr(mixer, "o_proj", None) + if isinstance(o_proj, nn.Linear): + setattr(mixer, "o_proj", _Fp8BlockScaleWeightReadLinear.from_linear(o_proj)) + o_proj.weight.data = o_proj.weight.data.new_empty(0) + count += 1 + + if count: + gc.collect() + torch.cuda.empty_cache() + return count + + +@torch.no_grad() +def _load_kimi_k3_mla_kv_b_proj( + mixer: nn.Module, + source: torch.Tensor, + *, + head_start: int, +) -> None: + """Load checkpoint-interleaved KV-B rows into MLA runtime tensors.""" + weight = mixer.kv_b_proj.weight + h = mixer.num_heads + n = mixer.qk_nope_head_dim + v = mixer.v_head_dim + kv = mixer.kv_lora_rank + head_width = n + v + + if source.ndim != 2 or source.shape[1] != kv: + raise ValueError( + "Kimi K3 MLA kv_b_proj checkpoint shape " + f"{tuple(source.shape)} is not [heads * {head_width}, {kv}]" + ) + if source.shape[0] % head_width != 0: + raise ValueError( + "Kimi K3 MLA kv_b_proj checkpoint rows " + f"{source.shape[0]} are not divisible by {head_width}" + ) + if head_start < 0: + raise ValueError(f"Kimi K3 MLA head_start must be non-negative, got {head_start}") + if weight.device.type == "meta": + raise RuntimeError("Kimi K3 MLA kv_b_proj loading requires materialized weights") + + source_heads = source.shape[0] // head_width + source = source.view(source_heads, head_width, kv) + num_local_source_heads = max(0, min(source_heads - head_start, h)) + if num_local_source_heads == h: + local = source[head_start : head_start + h] + else: + local = source.new_zeros((h, head_width, kv)) + if num_local_source_heads > 0: + local[:num_local_source_heads].copy_( + source[head_start : head_start + num_local_source_heads] + ) + + k_weight, v_weight = local.split([n, v], dim=1) + grouped = torch.cat( + [ + k_weight.reshape(h * n, kv), + v_weight.reshape(h * v, kv), + ], + dim=0, + ) + if grouped.shape != weight.shape: + raise ValueError( + "Kimi K3 MLA grouped kv_b_proj shape " + f"{tuple(grouped.shape)} does not match {tuple(weight.shape)}" + ) + weight.copy_(grouped.to(weight.dtype)) + + loaded_k, loaded_v = weight.split([h * n, h * v], dim=0) + k_b_proj_trans = loaded_k.view(h, n, kv).transpose(1, 2).contiguous() + if k_b_proj_trans.shape != mixer.k_b_proj_trans.shape: + raise ValueError( + "Kimi K3 MLA k_b_proj_trans shape " + f"{tuple(k_b_proj_trans.shape)} does not match " + f"{tuple(mixer.k_b_proj_trans.shape)}" + ) + mixer.k_b_proj_trans.copy_(k_b_proj_trans) + mixer.v_b_proj = nn.Parameter( + loaded_v.view(h, v, kv), + requires_grad=False, + ) + + +def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int: + """Swap the MLA q_a/q_b/o and output-gate projections to an FP8 weight read. + + Targets the large bias-free BF16 projections of every MLA (full-attention) + layer that are read only through their ``forward`` — ``q_a_proj``, + ``q_b_proj``, ``o_proj`` and, when the output gate is enabled, ``g_proj``, + each ``[out, in]`` with both dims a multiple of 128 — replicated attention + weights re-read in full by every rank each decode step under attention + data-parallelism. Two MLA projections are left in BF16 on purpose: + ``kv_a_proj_with_mqa`` outputs ``kv_lora_rank + qk_rope_head_dim`` (576, not + a multiple of 128), and ``kv_b_proj`` supplies ``k_b_proj_trans`` and + ``v_b_proj`` directly (the absorbed generation path never calls its + ``forward``), with no FP8 dequant path. Returns the number of projections + converted. + """ + import gc + + count = 0 + + def _swap(parent: nn.Module, attr: str) -> None: + nonlocal count + child = getattr(parent, attr, None) + if isinstance(child, (nn.Linear, TrtllmLinear)): + setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child)) + # Free the original BF16 storage now (as in the MLP/KDA conversions + # above): the loader's transient name->Parameter map would otherwise + # keep it alive until load returns, making the FP8 copy purely + # additive on the tight DEP16 pool. + child.weight.data = child.weight.data.new_empty(0) + count += 1 + + for layer in model.layers: + # MLA layers are the non-KDA layers (each layer is exactly one of the + # two); their projections live on the KimiK3MLAAttention mixer. + if getattr(layer, "is_kda", False): + continue + mixer = getattr(getattr(layer, "self_attn", None), "mixer", None) + if mixer is None: + continue + # g_proj exists only when the MLA output gate is enabled; a missing + # attr is a safe no-op. + for attr in ("q_a_proj", "q_b_proj", "o_proj", "g_proj"): + _swap(mixer, attr) + + if count: + gc.collect() + torch.cuda.empty_cache() + return count + + +# --------------------------------------------------------------------------- +# Latent MoE block using the unified ConfigurableMoE stack. +# --------------------------------------------------------------------------- + + +class KimiK3MoERuntime(nn.Module): + """Kimi K3 latent MoE block backed by ConfigurableMoE/TRTLLM-Gen.""" + + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ): + 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 + assert self.moe_hidden_size is not None, ( + "Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)" + ) + + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + dtype = torch.bfloat16 + + # Routing scores stay fp32; with attention-DP off 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). Under attention-DP the legacy upcast-to-fp32 + # GEMM is kept: the bf16-input min-latency GEMM's different reduction + # order flips borderline top-16 picks (GSM8K 96.7 -> 96.1/96.4, + # 3-run bisect on 62b20dd868), and the bs1-latency win is irrelevant + # at DEP batch sizes. KIMI_K3_ROUTER_BF16=1/0 forces either path. + _router_bf16_env = os.environ.get("KIMI_K3_ROUTER_BF16") + _router_bf16 = ( + _router_bf16_env == "1" + if _router_bf16_env is not None + else not model_config.mapping.enable_attention_dp + ) + 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 = QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) + self.routed_experts = create_moe( + 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, + reduce_results=True, + model_config=routed_moe_model_config, + override_quant_config=routed_quant_config, + layer_idx=layer_idx, + trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, + # Cubin alpha is the gate-side SiTU beta; cubin beta is the + # linear-side SiTU beta. + trtllm_gen_activation_alpha=float(situ_beta), + trtllm_gen_activation_beta=float( + situ_linear_beta if situ_linear_beta is not None else 1.0 + ), + # Let CommunicationFactory select the best available strategy. + communication_method=None, + ) + 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 experts stay replicated (DeepSeek's attention-DP + # semantics): ConfigurableMoE owns its own reduction, so there is + # no existing collective for column-shard partial sums to ride — + # the direct-path shared-expert TP (partials joining the MoE + # combine RS / routed allreduce) needs a partial-carry hook in the + # wrapper before it can be ported (follow-up). Instead their cost + # is hidden by running them on the aux stream, overlapped with the + # routed dispatch/expert/combine chain (see forward()). + shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts + self.shared_experts = KimiK3MLP( + hidden_size=cfg.hidden_size, + intermediate_size=shared_intermediate, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + use_fused_activation=True, + dtype=dtype, + ) + # Side stream (+ fork/join events) for overlapping the replicated + # shared-expert compute with the routed dispatch/expert/combine + # chain; see forward(). Only engaged when multi-stream is active + # (CUDA graphs on) and aux_stream is set, otherwise both run in + # order on the default stream. + self.aux_stream = aux_stream + 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 + ) + assert getattr(cfg, "latent_moe_use_norm", False), ( + "Kimi K3 runtime expects latent_moe_use_norm=True" + ) + # 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 + ) + + @staticmethod + def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: + """Resolve the routed-expert ``(moe_tp, moe_ep)`` split. + + Precedence: + + 1. ``TLLM_K3_MOE_TP_SIZE`` / ``TLLM_K3_MOE_EP_SIZE`` env overrides + (either alone; the other is derived from ``tp_size``). + 2. Explicit ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` + from the user config. Detected via + ``mapping.moe_tp_ep_user_specified`` so the auto-resolved mapping + default (``moe_tp=tp_size, moe_ep=1``) is NOT mistaken for a TP + request. + 3. Default: EP-only (``moe_tp=1, moe_ep=tp_size``), the historical + K3 layout. + """ + tp_size = mapping.tp_size + env_tp = os.environ.get(_K3_MOE_TP_ENV) + env_ep = os.environ.get(_K3_MOE_EP_ENV) + if env_tp is not None or env_ep is not None: + moe_tp = int(env_tp) if env_tp is not None else 0 + moe_ep = int(env_ep) if env_ep is not None else 0 + if moe_tp <= 0 and moe_ep > 0: + moe_tp = tp_size // moe_ep + elif moe_ep <= 0 and moe_tp > 0: + moe_ep = tp_size // moe_tp + return moe_tp, moe_ep + if getattr(mapping, "moe_tp_ep_user_specified", False): + return mapping.moe_tp_size, mapping.moe_ep_size + return 1, tp_size + + @staticmethod + def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: + """Build a private routed-expert mapping without mutating the shared + config. Default split is EP-only; see ``_select_moe_tp_ep``.""" + 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 + if getattr(mapping, "_dwdp_size", 0) > 1: + raise NotImplementedError("Kimi K3 packed-checkpoint streaming does not support DWDP.") + + moe_tp, moe_ep = KimiK3MoERuntime._select_moe_tp_ep(mapping) + if moe_tp < 1 or moe_ep < 1 or moe_tp * moe_ep != mapping.tp_size: + raise ValueError( + f"Kimi K3 routed MoE split moe_tp={moe_tp} x moe_ep={moe_ep} " + f"must multiply to tp_size={mapping.tp_size}." + ) + if moe_tp > 1 and mapping.enable_attention_dp: + raise NotImplementedError( + "Kimi K3 MoE tensor parallelism requires " + "enable_attention_dp=false (the attention-DP dispatch/combine " + "path is validated for EP-only splits)." + ) + 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 = "TRTLLM" + routed_model_config._frozen = True + return routed_model_config + + def forward(self, hidden_states: torch.Tensor, all_rank_num_tokens=None) -> torch.Tensor: + """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" + identity = hidden_states + router_logits = self.gate.compute_logits(hidden_states) + + 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). + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance( + self.routed_expert_down_proj, nn.Linear + ): + routed_in = self.routed_expert_down_proj(hidden_states) + else: + routed_in = torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, self.routed_expert_down_proj.weight.t(), None, None + ) + y = self.routed_experts( + routed_in, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + # EP partial latent sums are completed by the wrapper's own + # reduction BEFORE the (nonlinear) latent norm. + y = self.routed_expert_norm(y) + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance( + self.routed_expert_up_proj, nn.Linear + ): + return self.routed_expert_up_proj(y) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + y, self.routed_expert_up_proj.weight.t(), None, None + ) + + # Shared experts are replicated (computed once per rank) and depend + # only on the block input, not on the routed dispatch/expert/combine + # chain -- so run them on the aux stream to overlap with the serial + # EP dispatch/combine collectives. Multi-stream engages only under + # CUDA graphs; otherwise both run in order on the default stream + # with an identical result. Added after the routed combine so the + # replicated shared output is not double counted. + routed_out, shared_out = maybe_execute_in_parallel( + _routed_output, + lambda: self.shared_experts(identity), + self.moe_main_event, + self.moe_shared_event, + self.aux_stream, + disable_on_compile=True, + ) + return routed_out + shared_out + + +# --------------------------------------------------------------------------- +# KDA runtime (pool-backed prefill / decode via the FLA kernels). +# --------------------------------------------------------------------------- + + +def _kda_split_conv_sections( + cs: torch.Tensor, d: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split a gathered ``[N, 3D, W]`` conv-cache into contiguous q/k/v.""" + return (cs[:, :d].contiguous(), cs[:, d : 2 * d].contiguous(), cs[:, 2 * d :].contiguous()) + + +class KimiKDARuntime(nn.Module): + """Wraps the parity-tested ``KimiKDALinearAttention`` parameters with a + cache-pool-aware forward for the executor flow. + + Parameter names mirror the HF checkpoint 1:1 (the wrapped mixer is + registered under the layer as ``self_attn``, so e.g. + ``model.layers.N.self_attn.q_proj.weight`` maps identically). + """ + + def __init__( + self, cfg, layer_idx: int, mapping=None, allreduce_strategy=AllReduceStrategy.AUTO + ): + super().__init__() + # Lazy import: pulls in fla/einops. + from ..modules.kimi_kda.kimi_kda_mixer import KimiKDALinearAttention + + lin = cfg.linear_attn_config + self.layer_idx = layer_idx + self._use_indexed_ssm_pool = _KDA_INDEXED_STATE_POOL_ENABLED + # Attention-family TP semantics (Qwen3-Next GatedDeltaNet pattern, + # gdn_mixer.py): replicated under attention-DP — each rank runs + # its own batch with the full head set — and head-sharded across + # mapping.tp_size otherwise: every rank holds the same batch, runs + # its 1/tp head slice, and the row-sharded o_proj partials are + # all-reduced at the end of forward(). + if mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp: + self._kda_tp_size = mapping.tp_size + else: + self._kda_tp_size = 1 + self._kda_tp_rank = mapping.tp_rank if self._kda_tp_size > 1 else 0 + self._o_allreduce = ( + AllReduce(mapping=mapping, strategy=allreduce_strategy, dtype=torch.bfloat16) + if self._kda_tp_size > 1 + else None + ) + num_heads = lin["num_heads"] + assert num_heads % self._kda_tp_size == 0, ( + f"KDA num_heads {num_heads} not divisible by tp_size {self._kda_tp_size}" + ) + self.mixer = KimiKDALinearAttention( + hidden_size=cfg.hidden_size, + num_heads=num_heads // self._kda_tp_size, + head_dim=lin["head_dim"], + conv_kernel_size=lin["short_conv_kernel_size"], + use_full_rank_gate=lin.get("use_full_rank_gate", True), + gate_lower_bound=lin.get("gate_lower_bound", None), + rms_norm_eps=cfg.rms_norm_eps, + dtype=torch.bfloat16, + layer_idx=layer_idx, + # Use TLLM_KDA_ENABLE_OPT_PREFILL=0 to opt out of the optimized + # prefill kernel. + use_optimized_prefill=os.getenv("TLLM_KDA_ENABLE_OPT_PREFILL", "1") == "1", + use_optimized_decode=True, + ) + self.proj_size = (num_heads // self._kda_tp_size) * lin["head_dim"] + # Decode fast-path constants, built once by + # ``finalize_decode_weights()`` after checkpoint load: + # fused in-projection weight + kernel-layout conv weights + fp32 + # copies of the small parameters. ``None`` routes decode through + # the reference (module-level) path. + self._in_proj_weight: Optional[torch.Tensor] = None + # FP8 variant (``finalize_decode_weights_fp8``): q/k/v/g stay in the + # mixer's fused FP8 ``qkvg_proj`` GEMM and only the small [f_a | b] + # GEMV weight is fused here. ``None`` when the FP8 fast path is off. + self._in_proj_small_weight: Optional[torch.Tensor] = None + self._w_q_t = self._w_k_t = self._w_v_t = None + self._A_log_f32 = self._dt_bias_f32 = self._onorm_w_f32 = None + # Persistent batch-row-dense staging for the fused decode kernel's + # per-section conv windows (lazily sized to the pool slot count). + self._cs_dense: Optional[torch.Tensor] = None + + def finalize_decode_weights(self) -> None: + """Build the decode fast-path constants (once, after weight load). + + 1. Fused in-projection ``[q | k | v | g | f_a | b]``: all six + projections consume the same hidden state, so one GEMV replaces + five GEMV+splitK-reduce pairs plus a cublas dot per layer per + decode step. The source parameters are repointed to row views + of the fused buffer (no extra memory; prefill/verify paths keep + using them unchanged). + 2. Kernel-layout constants that ``_decode_via_optimized`` used to + rebuild with ~6 device kernels per layer per decode step: + transposed conv weights (bf16 ``[W, D]``) and fp32 copies of + ``A_log`` / ``dt_bias`` / ``o_norm.weight``. + """ + mixer = self.mixer + if mixer._dispatch.decode_kernel_path != "optimized" or not mixer.use_full_rank_gate: + return + if mixer.q_proj.weight.device.type != "cuda": + return + with torch.no_grad(): + mods = ( + mixer.q_proj, + mixer.k_proj, + mixer.v_proj, + mixer.g_proj, + mixer.f_a_proj, + mixer.b_proj, + ) + fused = torch.cat([m.weight.data for m in mods], dim=0).contiguous() + off = 0 + for m in mods: + n = m.weight.shape[0] + m.weight.data = fused[off : off + n] + off += n + self._build_decode_kernel_constants() + # Publish last: `_in_proj_weight is not None` gates the fast path. + self._in_proj_weight = fused + + def _build_decode_kernel_constants(self) -> None: + """Kernel-layout constants shared by both finalize variants.""" + mixer = self.mixer + self._w_q_t = ( + mixer.q_conv1d.weight.detach() + .squeeze(1) + .transpose(0, 1) + .to(torch.bfloat16) + .contiguous() + ) + self._w_k_t = ( + mixer.k_conv1d.weight.detach() + .squeeze(1) + .transpose(0, 1) + .to(torch.bfloat16) + .contiguous() + ) + self._w_v_t = ( + mixer.v_conv1d.weight.detach() + .squeeze(1) + .transpose(0, 1) + .to(torch.bfloat16) + .contiguous() + ) + self._A_log_f32 = mixer.A_log.detach().float().contiguous() + self._dt_bias_f32 = mixer.dt_bias.detach().float().contiguous() + self._onorm_w_f32 = mixer.o_norm.weight.detach().float().contiguous() + # Build the fused-verify conv constants eagerly too, so the first + # verify call never allocates (a capture-unsafe lazy allocation). + self._get_mtp_conv_weights() + + def finalize_decode_weights_fp8(self) -> None: + """FP8 counterpart of ``finalize_decode_weights()``. + + Runs AFTER ``_convert_kda_projections_to_fp8_weight_read``, so + q/k/v/g already live in the mixer's fused FP8 ``qkvg_proj`` GEMM + and the fast path's big projection is that GEMM as-is. Only the + two small BF16 projections reading the same hidden — ``f_a_proj`` + and ``b_proj`` (kept BF16 by the FP8 conversion: outputs are not + 128-multiples and feed the accuracy-sensitive recurrent decay) — + are fused here into one ``[f_a | b]`` GEMV weight, with the source + parameters repointed to row views. The decode step then issues two + GEMMs (FP8 qkvg + BF16 small) instead of the reference path's four + plus its per-step glue; the kernel-layout constants are shared with + the BF16 finalize. + """ + mixer = self.mixer + if mixer._dispatch.decode_kernel_path != "optimized" or not mixer.use_full_rank_gate: + return + fused_qkvg = getattr(mixer, "qkvg_proj", None) + split_sizes = getattr(mixer, "qkvg_split_sizes", None) + if fused_qkvg is None or split_sizes is None or len(split_sizes) != 4: + return + if mixer.f_a_proj.weight.device.type != "cuda": + return + with torch.no_grad(): + mods = (mixer.f_a_proj, mixer.b_proj) + fused = torch.cat([m.weight.data for m in mods], dim=0).contiguous() + off = 0 + for m in mods: + n = m.weight.shape[0] + m.weight.data = fused[off : off + n] + off += n + self._build_decode_kernel_constants() + # Publish last: gates the FP8 decode fast path. + self._in_proj_small_weight = fused + + def forward( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """``hidden_states``: flattened ``[num_tokens, hidden]`` (ctx tokens + first, then one token per generation request).""" + mamba_metadata = attn_metadata.mamba_metadata + num_prefills = attn_metadata.num_contexts + num_ctx_tokens = attn_metadata.num_ctx_tokens + batch_size = attn_metadata.seq_lens.shape[0] + # index_copy_/index_select need int64 indices; the int64 mirror is + # prepared once per step by Mamba2Metadata.prepare() so KDA layers + # do not each replay an int32->int64 cast inside the decode graph. + state_indices = getattr(mamba_metadata, "state_indices_long", None) + if state_indices is None or state_indices.shape[0] != batch_size: + state_indices = mamba_metadata.state_indices[:batch_size].long() + cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] + num_decodes = batch_size - num_prefills + + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + conv_pool = layer_cache.conv # [slots, 3D, W] bf16 + ssm_pool = layer_cache.temporal # [slots, H, V, K] fp32 + + outputs: List[torch.Tensor] = [] + if num_prefills > 0: + outputs.append( + self._forward_prefill( + hidden_states[:num_ctx_tokens], + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + state_indices[:num_prefills], + layer_cache, + ) + ) + if num_decodes > 0: + decode_rows = hidden_states.shape[0] - num_ctx_tokens + if decode_rows == num_decodes: + outputs.append( + self._forward_decode( + hidden_states[num_ctx_tokens:], + conv_pool, + ssm_pool, + state_indices[num_prefills:], + mamba_metadata, + layer_cache, + ssm_state_indices=( + mamba_metadata.state_indices[num_prefills:batch_size] + if self._use_indexed_ssm_pool + else None + ), + ) + ) + else: + # Speculative verification: each generation request carries + # 1 + draft_len tokens (drafts are padded to the static max, + # so T is uniform). Per-step states go to the manager's + # SpeculativeState scratch buffers — never the live pools — + # and kv_cache_manager.update_mamba_states() promotes the + # accepted step after sampling. + assert decode_rows % num_decodes == 0, ( + f"ragged generation batch: {decode_rows} tokens for {num_decodes} requests" + ) + outputs.append( + self._forward_verify( + hidden_states[num_ctx_tokens:], + decode_rows // num_decodes, + layer_cache, + conv_pool, + ssm_pool, + state_indices[num_prefills:], + ) + ) + out = outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) + if self._o_allreduce is not None: + # Head-sharded TP: every rank ran its head shard on the same + # local batch; sum the row-sharded o_proj partials. + out = self._o_allreduce(out) + return out + + def _has_kda_replay_caches(self, layer_cache) -> bool: + """True when the manager allocated the fused-verify replay caches.""" + return getattr(layer_cache, "kda_qkg_cache", None) is not None + + def _sync_kda_replay_conv_window( + self, layer_cache, slot_indices, conv_q, conv_k, conv_v + ) -> None: + """Seed the replay conv caches' committed window from FLA windows. + + The fused verify kernel keeps its own extended fp32 dim-contiguous + conv caches; their committed window (columns ``[0, W-1)``) must hold + the last ``W-1`` raw conv inputs whenever another path (prefill, + plain decode) advances the base conv pool. The FLA window's oldest + column drops out of every future convolution, so columns ``[1, W)`` + of the FLA cache map 1:1 onto the committed window. + """ + if not self._has_kda_replay_caches(layer_cache): + return + w = self.mixer.conv_size + for cache, window in ( + (layer_cache.kda_conv_q, conv_q), + (layer_cache.kda_conv_k, conv_k), + (layer_cache.kda_conv_v, conv_v), + ): + cache[:, :, : w - 1].index_copy_(0, slot_indices, window[:, :, 1:].to(cache.dtype)) + + def _forward_prefill( + self, + x2d, + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + slot_indices, + layer_cache=None, + ) -> torch.Tensor: + from einops import rearrange + + mixer = self.mixer + d = self.proj_size + x = x2d.unsqueeze(0) # [1, T, hidden] + + q_proj_states = mixer.q_proj(x) + k_proj_states = mixer.k_proj(x) + v_proj_states = mixer.v_proj(x) + + # Initial states: present for continuation chunks (chunked prefill) + # and for prefix-cache hits (block reuse), where the previous + # conv/recurrent state was onboarded into this request's slot. + conv_q_in = conv_k_in = conv_v_in = None + recurrent_in = None + if mamba_metadata.use_initial_states: + has_init = mamba_metadata.has_initial_states[:num_prefills] + cs = conv_pool.index_select(0, slot_indices) + cs[~has_init] = 0 + conv_q_in, conv_k_in, conv_v_in = _kda_split_conv_sections(cs, d) + recurrent_in = ssm_pool.index_select(0, slot_indices) + recurrent_in[~has_init] = 0 + + q, conv_q = mixer.q_conv1d( + q_proj_states, cache=conv_q_in, output_final_state=True, cu_seqlens=cu_seqlens + ) + k, conv_k = mixer.k_conv1d( + k_proj_states, cache=conv_k_in, output_final_state=True, cu_seqlens=cu_seqlens + ) + v, conv_v = mixer.v_conv1d( + v_proj_states, cache=conv_v_in, output_final_state=True, cu_seqlens=cu_seqlens + ) + + g = mixer.f_b_proj(mixer.f_a_proj(x)) + g = rearrange(g, "... (h d) -> ... h d", d=mixer.head_dim) + beta = mixer.b_proj(x).float() + + q = rearrange(q, "... (h d) -> ... h d", d=mixer.head_k_dim) + k = rearrange(k, "... (h d) -> ... h d", d=mixer.head_k_dim) + v = rearrange(v, "... (h d) -> ... h d", d=mixer.head_dim) + + # Kernel dispatch (in-tree trtllm::kda_prefill or FLA chunk_kda). + # Both paths exchange states in the pool's V-first [N, H, V, K] + # layout, so recurrent_in / final_state map to ssm_pool 1:1. + lower_bound = mixer.gate_lower_bound + o, final_state = mixer.prefill_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=mixer.A_log, + dt_bias=mixer.dt_bias, + scale=mixer.head_k_dim**-0.5, + initial_state=recurrent_in, + safe_gate=lower_bound is not None, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ) + + # Persist per-request states into the pools. + conv_pool.index_copy_( + 0, slot_indices, torch.cat([conv_q, conv_k, conv_v], dim=1).to(conv_pool.dtype) + ) + ssm_pool.index_copy_(0, slot_indices, final_state.to(ssm_pool.dtype)) + # Fused-verify replay caches: seed the committed conv window so the + # first verify round convolves the correct history (pending drafts + # are zero for a fresh request, so the tail columns are unused). + self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_q, conv_k, conv_v) + + return self._output_gate_and_proj(x, o) + + def _forward_decode( + self, + x2d, + conv_pool, + ssm_pool, + slot_indices, + mamba_metadata=None, + layer_cache=None, + ssm_state_indices=None, + ) -> torch.Tensor: + """Plain T=1 decode, fast path. + + Calls ``trtllm::kda_decode`` directly with kernel-native layouts + (nsys 07-24: the reference path spent ~70 us/layer on glue around + the 5 us kernel — 6 separate in-projection GEMV pairs, per-step + re-transposition of constant weights, conv-window slice/roll + copies, per-call torch.arange defaults, and redundant dtype + casts): + + * one fused in-projection GEMV (``finalize_decode_weights``); + * conv windows staged with one gather + one repack copy into a + persistent dense per-section buffer; + * conv-pool write-back with one cat + one index_copy_; + * constant tensors (transposed conv weights, fp32 A_log/dt_bias/ + o_norm weight) reused instead of rebuilt per step. + + The conv windows remain gathered batch-row-dense. When stable + int32 slot indices are supplied, the recurrent-state pool is passed + directly and the CUDA wrapper selects its indexed-state launch; + otherwise the state uses the batch-row-dense static layout. + """ + mixer = self.mixer + if mixer.decode_kernel_path != "optimized" or mixer.wrong_state_layout: + ssm_state_indices = None + if ssm_state_indices is not None: + logger.info_once( + "Kimi K3 KDA indexed recurrent-state pool path is active", + key="kimi_k3_kda_indexed_state_pool", + ) + else: + logger.info_once( + "Kimi K3 KDA static recurrent-state path is active", key="kimi_k3_kda_static_state" + ) + if ( + (self._in_proj_weight is None and self._in_proj_small_weight is None) + or mamba_metadata is None + or ssm_pool.dtype != torch.float32 + ): + return self._forward_decode_ref( + x2d, conv_pool, ssm_pool, slot_indices, layer_cache, ssm_state_indices + ) + + d = self.proj_size + hd = mixer.head_dim + H = mixer.num_heads + B = x2d.shape[0] + W = mixer.conv_size + + # Allocated ONCE at the pool slot count (== per-rank max batch on + # the Mixed manager) and never reallocated: captured CUDA graphs + # hold this pointer, so a later realloc would leave earlier graphs + # writing into freed memory. Footprint: slots x ~9(H=6)..222(H=96) + # KB per layer. + buf = self._cs_dense + if buf is None or buf.shape[1] < B: + if torch.cuda.is_current_stream_capturing(): + # Never allocate inside CUDA graph capture; the reference + # path is capture-safe (just slower). + return self._forward_decode_ref( + x2d, conv_pool, ssm_pool, slot_indices, layer_cache, ssm_state_indices + ) + buf = torch.empty( + 3, max(conv_pool.shape[0], B), d, W - 1, dtype=torch.bfloat16, device=x2d.device + ) + self._cs_dense = buf + + if self._in_proj_weight is not None: + # One GEMV over [q | k | v | g | f_a | b]; slices below are views. + proj = torch.nn.functional.linear(x2d, self._in_proj_weight) + x_qkv = proj[:, : 3 * d] + onorm_g = proj[:, 3 * d : 4 * d] + f_a = proj[:, 4 * d : 4 * d + hd] + beta = proj[:, 4 * d + hd : 4 * d + hd + H] + else: + # FP8 weight read (KIMI_K3_KDA_GLUE_FP8=1): the loader's fused + # FP8 [q | k | v | g] GEMM plus one BF16 GEMV over [f_a | b]; + # slices below are views. + qkvg = mixer.qkvg_proj(x2d) + small = torch.nn.functional.linear(x2d, self._in_proj_small_weight) + x_qkv = qkvg[:, : 3 * d] + onorm_g = qkvg[:, 3 * d : 4 * d] + f_a = small[:, :hd] + beta = small[:, hd : hd + H] + g = mixer.f_b_proj(f_a) # [B, d] + + # Gather the HF-layout conv windows once, then repack the + # historical W-1 columns into the kernel's dense per-section + # [B, d, W-1] layout (single strided copy kernel). + cs = conv_pool.index_select(0, slot_indices) # [B, 3d, W] + cs_dense = buf[:, :B] + cs_dense.copy_(cs.view(B, 3, d, W)[:, :, :, 1:].permute(1, 0, 2, 3)) + + state = ( + ssm_pool if ssm_state_indices is not None else ssm_pool.index_select(0, slot_indices) + ) + + o = mixer._dispatch.decode_kda( + x_q=x_qkv[:, :d].unflatten(-1, (H, hd)).unsqueeze(0), + x_k=x_qkv[:, d : 2 * d].unflatten(-1, (H, hd)).unsqueeze(0), + x_v=x_qkv[:, 2 * d :].unflatten(-1, (H, hd)).unsqueeze(0), + w_q_t=self._w_q_t, + w_k_t=self._w_k_t, + w_v_t=self._w_v_t, + bias_q=None, + bias_k=None, + bias_v=None, + cs_q=cs_dense[0], + cs_k=cs_dense[1], + cs_v=cs_dense[2], + A_log=self._A_log_f32, + g=g.unflatten(-1, (H, hd)).unsqueeze(0), + dt_bias=self._dt_bias_f32, + beta=beta.unsqueeze(0), + state=state, + onorm_g=onorm_g.unflatten(-1, (H, hd)).unsqueeze(0), + onorm_weight=self._onorm_w_f32, + out=None, + ssm_state_indices=ssm_state_indices, + cu_seqlens=mamba_metadata._arange_buffer[: B + 1], + scale=hd**-0.5, + onorm_eps=mixer.o_norm.eps, + lower_bound=mixer.gate_lower_bound, + use_beta_sigmoid_in_kernel=True, + verbose=False, + update_conv_cache=False, + ) + if ssm_state_indices is None: + ssm_pool.index_copy_(0, slot_indices, state) + + # Roll the HF-layout conv pool by one token: new window = + # [old columns 1..W-1, x_new]. One cat + one scatter. + new_win = torch.cat([cs[:, :, 1:], x_qkv.unsqueeze(-1)], dim=-1) + if new_win.dtype != conv_pool.dtype: + new_win = new_win.to(conv_pool.dtype) + conv_pool.index_copy_(0, slot_indices, new_win) + # Fused-verify replay caches (spec decoding only): keep the + # committed conv window in sync with the plain-decode advance. + self._sync_kda_replay_conv_window( + layer_cache, slot_indices, new_win[:, :d], new_win[:, d : 2 * d], new_win[:, 2 * d :] + ) + + return mixer.o_proj(o.view(B, d)) + + def _forward_decode_ref( + self, x2d, conv_pool, ssm_pool, slot_indices, layer_cache=None, ssm_state_indices=None + ) -> torch.Tensor: + from ..modules.kimi_kda.kimi_kda_mixer import KimiKDACachedState + + mixer = self.mixer + d = self.proj_size + x = x2d.unsqueeze(1) # [B, 1, hidden] + + cs = conv_pool.index_select(0, slot_indices) + conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) + cache = KimiKDACachedState( + conv_state_q=conv_q, + conv_state_k=conv_k, + conv_state_v=conv_v, + recurrent_state=( + ssm_pool + if ssm_state_indices is not None + else ssm_pool.index_select(0, slot_indices) + ), + ) + out, new_cache = mixer.forward_decode( + x, + cache, + ssm_state_indices=ssm_state_indices, + ) + + conv_pool.index_copy_( + 0, + slot_indices, + torch.cat( + [ + new_cache.conv_state_q, + new_cache.conv_state_k, + new_cache.conv_state_v, + ], + dim=1, + ).to(conv_pool.dtype), + ) + if ssm_state_indices is None: + ssm_pool.index_copy_(0, slot_indices, new_cache.recurrent_state.to(ssm_pool.dtype)) + # Fused-verify replay caches: keep the committed conv window in + # sync with the plain-decode advance. NOTE: this path is only + # correct for requests with no pending accepted drafts + # (prev_num_accepted_tokens == 0); with drafts pending, the live + # pools lag by the pending prefix and only the fused verify kernel + # can advance them. The spec workers pad drafts to the static max, + # so drafted batches always take the verify path. + self._sync_kda_replay_conv_window( + layer_cache, + slot_indices, + new_cache.conv_state_q, + new_cache.conv_state_k, + new_cache.conv_state_v, + ) + + return out.squeeze(1) + + def _forward_verify( + self, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices + ) -> torch.Tensor: + """Speculative verification: advance each request ``num_steps`` + tokens (1 golden + ``num_steps - 1`` padded drafts). + + Two paths: + + * Fused (``trtllm::kda_mtp_decode``, when the manager allocated the + KDA replay caches): one kernel launch replays the previous + round's accepted drafts from the per-slot replay caches, then + processes the new tokens, committing the recurrent state and conv + windows **in place** after the golden token and caching the new + drafts. ``update_mamba_states()`` afterwards only records the + accepted count for the next round's replay. + * Legacy (sequential per-step FLA): per-step states go to the + manager's batch-row-indexed intermediate scratch buffers and + ``update_mamba_states()`` promotes the accepted step's state + after sampling. + """ + if self._has_kda_replay_caches(layer_cache): + assert self.mixer.verify_kernel_path == "optimized", ( + "KDA replay caches are allocated but the fused verify " + "kernel is unavailable; the legacy intermediate buffers " + "were not allocated so there is no fallback" + ) + return self._forward_verify_fused(x2d, num_steps, layer_cache, ssm_pool, slot_indices) + return self._forward_verify_sequential( + x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices + ) + + def _forward_verify_fused( + self, x2d, num_steps, layer_cache, ssm_pool, slot_indices + ) -> torch.Tensor: + """Fused multi-token verify via ``trtllm::kda_mtp_decode``. + + Token layout: the kernel indexes each request's new tokens at + ``cu_seqlens[n] + num_accepted[n] + i``. The runtime packs the + ``num_steps`` new tokens per request contiguously, so we pass + ``cu_seqlens[n] = n * num_steps - num_accepted[n]`` — the shift + lands the kernel's reads/writes exactly on the packed rows. A + negative entry for request 0 is fine: ``bos`` is only ever used + additively with a token offset ``>= num_accepted``. + """ + mixer = self.mixer + num_decodes = x2d.shape[0] // num_steps + num_spec = num_steps - 1 + H = mixer.num_heads + K = mixer.head_k_dim + x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] + T_total = num_decodes * num_steps + + x_q = mixer.q_proj(x).view(1, T_total, H, K) + x_k = mixer.k_proj(x).view(1, T_total, H, K) + x_v = mixer.v_proj(x).view(1, T_total, H, mixer.head_dim) + # Raw gate / beta: the kernel applies dt_bias, A_log, the + # lower-bound sigmoid gate, and the beta sigmoid itself. + g = mixer.f_b_proj(mixer.f_a_proj(x)).view(1, T_total, H, K) + beta = mixer.b_proj(x).view(1, T_total, H) + + w_q, w_k, w_v = self._get_mtp_conv_weights() + lower_bound = ( + mixer.gate_lower_bound_override + if mixer.gate_lower_bound_override is not None + else mixer.gate_lower_bound + ) + + pending = layer_cache.prev_num_accepted_tokens[slot_indices].to( + torch.int32 + ) # accepted drafts of the previous round, per req + cu_seqlens = torch.arange( + 0, (num_decodes + 1) * num_steps, num_steps, dtype=torch.int32, device=x2d.device + ) + cu_seqlens[:num_decodes].sub_(pending) + + out = mixer._dispatch.mtp_verify( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=layer_cache.kda_conv_q, + cs_k=layer_cache.kda_conv_k, + cs_v=layer_cache.kda_conv_v, + g=g, + beta=beta, + # .detach(): the CuTe DSL DLPack bridge rejects grad-tracking + # tensors. + A_log=mixer.A_log.detach(), + dt_bias=mixer.dt_bias.detach(), + recurrent_state=ssm_pool, + qkg_cache=layer_cache.kda_qkg_cache, + v_cache=layer_cache.kda_v_cache, + beta_cache=layer_cache.kda_beta_cache, + ssm_state_indices=slot_indices.to(torch.int32), + cu_seqlens=cu_seqlens, + num_spec=num_spec, + num_accepted_tokens=pending, + lower_bound=lower_bound, + scale=mixer.head_k_dim**-0.5, + ) + o = out.view(num_decodes, num_steps, H, mixer.head_dim) + return self._output_gate_and_proj(x, o) + + def _get_mtp_conv_weights(self): + """fp32 ``[dim, W]`` conv weights for the fused verify kernel, + computed once per runtime instance.""" + cached = getattr(self, "_mtp_conv_weights", None) + if cached is None: + if torch.cuda.is_current_stream_capturing(): + # Allocating inside CUDA graph capture would bake + # capture-pool pointers into the cached tuple; the constants + # are normally prebuilt by _build_decode_kernel_constants(). + raise RuntimeError( + "Kimi K3 fused-verify conv weights were not prebuilt " + "before CUDA graph capture; call finalize_decode_weights" + "() / finalize_decode_weights_fp8() after weight load." + ) + mixer = self.mixer + cached = tuple( + conv.weight.detach().squeeze(1).float().contiguous() + for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d) + ) + self._mtp_conv_weights = cached + return cached + + def _forward_verify_sequential( + self, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices + ) -> torch.Tensor: + """Sequential per-step FLA verification (legacy intermediate-buffer + path). Live pools are read-only here; ``update_mamba_states()`` + commits the accepted step's state after sampling. + """ + from einops import rearrange + from fla.ops.kda import fused_recurrent_kda + + intermediate_conv = layer_cache.intermediate_conv_window + intermediate_ssm = layer_cache.intermediate_ssm + assert intermediate_conv is not None and intermediate_ssm is not None, ( + "speculative verification requires the cache manager's " + "SpeculativeState (legacy intermediate-buffer path)" + ) + + mixer = self.mixer + d = self.proj_size + num_decodes = x2d.shape[0] // num_steps + x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] + + q_proj_states = mixer.q_proj(x) + k_proj_states = mixer.k_proj(x) + v_proj_states = mixer.v_proj(x) + g = mixer.f_b_proj(mixer.f_a_proj(x)) + g = rearrange(g, "... (h d) -> ... h d", d=mixer.head_dim) + beta = mixer.b_proj(x).float() + + # Gathered copies — mutated across steps, never written back to the + # live pools. + cs = conv_pool.index_select(0, slot_indices) + conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) + state = ssm_pool.index_select(0, slot_indices) + + step_outputs: List[torch.Tensor] = [] + for t in range(num_steps): + # ShortConvolution.step updates the (gathered) caches in place. + q_t, conv_q = mixer.q_conv1d( + q_proj_states[:, t : t + 1], cache=conv_q, output_final_state=True + ) + k_t, conv_k = mixer.k_conv1d( + k_proj_states[:, t : t + 1], cache=conv_k, output_final_state=True + ) + v_t, conv_v = mixer.v_conv1d( + v_proj_states[:, t : t + 1], cache=conv_v, output_final_state=True + ) + + q_t = rearrange(q_t, "... (h d) -> ... h d", d=mixer.head_k_dim) + k_t = rearrange(k_t, "... (h d) -> ... h d", d=mixer.head_k_dim) + v_t = rearrange(v_t, "... (h d) -> ... h d", d=mixer.head_dim) + + o_t, state = fused_recurrent_kda( + q=q_t, + k=k_t, + v=v_t, + g=g[:, t : t + 1], + beta=beta[:, t : t + 1], + A_log=mixer.A_log, + dt_bias=mixer.dt_bias, + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=mixer.gate_lower_bound, + state_v_first=True, + ) + step_outputs.append(o_t) + + # Batch-row indexed ([:num_decodes] prefix), matching + # update_mamba_states()'s intermediate_state_indices. + intermediate_conv[:num_decodes, t] = torch.cat([conv_q, conv_k, conv_v], dim=1).to( + intermediate_conv.dtype + ) + intermediate_ssm[:num_decodes, t] = state.to(intermediate_ssm.dtype) + + o = torch.cat(step_outputs, dim=1) # [B, T, H, V] + return self._output_gate_and_proj(x, o) + + def _output_gate_and_proj(self, x: torch.Tensor, o: torch.Tensor) -> torch.Tensor: + from einops import rearrange + + mixer = self.mixer + if mixer.use_full_rank_gate: + g_out = mixer.g_proj(x) + else: + g_out = mixer.g_b_proj(mixer.g_a_proj(x)) + g_out = rearrange(g_out, "... (h d) -> ... h d", d=mixer.head_dim) + o = mixer.o_norm(o, g_out) + o = rearrange(o, "b t h d -> (b t) (h d)") + return mixer.o_proj(o) + + +class KimiMLARuntime(nn.Module): + """Owns K3 MLA head padding, TP sharding, and output reduction.""" + + def __init__( + self, + cfg, + layer_idx: int, + mapping=None, + quant_config=None, + allreduce_strategy=AllReduceStrategy.AUTO, + ): + super().__init__() + + from ..modules.kimi_k3_mla import KimiK3MLAAttention + + max_positions = int( + os.environ.get( + _KIMI_K3_MLA_MAX_POSITIONS_ENV, + cfg.max_position_embeddings, + ) + ) + self.layer_idx = layer_idx + # The trtllm-gen MLA generation kernels group query heads per CTA and + # require numHeadsQ divisible by the group size (a power of two, up + # to 128). K3's 96 query heads are unsupported, so pad to the next + # power of two (128) with zero weights: padded heads produce exactly + # zero output (their kv_b_proj "v_absorb" rows are zero), so the + # numerics are unchanged at ~33% extra MLA-layer q compute. + self.num_real_heads = cfg.num_attention_heads + padded_heads = 1 + while padded_heads < self.num_real_heads: + padded_heads *= 2 + self.num_padded_heads = padded_heads + # Attention-family TP semantics (DeepSeek MLA pattern, mla.py): + # replicated under attention-DP, head-sharded otherwise. Pad + # FIRST, then divide, so every rank gets a power-of-two head + # count (128/16 = 8; sharding the real 96 would give 6/rank and + # break the generation-FMHA per-CTA head grouping). q_b/kv_b/g + # column-shard and o_proj row-shards by padded head range; ranks + # holding only padded heads contribute exact zeros. The latent KV + # cache and both *_a_proj down-projections stay replicated — with + # a single latent KV head the TP ranks hold duplicated KV cache, + # exactly like DeepSeek MLA under TP (attention-DP dedups it). + self._mla_tp_size = ( + mapping.tp_size + if (mapping is not None and not mapping.enable_attention_dp and mapping.tp_size > 1) + else 1 + ) + self._mla_tp_rank = mapping.tp_rank if self._mla_tp_size > 1 else 0 + assert padded_heads % self._mla_tp_size == 0, ( + f"padded MLA heads {padded_heads} not divisible by tp_size {self._mla_tp_size}" + ) + self._o_allreduce = ( + AllReduce(mapping=mapping, strategy=allreduce_strategy, dtype=torch.bfloat16) + if self._mla_tp_size > 1 + else None + ) + self.mixer = KimiK3MLAAttention( + hidden_size=cfg.hidden_size, + num_heads=padded_heads // self._mla_tp_size, + 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, + quant_config=quant_config, + ) + + def forward( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + out = self.mixer(hidden_states, attn_metadata) + if self._o_allreduce is not None: + # Head-sharded TP: sum the row-sharded o_proj partials across + # the head-shard group. + out = self._o_allreduce(out) + return out + + +# --------------------------------------------------------------------------- +# Decoder layer. +# --------------------------------------------------------------------------- + + +class KimiLinearDecoderLayer(nn.Module): + def __init__( + self, + model_config: ModelConfig, + cfg, + layer_idx: int, + aux_stream: Optional[torch.cuda.Stream] = None, + ): + 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: + self.self_attn = KimiKDARuntime( + cfg, + layer_idx, + mapping=model_config.mapping, + allreduce_strategy=model_config.allreduce_strategy, + ) + else: + # Forward only the KV-cache quantization to the MLA attention + # backends (enables FP8 KV cache). The attention projection + # weights themselves stay BF16 — the model-level weight-quant + # algo must not leak into the attention backend's weight paths. + mla_quant_config = None + kv_quant_algo = ( + model_config.quant_config.kv_cache_quant_algo + if model_config.quant_config is not None + else None + ) + if kv_quant_algo is not None: + mla_quant_config = QuantConfig(kv_cache_quant_algo=kv_quant_algo) + self.self_attn = KimiMLARuntime( + cfg, + layer_idx, + mapping=model_config.mapping, + quant_config=mla_quant_config, + allreduce_strategy=model_config.allreduce_strategy, + ) + + 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) + else: + situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + # Dense-MLP TP semantics (DeepSeek _compute_mlp_tp_size + # pattern): replicated under attention-DP — each rank runs + # only its own tokens, so a weight shard would need an extra + # gather/scatter — and sharded like the shared experts + # otherwise (column gate_up, row down); the partial sums are + # all-reduced right after the call in forward(). + self._mlp_tp_size = ( + model_config.mapping.tp_size + if ( + not model_config.mapping.enable_attention_dp + and model_config.mapping.tp_size > 1 + and cfg.intermediate_size % model_config.mapping.tp_size == 0 + ) + else 1 + ) + self.mlp = KimiK3MLP( + hidden_size=cfg.hidden_size, + intermediate_size=cfg.intermediate_size // self._mlp_tp_size, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + use_fused_activation=True, + dtype=dtype, + ) + self._mlp_allreduce = ( + AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=dtype, + ) + if self._mlp_tp_size > 1 + else 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) + + def forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Port of HF ``KimiDecoderLayer._forward_attn_residual`` (per token). + + Returns ``(prefix_sum, block_residual)`` with the snapshot stack in + kernel-native ``[K, M, H]`` layout; the running prefix sum is the + hidden state handed to the next layer. + """ + prefix_sum = hidden_states + + if block_residual.shape[0] > 0: + hidden_states = _apply_attn_res( + prefix_sum, + block_residual, + self.self_attention_res_proj, + self.self_attention_res_norm, + ) + + if self.layer_idx % self.attn_res_block_size == 0: + block_residual = torch.cat((block_residual, prefix_sum.unsqueeze(0)), dim=0) + prefix_sum = None + + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, attn_metadata) + + if prefix_sum is not None: + prefix_sum = prefix_sum + hidden_states + else: + prefix_sum = hidden_states + + hidden_states = _apply_attn_res( + prefix_sum, block_residual, self.mlp_res_proj, self.mlp_res_norm + ) + + hidden_states = self.post_attention_layernorm(hidden_states) + if self.is_moe: + hidden_states = self.block_sparse_moe( + hidden_states, getattr(attn_metadata, "all_rank_num_tokens", None) + ) + else: + hidden_states = self.mlp(hidden_states) + if getattr(self, "_mlp_allreduce", None) is not None: + # TEP-sharded dense MLP: sum the row-parallel partials. + hidden_states = self._mlp_allreduce(hidden_states) + + prefix_sum = prefix_sum + hidden_states + return prefix_sum, block_residual + + +# --------------------------------------------------------------------------- +# Model. +# --------------------------------------------------------------------------- + + +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 + + # One side stream shared across all layers, used by KimiK3MoERuntime + # to overlap the replicated shared-expert compute with the routed EP + # dispatch/combine collectives. + self.aux_stream = torch.cuda.Stream() + + 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) + 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) + + 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") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + + num_tokens = attn_metadata.num_tokens + assert hidden_states.shape[0] == num_tokens, ( + f"Kimi K3 does not support padded batches " + f"(got {hidden_states.shape[0]} rows, metadata says {num_tokens} " + "tokens); disable CUDA graphs and the overlap scheduler." + ) + + block_residual = hidden_states.new_zeros(0, hidden_states.shape[0], hidden_states.shape[1]) + for layer in self.layers: + hidden_states, block_residual = layer(hidden_states, block_residual, attn_metadata) + if spec_metadata is not None: + # DFlash hidden-state capture. K3's attn-residual scheme + # already folds the residual into the running prefix sum + # returned by each layer, so unlike Qwen3/Llama we pass the + # full hidden state with residual=None. Whether the drafter + # is trained against this prefix sum or some other tap point + # must be confirmed against the K3 drafter training recipe + # before real weights are used. + spec_metadata.maybe_capture_hidden_states(layer.layer_idx, hidden_states, None) + + hidden_states = _apply_attn_res( + hidden_states, block_residual, self.output_attn_res_proj, self.output_attn_res_norm + ) + return self.norm(hidden_states) + + +# --------------------------------------------------------------------------- +# Causal LM wrapper + weight loading. +# --------------------------------------------------------------------------- + + +def _materialize(value) -> torch.Tensor: + """Materialize a (possibly lazy safetensors slice) weight value.""" + if isinstance(value, torch.Tensor): + return value + return value[:] + + +@register_auto_model("KimiK3ForConditionalGeneration") +@register_auto_model("KimiLinearForCausalLM") +class KimiLinearForCausalLM(SpecDecOneEngineForCausalLM[KimiLinearModel, Any]): + """Kimi K3 text model (the vision tower is ignored; text-only serving).""" + + def __init__(self, model_config: ModelConfig): + cfg = _get_text_config(model_config.pretrained_config) + assert model_config.mapping.pp_size == 1, "Kimi K3 does not support pipeline parallelism" + spec_config = getattr(model_config, "spec_config", None) + # Supported spec-dec modes: + # - SA (suffix automaton): one-engine in-forward drafting, no draft + # weights; the KDA/MLA verify paths below implement multi-token + # verification for it. + # - DFlash: external-drafter parallel drafting; the drafter is a + # separate dense checkpoint (K2.7-Code-DFlash schema) consumed by + # the generic DFlashForCausalLM wrapper, and the target only has + # to expose per-layer hidden states via maybe_capture_hidden_states + # (see KimiLinearModel.forward). No trained K3 drafter exists yet; + # this path is exercised with synthetic weights + # (examples/kimi_k3/make_synthetic_dflash_drafter.py). + # Modes needing draft heads (MTP/Eagle) are blocked until a + # draft-head checkpoint exists. + assert ( + spec_config is None + or spec_config.spec_dec_mode.is_sa() + or spec_config.spec_dec_mode.is_dflash() + ), "Kimi K3 supports speculative decoding only with SA or DFlash" + super().__init__( + KimiLinearModel(model_config), + model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + + @classmethod + def get_model_defaults(cls, llm_args) -> dict: + # - enable_block_reuse defaults off: reuse is supported as an + # explicit opt-in (routes to CppMambaHybridCacheManager with + # per-block KDA state snapshots); the default stays on the + # Mixed manager, which SA speculative decoding requires. + # - tokens_per_block=64: with 32, the flashinfer trtllm-gen FMHA lib + # rejects the MLA (576, 512) generation kernel (marked slower) and + # the fallback C++ path requires num_heads % 64 == 0, which K3's + # 96 query heads violate. + return { + "kv_cache_config": { + "enable_block_reuse": False, + "tokens_per_block": 64, + } + } + + # ------------------------------------------------------------------ + # Weight loading (streams the 1.5TB checkpoint; only the rank-local + # expert slice of each MoE layer is kept: whole experts under MoE EP, + # the intra-expert intermediate shard of ALL experts under MoE TP — + # in the TP case every expert tensor is read and sliced, so expect a + # correspondingly longer load). + # ------------------------------------------------------------------ + + def _trunk_parameters(self): + """Named parameters of the trunk only. Spec-dec draft modules + (e.g. the DFlash drafter attached by SpecDecOneEngineForCausalLM) + live in a separate checkpoint loaded by + ModelLoader.load_draft_weights, not in the target checkpoint. MLA + K/V absorb Parameters are derived by the KV-B loader and are likewise + excluded from checkpoint jobs.""" + return { + name: param + for name, param in self.named_parameters() + if not name.startswith("draft_model.") + and not name.endswith(_KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES) + } + + def checkpoint_name_plan(self, prefix: str): + """Return ``(name_map, expected_keys, expert_jobs)``. + + ``name_map`` maps every model parameter name to its checkpoint key + (for fused ``gate_up_proj`` parameters the mapped key is virtual; + the two real per-half keys come from ``_gate_up_ckpt_keys``); + ``expected_keys`` additionally covers the rank-local per-expert MXFP4 + tensors; ``expert_jobs`` lists ``(layer_idx, moe_module, key_base)`` + for backend-owned expert slots. Exposed separately so the weight-name + mapping can be dry-run without touching any tensor data. + """ + params = self._trunk_parameters() + expected_keys = set() + name_map: Dict[str, str] = {} + for name in params: + # ConfigurableMoE's backend owns already-packed runtime weights, + # generated zero biases, and SiTU constants. They do not have + # one-to-one checkpoint parameter names. + if ".routed_experts.backend." in name: + continue + if name == "lm_head.weight": + ckpt_key = prefix + "lm_head.weight" + else: + # Runtime wrapper modules hold the parity-tested mixers as a + # "mixer" submodule; the checkpoint names have no such scope. + ckpt_key = prefix + name.replace(".self_attn.mixer.", ".self_attn.") + name_map[name] = ckpt_key + if name.endswith(_GATE_UP_FUSED_SUFFIX): + # Fused [gate | up] MLP layout (dense mlp / shared_experts): + # the checkpoint stores two separate tensors. + expected_keys.update(_gate_up_ckpt_keys(ckpt_key)) + else: + expected_keys.add(ckpt_key) + + # Backend-owned expert slots (per-expert checkpoint tensors; the + # rank-local id range — an EP slice of whole experts, or ALL experts + # when the routed MoE is TP-sharded (moe_ep=1 -> ids 0..num_experts)). + expert_jobs = [] + for layer_idx, layer in enumerate(self.model.layers): + if not getattr(layer, "is_moe", False): + continue + moe = layer.block_sparse_moe + base = f"{prefix}model.layers.{layer_idx}.block_sparse_moe.experts" + for expert_idx in moe.local_expert_ids: + for w in ("w1", "w2", "w3"): + expected_keys.add(f"{base}.{expert_idx}.{w}.weight_packed") + expected_keys.add(f"{base}.{expert_idx}.{w}.weight_scale") + expert_jobs.append((layer_idx, moe, base)) + return name_map, expected_keys, expert_jobs + + def load_weights(self, weights: Dict): + from .modeling_utils import run_concurrently + + prefix = "language_model." if any(k.startswith("language_model.") for k in weights) else "" + + # The checkpoint stores every MLA KV-B head as interleaved [K | V] + # rows. Runtime keeps one DeepSeek-style [all K | all V] parameter + # instead, so context can project directly into the FMHA layout and + # absorbed decode can take zero-copy K/V views. + mla_mixers = [ + layer.self_attn.mixer + for layer in self.model.layers + if not getattr(layer, "is_kda", True) + ] + mla_kv_b_mixers = {id(mixer.kv_b_proj.weight): mixer for mixer in mla_mixers} + + params = self._trunk_parameters() + name_map, expected_keys, expert_jobs = self.checkpoint_name_plan(prefix) + + # ---- key-set validation (both directions) ---- + ckpt_keys = set(weights.keys()) + relevant_ckpt_keys = { + k + for k in ckpt_keys + if not (k.startswith("vision_tower.") or k.startswith("mm_projector.")) + } + missing = sorted(expected_keys - ckpt_keys) + if missing: + raise KeyError( + f"Kimi K3 load_weights: {len(missing)} expected checkpoint " + f"keys are missing, e.g. {missing[:10]}" + ) + unexpected = relevant_ckpt_keys - expected_keys + # Non-local experts and (in layer-truncated debug mode) extra layers + # are expected leftovers. + surprising = sorted( + k + for k in unexpected + if ".block_sparse_moe.experts." not in k and not k.startswith(f"{prefix}model.layers.") + ) + if surprising: + logger.warning( + f"Kimi K3 load_weights: {len(surprising)} unmatched " + f"checkpoint keys, e.g. {surprising[:10]}" + ) + + device = next(self.parameters()).device + + # MLP TP shard index (used only when a param's checkpoint shape is a + # tp_size multiple of the param shape — the dense L0 MLP with + # attention-DP off; shapes match and no slicing runs otherwise). + # Every mode that shards these fused-MLP tensors shards by tp_rank. + shared_tp_rank = self.model_config.mapping.tp_rank + # KDA head-shard (attention-DP off): rank r loads head rows/cols + # [r*local : (r+1)*local] of every head-major KDA tensor. + kda_tp_size, kda_tp_rank = 1, 0 + for layer in self.model.layers: + if getattr(layer, "is_kda", False): + kda_tp_size = layer.self_attn._kda_tp_size + kda_tp_rank = layer.self_attn._kda_tp_rank + break + # MLA head-shard (attention-DP off): rank slice of the zero-padded + # 128-head layout (pad before divide). + mla_tp_size, mla_tp_rank = 1, 0 + for layer in self.model.layers: + if not getattr(layer, "is_kda", True): + mla_tp_size = layer.self_attn._mla_tp_size + mla_tp_rank = layer.self_attn._mla_tp_rank + break + + def load_param(name: str, param: torch.nn.Parameter): + if device.type == "cuda": + torch.cuda.set_device(device) + if name.endswith(_GATE_UP_FUSED_SUFFIX): + # Row-concat the checkpoint's separate gate_proj / up_proj + # tensors into the fused [gate | up] parameter. + gate_key, up_key = _gate_up_ckpt_keys(name_map[name]) + gate = _materialize(weights[gate_key]) + up = _materialize(weights[up_key]) + inter = param.shape[0] // 2 + if gate.shape[0] != inter and gate.shape[0] % inter == 0: + # TP-sharded fused MLP (shared experts on the direct + # MoE path, dense MLP with attention-DP off): take + # this rank's MATCHING row block from each half so + # the SiTU gate/up pairs stay aligned. shared_tp_rank + # == tp_rank in every mode that shards these. + lo = shared_tp_rank * inter + gate = gate[lo : lo + inter] + up = up[lo : lo + inter] + if gate.shape != (inter, param.shape[1]) or up.shape != gate.shape: + raise ValueError( + f"{name}: checkpoint gate/up shapes " + f"{tuple(gate.shape)} / {tuple(up.shape)} do not " + f"concat to param shape {tuple(param.shape)}" + ) + param.data[:inter].copy_(gate.to(param.dtype)) + param.data[inter:].copy_(up.to(param.dtype)) + return + src = _materialize(weights[name_map[name]]) + if name == "lm_head.weight": + # LMHead is vocab-sharded (TP column) + gathered; its + # load_weights shards the full checkpoint tensor. + self.lm_head.load_weights(weights=[{"weight": src}]) + return + mla_mixer = mla_kv_b_mixers.get(id(param)) + if mla_mixer is not None: + _load_kimi_k3_mla_kv_b_proj( + mla_mixer, + src, + head_start=mla_tp_rank * mla_mixer.num_heads, + ) + return + if name.endswith(".A_log") and src.numel() != param.numel(): + # The checkpoint pads A_log from [num_heads] to [head_dim] + # (e.g. [96] -> [128]); the tail must be zeros. Under KDA + # head-shard TP the param holds this rank's head range + # instead of the full [num_heads]. + assert src.numel() > param.numel(), (name, src.shape) + if kda_tp_size > 1: + lo = kda_tp_rank * param.numel() + src = src[lo : lo + param.numel()] + else: + tail = src[param.numel() :] + if tail.abs().max().item() != 0.0: + raise ValueError( + f"{name}: expected zero padding beyond " + f"{param.numel()} entries, got nonzero tail" + ) + src = src[: param.numel()] + if src.shape != param.shape: + # KDA head-shard (attention-DP off): every mismatching KDA + # tensor is head-major with the checkpoint exactly + # kda_tp_size times larger on one axis — q/k/v/g/f_b + # projections, b_proj, dt_bias, and the depthwise conv + # weights on dim 0 (rows), o_proj on dim 1 (columns). + # MLA layers never produce a x-tp_size ratio (their + # mismatches are the padding branches below), so shape + # ratios alone identify the KDA slices. + if kda_tp_size > 1 and ".self_attn." in name: + if ( + src.shape[0] == param.shape[0] * kda_tp_size + and src.shape[1:] == param.shape[1:] + ): + s = param.shape[0] + lo = kda_tp_rank * s + param.data.copy_(src[lo : lo + s].to(param.dtype)) + return + if ( + src.dim() == 2 + and src.shape[0] == param.shape[0] + and src.shape[1] == param.shape[1] * kda_tp_size + ): + s = param.shape[1] + lo = kda_tp_rank * s + param.data.copy_(src[:, lo : lo + s].to(param.dtype)) + return + # MLA head-shard (attention-DP off): the checkpoint holds + # the real 96 heads; the param holds this rank's slice of + # the zero-PADDED 128-head layout (pad before divide, so + # per-rank counts stay a power of two). Head-major output + # rows for q_b/g_proj and input columns for o_proj. KV-B has + # a dedicated head-aware loader above. A rank whose padded + # range lies beyond the real heads gets zeros. KDA layers + # never reach here: their identically named g/o projections + # match the exact-ratio branch above. + if mla_tp_size > 1 and ".self_attn." in name: + if ( + name.endswith((".q_b_proj.weight", ".g_proj.weight")) + and src.shape[1:] == param.shape[1:] + and src.shape[0] < param.shape[0] * mla_tp_size + ): + s = param.shape[0] + lo = mla_tp_rank * s + param.data.zero_() + n = max(0, min(src.shape[0] - lo, s)) + if n > 0: + param.data[:n].copy_(src[lo : lo + n].to(param.dtype)) + return + if ( + name.endswith(".o_proj.weight") + and src.dim() == 2 + and src.shape[0] == param.shape[0] + and src.shape[1] < param.shape[1] * mla_tp_size + ): + s = param.shape[1] + lo = mla_tp_rank * s + param.data.zero_() + n = max(0, min(src.shape[1] - lo, s)) + if n > 0: + param.data[:, :n].copy_(src[:, lo : lo + n].to(param.dtype)) + return + # Shared-expert TP (direct MoE path): the module holds a + # 1/tp shard of the FFN dim — column shard for gate/up + # (output rows), row shard for down (input columns). + if ".shared_experts." in name or ".mlp." in name: + # Shared experts (direct MoE path) and the dense L0 + # MLP (attention-DP off): the fused gate_up_proj is + # sliced in its dedicated branch above; here the + # unfused halves (if ever configured) and down_proj. + if ( + name.endswith((".gate_proj.weight", ".up_proj.weight")) + and src.shape[0] % param.shape[0] == 0 + and src.shape[1:] == param.shape[1:] + ): + lo = shared_tp_rank * param.shape[0] + param.data.copy_(src[lo : lo + param.shape[0]].to(param.dtype)) + return + if ( + name.endswith(".down_proj.weight") + and src.shape[1] % param.shape[1] == 0 + and src.shape[0] == param.shape[0] + ): + lo = shared_tp_rank * param.shape[1] + param.data.copy_(src[:, lo : lo + param.shape[1]].to(param.dtype)) + return + # MLA head padding (96 -> 128 query heads, see + # KimiMLARuntime): pad the head-major output rows + # (q_b_proj / g_proj) or the head-major input columns + # (o_proj) with zeros. KV-B is handled above. KDA layers' + # identically named projections match exactly and never take + # this path. + if ( + ".self_attn." in name + and name.endswith((".q_b_proj.weight", ".g_proj.weight")) + and src.shape[1:] == param.shape[1:] + and src.shape[0] < param.shape[0] + ): + param.data.zero_() + param.data[: src.shape[0]].copy_(src.to(param.dtype)) + return + if ( + ".self_attn." in name + and name.endswith(".o_proj.weight") + and src.shape[0] == param.shape[0] + and src.shape[1] < param.shape[1] + ): + param.data.zero_() + param.data[:, : src.shape[1]].copy_(src.to(param.dtype)) + return + raise ValueError( + f"{name}: checkpoint shape " + f"{tuple(src.shape)} != param shape " + f"{tuple(param.shape)}" + ) + param.data.copy_(src.to(param.dtype)) + + def load_expert( + moe: KimiK3MoERuntime, base: str, local_slot_id: int, expert_idx: int, get_tensor + ): + if device.type == "cuda": + torch.cuda.set_device(device) + backend = moe.routed_experts.backend + 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_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str): + del layer_idx + for local_slot_id, expert_idx in enumerate(moe.local_expert_ids): + load_expert( + moe, + base, + local_slot_id, + expert_idx, + lambda key: _materialize(weights[key]), + ) + + param_jobs = [(name, params[name]) for name in name_map] + run_concurrently(load_param, param_jobs, num_workers=8) + + logger.info( + f"Kimi K3: loaded {len(mla_mixers)} MLA KV-B projections in grouped runtime layout" + ) + + # ---- backend expert slots: file-grouped streaming ---- + # The shared lazy ``weights`` dict keeps every shard mmapped for the + # whole load, so pages it touches cannot be dropped until the load + # ends (fadvise skips mapped pages). The expert slices are ~90 GB of + # DISTINCT pages per rank — with 4 ranks/node that overruns the job + # cgroup and OOM-kills the step (observed repeatedly on GB300 + # trays). Instead, group the rank-local expert tensors by shard file + # and stream each file through a short-lived handle: + # open -> copy -> close (unmap) -> fadvise(DONTNEED). + ckpt_dir = getattr(self.model_config.pretrained_config, "_name_or_path", None) + index_path = os.path.join(ckpt_dir or "", "model.safetensors.index.json") + if expert_jobs and ckpt_dir and os.path.isfile(index_path): + import json as _json + from contextlib import ExitStack + + from safetensors import safe_open + + with open(index_path) as f: + weight_map = _json.load(f)["weight_map"] + per_file: Dict[str, list] = {} + split_file_jobs = [] + for layer_idx, moe, base in expert_jobs: + del layer_idx + for local_slot_id, expert_idx in enumerate(moe.local_expert_ids): + keys = [ + f"{base}.{expert_idx}.{w}.{kind}" + for w in ("w1", "w2", "w3") + for kind in ("weight_packed", "weight_scale") + ] + files = {weight_map[key] for key in keys} + job = (moe, base, local_slot_id, expert_idx) + if len(files) == 1: + per_file.setdefault(files.pop(), []).append(job) + else: + split_file_jobs.append((job, files)) + + def drop_file_pages(file_name: str): + path = os.path.join(ckpt_dir, file_name) + try: + fd = os.open(path, os.O_RDONLY) + try: + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(fd) + except OSError: + pass + + def load_expert_file(file_name: str, jobs: list): + if device.type == "cuda": + torch.cuda.set_device(device) + path = os.path.join(ckpt_dir, file_name) + with safe_open(path, framework="pt", device="cpu") as fh: + for moe, base, local_slot_id, expert_idx in jobs: + load_expert(moe, base, local_slot_id, expert_idx, fh.get_tensor) + # Handle closed -> pages unmapped -> the drop takes effect. + drop_file_pages(file_name) + + def load_split_file_expert(job, files): + if device.type == "cuda": + torch.cuda.set_device(device) + with ExitStack() as stack: + handles = { + file_name: stack.enter_context( + safe_open( + os.path.join(ckpt_dir, file_name), framework="pt", device="cpu" + ) + ) + for file_name in files + } + + def get_tensor(key): + return handles[weight_map[key]].get_tensor(key) + + load_expert(*job, get_tensor) + for file_name in files: + drop_file_pages(file_name) + + run_concurrently(load_expert_file, sorted(per_file.items()), num_workers=4) + run_concurrently(load_split_file_expert, split_file_jobs, num_workers=4) + else: + run_concurrently(load_experts_from_weights, expert_jobs, num_workers=4) + + for _, moe, _ in expert_jobs: + backend = moe.routed_experts.backend + loaded_slots = getattr(backend, "_packed_mxfp4_loaded_slots", set()) + expected_slots = set(range(backend.expert_size_per_partition)) + if loaded_slots != expected_slots: + missing_slots = sorted(expected_slots - loaded_slots) + raise RuntimeError( + "Kimi K3 packed expert loading did not fill all backend " + f"slots; missing {missing_slots[:10]}." + ) + backend._weights_transformed = False + + # FP8 weight-read master switch (see the conversion block below). + # The KDA conversion replaces the decode in-projection GEMV with a + # fused FP8 qkvg GEMM in the mixer decode path, so when it is enabled + # the bf16 wrapper fast path (finalize_decode_weights) is NOT built: + # both fuse the same projections and the wrapper path — checked first + # at decode — would bypass the FP8 modules entirely, leaving the FP8 + # copies resident but inert. KIMI_K3_FP8_WEIGHT_READ_KDA=0 restores + # the bf16 wrapper fast path; KIMI_K3_KDA_GLUE_FP8=1 instead rebuilds + # the wrapper fast path on top of the FP8 modules after the + # conversion (finalize_decode_weights_fp8), so neither is traded away. + fp8_weight_read = is_sm_100f() and os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_ENV, "1") != "0" + kda_fp8 = fp8_weight_read and os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0" + kda_glue_fp8 = kda_fp8 and os.environ.get(_KIMI_K3_KDA_GLUE_FP8_ENV, "1") != "0" + + # Build the KDA decode fast-path constants (fused in-projection + # weight views + kernel-layout conv weights + fp32 params). Must + # run after every KDA parameter is loaded/sharded. + num_kda_fused = 0 + for layer in self.model.layers: + if getattr(layer, "is_kda", False): + if not kda_fp8: + layer.self_attn.finalize_decode_weights() + num_kda_fused += int(layer.self_attn._in_proj_weight is not None) + logger.info( + f"Kimi K3: loaded {len(param_jobs)} parameters and the expert " + f"slices of {len(expert_jobs)} MoE layers; fused decode " + f"in-projections on {num_kda_fused} KDA layers" + ) + + # FP8 block-scale weight read for the replicated MoE-layer MLPs. The + # DeepGEMM fp8_swap_ab_gemm kernel is Blackwell-only; keep BF16 on any + # other SM or when explicitly disabled. + if fp8_weight_read: + gate_up_default = "1" if self.model_config.mapping.enable_attention_dp else "0" + n_fp8 = _convert_moe_mlps_to_fp8_weight_read( + self.model, + include_fused_gate_up=os.environ.get( + _KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV, gate_up_default + ) + != "0", + ) + logger.info( + f"Kimi K3: reading {n_fp8} MoE-layer MLP projections " + f"(shared-expert + latent) at FP8 block-scale" + ) + # The KDA q/k/v/g/o projections are the largest single replicated + # weight read; convert them to the same FP8 block-scale read unless + # kept in BF16 for accuracy (their own switch — the recurrent core + # is the most precision-sensitive slice). + if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0": + n_kda = _convert_kda_projections_to_fp8_weight_read(self.model) + logger.info( + f"Kimi K3: reading {n_kda} KDA q/k/v/g/o projections " + f"at FP8 block-scale (q/k/v/g fused into one decode GEMM " + f"per layer)" + ) + if kda_glue_fp8: + # Rebuild the wrapper decode fast path on top of the FP8 + # modules (must run after the conversion above so the + # fused FP8 qkvg_proj exists and only [f_a | b] is fused + # in bf16). + n_glue = 0 + for layer in self.model.layers: + if getattr(layer, "is_kda", False): + layer.self_attn.finalize_decode_weights_fp8() + n_glue += int(layer.self_attn._in_proj_small_weight is not None) + logger.info( + f"Kimi K3: FP8 fused decode glue on {n_glue} KDA " + f"layers ({_KIMI_K3_KDA_GLUE_FP8_ENV}=1)" + ) + # The MLA q_a/q_b/o and output-gate projections are the remaining + # replicated attention weight read the MLP and KDA passes above + # leave in BF16; convert them to the same FP8 block-scale read + # (kv_a/kv_b stay BF16 — see the switch's comment) unless kept in + # BF16 for accuracy. + if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_MLA_ENV, "1") != "0": + n_mla = _convert_mla_projections_to_fp8_weight_read(self.model) + logger.info( + f"Kimi K3: reading {n_mla} MLA q_a/q_b/o/g projections at FP8 block-scale" + ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index 92b0425384b4..4ab178c75239 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -88,6 +88,7 @@ def create_strategy( alltoall_result_do_sum: bool = True, use_flashinfer: bool = False, hidden_size: Optional[int] = None, + communication_method: Optional[str] = None, ) -> Optional[Communication]: """ Create the best communication method for the given configuration @@ -113,6 +114,8 @@ def create_strategy( hidden_size: Actual MoE activation dimension (the A2A payload width). For latent-MoE models this is moe_latent_size, not pretrained_config.hidden_size. Falls back to pretrained_config.hidden_size when not provided. + communication_method: Optional model-selected communication method. + ``TRTLLM_FORCE_COMM_METHOD`` takes precedence when set. # TODO: Need a way to indicate whether EPLB is enabled. Returns: @@ -143,7 +146,7 @@ def create_strategy( return AllGatherReduceScatter(mapping) # Check if forced method is specified via environment variable - force_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD") + force_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD", communication_method) if force_method is not None: return CommunicationFactory._create_forced_method( @@ -265,6 +268,13 @@ def create_strategy( # Try DeepEPLowLatency as fallback when DeepEP is not available try: + if top_k > DeepEPLowLatency.MAX_TOP_K: + raise ValueError( + f"top_k={top_k} exceeds the low-latency kernels' " + f"compile-time cap MAX_TOP_K={DeepEPLowLatency.MAX_TOP_K} " + "(kNumMaxTopK in internode_ll.cu); the kernel-side " + "EP_HOST_ASSERT would abort on the first dispatch/combine" + ) strategy = DeepEPLowLatency( mapping, num_slots, @@ -378,6 +388,10 @@ def _create_forced_method( use_cuda_graph, ) elif method == "DEEPEPLOWLATENCY": + if top_k > DeepEPLowLatency.MAX_TOP_K: + raise ValueError( + f"DeepEPLowLatency supports top_k <= {DeepEPLowLatency.MAX_TOP_K}, got {top_k}." + ) return DeepEPLowLatency( mapping, num_slots, diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py index 2c0f366febf9..a9923d61ce7d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.py @@ -48,6 +48,15 @@ class DeepEPLowLatency(Communication): Sourced from SWITCH_HIDDEN_FOR_EXTENSION_KERNELS in extension_kernels.cu. """ + MAX_TOP_K: int = 9 + """int: Compile-time top-k cap of the low-latency kernels. + + ``kNumMaxTopK``/``kNumMaxTopk`` in internode_ll.cu (dispatch and combine) + size per-thread register arrays with it and guard it with + ``EP_HOST_ASSERT(num_topk <= kNumMaxTopK)`` — a larger top_k aborts on the + first dispatch/combine, so selection must reject it up front. + """ + def __init__( self, mapping: Mapping, diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 5db7e5d90523..11c648d59f49 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -37,7 +37,12 @@ from tensorrt_llm._torch.modules.fused_moe.interface import MoE, MoESchedulerKind from tensorrt_llm._torch.modules.fused_moe.routing import BaseMoeRoutingMethod from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager -from tensorrt_llm._torch.utils import AuxStreamType, EventType, Fp4QuantizedTensor +from tensorrt_llm._torch.utils import ( + ActType_TrtllmGen, + AuxStreamType, + EventType, + Fp4QuantizedTensor, +) from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig @@ -158,6 +163,10 @@ def __init__( apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, override_quant_config: Optional["QuantConfig"] = None, + trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, + trtllm_gen_activation_alpha: Optional[float] = None, + trtllm_gen_activation_beta: Optional[float] = None, + communication_method: Optional[str] = None, **kwargs, ): super().__init__( @@ -179,6 +188,7 @@ def __init__( # Store model_config and aux_stream_dict for later use (e.g., backend setter) self.model_config = model_config self.aux_stream_dict = aux_stream_dict + self.communication_method = communication_method # If True, the router weight will be multiplied on the input rather than at the end of FC2 self.apply_router_weight_on_input = apply_router_weight_on_input @@ -188,6 +198,9 @@ def __init__( model_config=model_config, routing_method=routing_method, override_quant_config=override_quant_config, + trtllm_gen_activation_type=trtllm_gen_activation_type, + trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, + trtllm_gen_activation_beta=trtllm_gen_activation_beta, **kwargs, ) @@ -271,6 +284,9 @@ def _create_and_sync_backend( model_config: ModelConfig, routing_method: BaseMoeRoutingMethod, override_quant_config: Optional["QuantConfig"], + trtllm_gen_activation_type: Optional[ActType_TrtllmGen], + trtllm_gen_activation_alpha: Optional[float], + trtllm_gen_activation_beta: Optional[float], **kwargs, ) -> None: """Build the MoE backend, mirror EPLB attrs, then create weights. @@ -325,6 +341,9 @@ def _create_and_sync_backend( swiglu_limit_scalar=kwargs.get("swiglu_limit_scalar"), init_load_balancer=False, activation_type=self.activation_type, + trtllm_gen_activation_type=trtllm_gen_activation_type, + trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, + trtllm_gen_activation_beta=trtllm_gen_activation_beta, ) # Backend acceptance is validated at the end of ``__init__`` instead @@ -557,6 +576,7 @@ def _create_comm_strategy_auto(self) -> Optional[Communication]: alltoall_result_do_sum=True, use_flashinfer=self.use_flashinfer, hidden_size=self.hidden_size, + communication_method=self.communication_method, ) def forward_impl( diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 1d22fc1ac47c..d8c512769b84 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -9,7 +9,7 @@ from ...model_config import ModelConfig from ...peft.lora.validation import check_moe_lora_supported -from ...utils import ActivationType, AuxStreamType +from ...utils import ActivationType, ActType_TrtllmGen, AuxStreamType from .configurable_moe import ConfigurableMoE from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cute_dsl_b12x import CuteDslB12xFusedMoE @@ -261,6 +261,9 @@ def create_moe_backend( swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = True, activation_type: ActivationType = ActivationType.Swiglu, + trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, + trtllm_gen_activation_alpha: Optional[float] = None, + trtllm_gen_activation_beta: Optional[float] = None, ) -> MoE: """ Create MoE backend instance with validation. @@ -284,6 +287,9 @@ def create_moe_backend( swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4) swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8) activation_type: Activation type + trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type + trtllm_gen_activation_alpha: Optional backend-local activation alpha + trtllm_gen_activation_beta: Optional backend-local activation beta Returns: MoE: MoE backend instance @@ -369,7 +375,17 @@ def create_moe_backend( swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, activation_type=activation_type, + trtllm_gen_activation_type=trtllm_gen_activation_type, + trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, + trtllm_gen_activation_beta=trtllm_gen_activation_beta, ) + + if any(value is not None for value in (trtllm_gen_activation_type, + trtllm_gen_activation_alpha, + trtllm_gen_activation_beta)): + raise ValueError( + "TRTLLM-Gen backend-local activation options are only supported " + f"by TRTLLMGenFusedMoE, got {moe_cls.__name__}") elif moe_cls in (CutlassFusedMoE, MarlinFusedMoE): # CuteDslFusedMoE, DeepGemmFusedMoE, and CuteDslB12xFusedMoE # also subclass CutlassFusedMoE but have narrower constructors, so @@ -549,6 +565,10 @@ def create_moe( swiglu_limit: Optional[torch.Tensor] = None, swiglu_limit_scalar: Optional[float] = None, activation_type: ActivationType = ActivationType.Swiglu, + trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, + trtllm_gen_activation_alpha: Optional[float] = None, + trtllm_gen_activation_beta: Optional[float] = None, + communication_method: Optional[str] = None, ) -> MoE: """ Create MoE instance with automatic parameter inference from model_config. @@ -572,6 +592,10 @@ def create_moe( swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4) swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8) activation_type: Activation type + trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type + trtllm_gen_activation_alpha: Optional backend-local activation alpha + trtllm_gen_activation_beta: Optional backend-local activation beta + communication_method: Optional ConfigurableMoE communication method Returns: MoE: MoE instance @@ -597,6 +621,14 @@ def create_moe( moe_cls = resolve_moe_cls(model_config, routing_method, dtype, override_quant_config, layer_idx) + if (any(value is not None for value in (trtllm_gen_activation_type, + trtllm_gen_activation_alpha, + trtllm_gen_activation_beta)) + and moe_cls is not TRTLLMGenFusedMoE): + raise ValueError( + "A TRTLLM-Gen backend-local activation requires " + "TRTLLMGenFusedMoE without backend fallback, but resolved " + f"{moe_cls.__name__}.") if moe_cls in (DeepGemmFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, CuteDslB12xFusedMoE, CutlassFusedMoE, DenseGEMMFusedMoE, @@ -620,10 +652,16 @@ def create_moe( swiglu_limit=swiglu_limit, swiglu_limit_scalar=swiglu_limit_scalar, activation_type=activation_type, + trtllm_gen_activation_type=trtllm_gen_activation_type, + trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, + trtllm_gen_activation_beta=trtllm_gen_activation_beta, + communication_method=communication_method, ) # WideEPMoE, TritonFusedMoE and VanillaMoE are not wrapped by ConfigurableMoE # and own their communication and forward paths. + if communication_method is not None: + raise ValueError("communication_method requires ConfigurableMoE.") return create_moe_backend( moe_cls=moe_cls, routing_method=routing_method, @@ -643,4 +681,7 @@ def create_moe( swiglu_limit=swiglu_limit, swiglu_limit_scalar=swiglu_limit_scalar, activation_type=activation_type, + trtllm_gen_activation_type=trtllm_gen_activation_type, + trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, + trtllm_gen_activation_beta=trtllm_gen_activation_beta, ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 384fbe0a48cb..a1bf48be87a6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -28,7 +28,8 @@ from ...custom_ops.trtllm_gen_custom_ops import \ fp4_block_scale_fake_output_without_finalize from ...model_config import ModelConfig -from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor +from ...utils import (ActivationType, ActType_TrtllmGen, AuxStreamType, + Fp4QuantizedTensor) from ..gated_mlp import GatedMLP from .interface import MoE, MoEWeightLoadingMode from .moe_op_backend import MoEOpBackend, TRTLLMOpBackend, get_op_backend @@ -206,6 +207,9 @@ def __init__( swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = True, activation_type: ActivationType = ActivationType.Swiglu, + trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, + trtllm_gen_activation_alpha: Optional[float] = None, + trtllm_gen_activation_beta: Optional[float] = None, ): super().__init__( routing_method=routing_method, @@ -227,6 +231,13 @@ def __init__( activation_type=activation_type, ) + self.trtllm_gen_activation_type = ( + ActType_TrtllmGen(trtllm_gen_activation_type) + if trtllm_gen_activation_type is not None else None) + self.trtllm_gen_activation_alpha = trtllm_gen_activation_alpha + self.trtllm_gen_activation_beta = trtllm_gen_activation_beta + self._validate_backend_local_activation() + # Cached for autotune profile sizing (forward path passes # tune_max_num_tokens to the MoE op). self.max_num_tokens = model_config.max_num_tokens @@ -297,6 +308,8 @@ def __init__( def _to_trtllm_gen_activation_type(self, activation_type: ActivationType) -> int: + if self.trtllm_gen_activation_type is not None: + return int(self.trtllm_gen_activation_type) if activation_type == ActivationType.Swiglu: return 0 elif activation_type == ActivationType.SwigluBias: @@ -310,6 +323,76 @@ def _to_trtllm_gen_activation_type(self, else: raise ValueError(f"Unsupported activation type: {activation_type}") + @property + def is_situ_activation(self) -> bool: + return self.trtllm_gen_activation_type == ActType_TrtllmGen.SiTu + + def _validate_backend_local_activation(self) -> None: + if self.trtllm_gen_activation_type is None: + if (self.trtllm_gen_activation_alpha is not None + or self.trtllm_gen_activation_beta is not None): + raise ValueError( + "TRTLLM-Gen backend-local activation alpha/beta require " + "trtllm_gen_activation_type.") + return + + if not self.is_situ_activation: + raise ValueError( + "Only the SiTu TRTLLM-Gen backend-local activation is " + f"supported, got {self.trtllm_gen_activation_type.name}.") + if self.dtype != torch.bfloat16: + raise ValueError( + "TRTLLM-Gen SiTu requires bfloat16 activations, got " + f"{self.dtype}.") + if get_sm_version() not in {100, 103}: + raise ValueError("TRTLLM-Gen SiTu requires SM100 or SM103, got " + f"SM{get_sm_version()}.") + if (self.quant_config is None + or self.quant_config.quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8): + quant_algo = (None if self.quant_config is None else + self.quant_config.quant_algo) + raise ValueError( + "TRTLLM-Gen SiTu requires W4A8_MXFP4_MXFP8 quantization, " + f"got {quant_algo}.") + if self.tp_size > 1: + # Intra-expert MoE TP: w1/w3 column-shard and w2 row-shard along + # the intermediate dim (the stock MXFP4 quant-method loaders slice + # the group-32 packed bytes and scales per rank). Require the + # per-rank shard to stay a whole multiple of the quant method's + # weight alignment so per-shard scale groups and the padded + # weight buffers line up without fractional groups. + alignment = W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod.weight_alignment + if (self.intermediate_size % self.tp_size != 0 + or self.intermediate_size_per_partition % alignment != 0): + raise ValueError( + "TRTLLM-Gen SiTu MoE TP requires intermediate_size " + f"({self.intermediate_size}) divisible by moe_tp_size " + f"({self.tp_size}) with the per-rank shard a multiple of " + f"{alignment}, got " + f"{self.intermediate_size_per_partition}.") + if self.activation_type != ActivationType.Swiglu: + raise ValueError( + "TRTLLM-Gen SiTu must use generic SwiGLU geometry so FC1 " + "contains gate and up projections.") + if self.bias or any( + value is not None + for value in (self.swiglu_alpha, self.swiglu_beta, + self.swiglu_limit, self.swiglu_limit_scalar)): + raise ValueError( + "TRTLLM-Gen SiTu does not support bias or SwiGLU-specific " + "alpha/beta/limit parameters.") + if (self.trtllm_gen_activation_alpha is None + or self.trtllm_gen_activation_beta is None): + raise ValueError( + "TRTLLM-Gen SiTu requires both backend-local activation " + "alpha and beta.") + if (self.trtllm_gen_activation_alpha <= 0.0 + or self.trtllm_gen_activation_beta <= 0.0): + raise ValueError( + "TRTLLM-Gen SiTu activation alpha/beta must be positive, got " + f"{self.trtllm_gen_activation_alpha} and " + f"{self.trtllm_gen_activation_beta}.") + @staticmethod def _is_flashinfer_fused_moe_available() -> bool: try: @@ -331,6 +414,11 @@ def _requires_separated_routing(self) -> bool: return not isinstance(self.routing_method, DeepSeekV3MoeRoutingMethod) def _check_flashinfer_backend_support(self) -> bool: + # SiTu is provided by the native TRTLLM-Gen cubin and is not part of + # FlashInfer's activation enum. + if self.is_situ_activation: + return False + # For BF16 (unquantized) path, we will use FlashInfer regardless whether # env TRTLLM_GEN_FUSED_MOE_USE_FLASHINFER=1 is set or not as it's the only way. if self._is_unquantized_path(): @@ -427,6 +515,26 @@ def _check_configs(self): "TRTLLMGenFusedMoE FP8 block-scale path only supports the uniform " \ "swiglu_limit_scalar, not a per-expert swiglu_limit tensor." + if self.is_situ_activation: + if not isinstance(self.op_backend, TRTLLMOpBackend): + raise ValueError( + "TRTLLM-Gen SiTu requires the native TRTLLM op backend.") + if not self.has_w4a8_mxfp4_mxfp8: + raise ValueError( + "TRTLLM-Gen SiTu requires the W4A8_MXFP4_MXFP8 path.") + if self.scaling_vector_size != 32: + raise ValueError( + "TRTLLM-Gen SiTu requires MXFP4 scaling vector size 32, " + f"got {self.scaling_vector_size}.") + for name in ("situ_alpha", "situ_beta"): + value = getattr(self, name) + if (value.dtype != torch.float32 + or value.shape != (self.expert_size_per_partition, ) + or not value.is_contiguous()): + raise ValueError( + f"{name} must be a contiguous float32 tensor with " + "one value per local expert/slot.") + def _get_quant_method(self): if self.quant_config is not None and self.quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True): @@ -463,6 +571,20 @@ def create_weights(self): else: self.quant_method.create_weights(self) + if self.is_situ_activation: + situ_alpha = nn.Parameter(torch.full( + (self.expert_size_per_partition, ), + float(self.trtllm_gen_activation_alpha), + dtype=torch.float32), + requires_grad=False) + situ_beta = nn.Parameter(torch.full( + (self.expert_size_per_partition, ), + float(self.trtllm_gen_activation_beta), + dtype=torch.float32), + requires_grad=False) + self.register_parameter("situ_alpha", situ_alpha) + self.register_parameter("situ_beta", situ_beta) + self._weights_created = True self._check_configs() @@ -480,6 +602,14 @@ def create_weights(self): requires_grad=False) self.register_parameter("w2_bias", self.w2_bias) + def cache_derived_state(self) -> None: + super().cache_derived_state() + if self.is_situ_activation: + # Reinitialize constants after meta-device materialization. These + # are backend configuration, not checkpoint weights. + self.situ_alpha.data.fill_(float(self.trtllm_gen_activation_alpha)) + self.situ_beta.data.fill_(float(self.trtllm_gen_activation_beta)) + def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): @@ -742,12 +872,16 @@ def run_moe( # When output is provided, use it directly as the result final_hidden_states = moe_output if moe_output is not None else result elif self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_mxfp4_mxfp8: - factor = 1 if self.activation_type in [ - ActivationType.Relu2, ActivationType.Silu + act_type = self._to_trtllm_gen_activation_type(self.activation_type) + factor = 1 if act_type in [ + ActType_TrtllmGen.Relu2, ActType_TrtllmGen.Silu ] else 2 intermediate_size_per_partition_padded = self.w3_w1_weight.shape[ -2] // factor - act_type = self._to_trtllm_gen_activation_type(self.activation_type) + gemm1_alpha = (self.situ_alpha + if self.is_situ_activation else self.swiglu_alpha) + gemm1_beta = (self.situ_beta + if self.is_situ_activation else self.swiglu_beta) output1_scale_scalar = self._get_data_or_none("fc31_scale_c") output1_scale_gate_scalar = self._get_data_or_none("fc31_alpha") @@ -761,8 +895,8 @@ def run_moe( self.w3_w1_weight, self.w3_w1_weight_scale, self.w3_w1_bias if self.bias else None, - self.swiglu_alpha, - self.swiglu_beta, + gemm1_alpha, + gemm1_beta, self.swiglu_limit, self.w2_weight, self.w2_weight_scale, diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py index 79be5d74881a..3f37c694615d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py @@ -25,6 +25,8 @@ import torch +from tensorrt_llm._torch.utils import ActType_TrtllmGen + # Global registry for MoE backends _MOE_OP_BACKEND_REGISTRY: Dict[str, Type["MoEOpBackend"]] = {} @@ -332,6 +334,37 @@ def run_fp4_block_scale_moe( use_dp=False, ): hidden_size = gemm1_weights.shape[-1] * 2 + if gated_act_type == int(ActType_TrtllmGen.SiTu): + if hidden_states.dtype != torch.float8_e4m3fn or hidden_states_scale is None: + raise ValueError( + "TRTLLM-Gen SiTu expects dynamically quantized MXFP8 " + "activations with block scales, got " + f"dtype={hidden_states.dtype}, " + f"has_scale={hidden_states_scale is not None}." + ) + if gemm1_weights_scale is None or gemm1_weights_scale.shape[-1] != hidden_size // 32: + raise ValueError("TRTLLM-Gen SiTu requires MXFP4 weights with group-32 scales.") + if gemm1_alpha is None or gemm1_beta is None: + raise ValueError( + "TRTLLM-Gen SiTu requires per-local-expert alpha and beta tensors." + ) + for name, value in (("gemm1_alpha", gemm1_alpha), ("gemm1_beta", gemm1_beta)): + if ( + value.dtype != torch.float32 + or not value.is_contiguous() + or value.numel() != local_num_experts + ): + raise ValueError( + f"{name} must be contiguous float32 with {local_num_experts} elements." + ) + if gemm1_clamp_limit is not None: + raise ValueError("TRTLLM-Gen SiTu does not use gemm1_clamp_limit.") + if valid_hidden_size is None or valid_intermediate_size is None: + raise ValueError( + "TRTLLM-Gen SiTu requires valid hidden and intermediate " + "sizes for padded MXFP4 weights." + ) + if hidden_states.dtype == torch.uint8 or hidden_states.dtype == torch.float8_e4m3fn: if ( gemm1_weights_scale is not None diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index b0dccb1dc57b..72d355f97289 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -42,6 +42,7 @@ from __future__ import annotations +import os from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union @@ -61,6 +62,13 @@ from .fused_moe_trtllm_gen import TRTLLMGenFusedMoE from .interface import MoESchedulerKind +# Route on the host (fused noaux_tc + post-topk pipeline) instead of inside +# the trtllm-gen cubin. The in-cubin top-k tier for large expert counts +# (896 experts / top-16) register-spills and costs ~33 us/layer at decode +# batch 5..64 vs ~10 us for the post-topk pipeline; the separated path is +# the same math the attention-DP deployments already run. +FORCE_SEPARATED_ROUTING = os.environ.get("TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING", "0") == "1" + __all__ = [ "MoEScheduler", "ExternalCommMoEScheduler", @@ -382,6 +390,7 @@ def _forward_chunk_impl( moe.backend._supports_load_balancer() or moe.routing_method.requires_separated_routing or moe.comm is not None + or FORCE_SEPARATED_ROUTING ) if requires_separated_routing: # Separated routing: ConfigurableMoE calls routing_method @@ -831,7 +840,15 @@ def _get_backend_kwargs( # When the scheduler precomputes top-k for DP/load-balancer paths, # the backend must not route again. Single-rank TRTLLMGen paths do # not get precomputed top-k, so they still need router_logits. - router_logits_arg = None if moe.backend._supports_load_balancer() else router_logits + router_logits_arg = ( + None + if ( + moe.backend._supports_load_balancer() + or moe.routing_method.requires_separated_routing + or FORCE_SEPARATED_ROUTING + ) + else router_logits + ) kwargs["router_logits"] = router_logits_arg kwargs["do_finalize"] = do_finalize kwargs["moe_output"] = self._get_nvlink_onesided_moe_output( diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 943e629ad396..e1afc5950807 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -5916,6 +5916,11 @@ def load_quant_scales(self, module: torch.nn.Module, weights: Dict): self.BLOCK_SCALES_DTYPE).reshape(orig_w2_int32_shape)) +# Serializes the duplicate-check + slot-claim step of +# ``load_packed_mxfp4_expert`` across the loader thread pool. +_PACKED_MXFP4_SLOT_CLAIM_LOCK = threading.Lock() + + class MXFP4WeightTRTLLMGenFusedMoEMethod(MXFP4WeightFusedMoEMethod): weight_dtype = torch.uint8 block_scales_dtype = torch.uint8 @@ -6230,6 +6235,73 @@ def load_expert_w2_weight_scale_mxfp4(self, module: torch.nn.Module, if not dst_on_gpu: dst_w2_weight_scale.copy_(dst_w2_weight_scale_gpu) + def load_packed_mxfp4_expert( + self, + module: torch.nn.Module, + *, + global_expert_id: int, + local_slot_id: int, + w1_weight: torch.Tensor, + w1_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + w3_weight: torch.Tensor, + w3_weight_scale: torch.Tensor, + ) -> None: + """Load one group-32 packed MXFP4 checkpoint expert into a local slot. + + This adapter is intentionally per-expert so model-specific streaming + loaders can keep safetensors mappings short-lived while reusing the + TRTLLM-Gen padding, sharding, shuffle, and scale-interleave lifecycle. + """ + if not 0 <= local_slot_id < module.expert_size_per_partition: + raise IndexError(f"local_slot_id={local_slot_id} is outside " + f"[0, {module.expert_size_per_partition}).") + expected_expert_id = module.initial_local_expert_ids[local_slot_id] + if global_expert_id != expected_expert_id: + raise ValueError( + f"local slot {local_slot_id} expects global expert " + f"{expected_expert_id}, got {global_expert_id}.") + + tensors = { + "w1_weight": w1_weight, + "w1_weight_scale": w1_weight_scale, + "w2_weight": w2_weight, + "w2_weight_scale": w2_weight_scale, + "w3_weight": w3_weight, + "w3_weight_scale": w3_weight_scale, + } + for name, value in tensors.items(): + if value.dtype != torch.uint8: + raise TypeError( + f"{name} must contain packed MXFP4 uint8 data, got " + f"{value.dtype}.") + + # Callers stream experts from a thread pool, so the duplicate check + # and slot claim must be one atomic step. The slot is claimed BEFORE + # the loaders run: a failed load leaves the destination buffer + # partially transformed, and a retry must not transform it again. + with _PACKED_MXFP4_SLOT_CLAIM_LOCK: + loaded_slots = getattr(module, "_packed_mxfp4_loaded_slots", None) + if loaded_slots is None: + loaded_slots = set() + module._packed_mxfp4_loaded_slots = loaded_slots + if local_slot_id in loaded_slots: + raise ValueError( + f"Packed MXFP4 local slot {local_slot_id} was loaded twice." + ) + loaded_slots.add(local_slot_id) + + self.load_expert_w3_w1_weight(module, w1_weight, w3_weight, + module.w3_w1_weight.data[local_slot_id]) + self.load_expert_w2_weight(module, w2_weight, + module.w2_weight.data[local_slot_id]) + self.load_expert_w3_w1_weight_scale_mxfp4( + module, w1_weight_scale, w3_weight_scale, + module.w3_w1_weight_scale.data[local_slot_id]) + self.load_expert_w2_weight_scale_mxfp4( + module, w2_weight_scale, module.w2_weight_scale.data[local_slot_id]) + class W4A16MXFP4TRTLLMGenFusedMoEMethod(MXFP4WeightTRTLLMGenFusedMoEMethod): pass diff --git a/tensorrt_llm/_torch/modules/kimi_k3_attn_res/__init__.py b/tensorrt_llm/_torch/modules/kimi_k3_attn_res/__init__.py new file mode 100644 index 000000000000..903729e57502 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_attn_res/__init__.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 Attention Residual fused decoder-layer op. + +Wraps the sm_100 ``attn_res_fwd`` (packed forward) kernel from +``exisiting_optimization_work/Attention_residual`` as an in-tree +decoder-layer fused op. K3 sets ``attn_res_block_size=12`` and uses this +op in two positions per decoder layer (before self-attention and after +the sub-block that produces the running ``prefix_sum``), plus once at +the top of the model. + +The kernel is Blackwell (sm_100a) only. On other devices the module +routes to a pure-torch reference implementation identical to +``exisiting_optimization_work/Attention_residual/tests/util/attn_res_ref.py`` +and to HF ``modeling_kimi._apply_attn_res``. +""" + +from .kimi_k3_attn_res import ( + KimiK3AttnResidualKernelPath, + KimiK3AttnResidualOp, + apply_attn_res_reference, + attn_res_fwd_chunked_reference, +) + +__all__ = [ + "KimiK3AttnResidualKernelPath", + "KimiK3AttnResidualOp", + "apply_attn_res_reference", + "attn_res_fwd_chunked_reference", +] diff --git a/tensorrt_llm/_torch/modules/kimi_k3_attn_res/_attn_res_kernels.py b/tensorrt_llm/_torch/modules/kimi_k3_attn_res/_attn_res_kernels.py new file mode 100644 index 000000000000..541f19ebec6d --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_attn_res/_attn_res_kernels.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernel dispatch for the in-tree Kimi K3 Attention Residual fused op. + +The optimized ``attn_res_fwd`` kernel (SM100 CuTe TMA warp-specialised +online-softmax + RMSNorm, Blackwell sm_100/sm_103 only) is source-integrated +into TensorRT-LLM as the ``trtllm::attn_res_fwd`` Torch op +(``cpp/tensorrt_llm/kernels/kimiK3AttnRes`` + ``cpp/tensorrt_llm/thop/attnResOp.cpp``). + +Dispatch: on sm_100/sm_103 with the compiled Torch bindings loaded the +module uses the fused op; otherwise it falls back to the pure-torch chunked +reference in :mod:`kimi_k3_attn_res.kimi_k3_attn_res`. +""" + +from __future__ import annotations + +import torch + +try: + from tensorrt_llm._utils import get_sm_version as _tllm_get_sm_version +except ImportError: # pragma: no cover — source-loader stub path + _tllm_get_sm_version = None + + +def _default_get_sm_version() -> int: + if not torch.cuda.is_available() or torch.cuda.device_count() == 0: + return -1 + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor + + +def get_attn_res_sm_version() -> int: + """Return the runtime SM version used for kernel selection. + + Prefers ``tensorrt_llm._utils.get_sm_version`` when the real package is + importable so environment-side overrides propagate. Falls back to a + plain CUDA-property probe otherwise. + """ + if _tllm_get_sm_version is not None: + try: + return int(_tllm_get_sm_version()) + except RuntimeError: + # torch raises RuntimeError when no CUDA device is usable; + # the property probe below handles that case itself. + return _default_get_sm_version() + return _default_get_sm_version() + + +def is_attn_res_optimized_supported() -> bool: + """The optimized ``attn_res_fwd`` kernel is Blackwell sm_100 only.""" + return get_attn_res_sm_version() in (100, 103) + + +def is_intree_attn_res_available() -> bool: + """True when the in-tree ``trtllm::attn_res_fwd`` Torch op is registered. + + The op is registered when TensorRT-LLM's compiled Torch bindings + (``libth_common``) are loaded, which happens on ``import tensorrt_llm``. + Under the source-loader stub subtree the bindings may be absent — then + this returns False and callers fall back to the reference (or the legacy + external-loader) path. + """ + try: + torch.ops.trtllm.attn_res_fwd # noqa: B018 — probe schema lookup + return True + except (AttributeError, RuntimeError): + return False + + +def intree_attn_res_fwd( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, +): + """Run the in-tree fused op. Returns ``(output, rsigma, probs, logits)``.""" + return torch.ops.trtllm.attn_res_fwd( + layer_residual, + block_residual, + res_weight.reshape(-1).contiguous(), + rms_weight.contiguous(), + float(rms_eps), + ) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py b/tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py new file mode 100644 index 000000000000..df90b4b41a40 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KimiK3AttnResidualOp — decoder-layer fused Attention Residual op. + +Kimi K3 uses ``attn_res_block_size=12``. Every layer emits a running +``prefix_sum``; every ``attn_res_block_size``-th layer accumulates that +prefix into a growing block-residual stack. Two positions per layer +consume the stack: once before self-attention (``self_attention_res``) +and once after the sub-block that produces the running ``prefix_sum`` +(``mlp_res``). The model also emits one final ``output_attn_res`` at the +top of the decoder tower. + +Each residual-selection step evaluates the same algebra, expressed by HF +``modeling_kimi._apply_attn_res``: + + v = concat(block_residual, prefix_sum.unsqueeze(1)) # [M, K+1, H] + variance = v.pow(2).mean(-1, keepdim=True) + k = v * rsqrt(variance + eps) + score_weight = norm.weight * proj.weight.squeeze(0) + scores = (k * score_weight).sum(-1) + probs = softmax(scores, dim=-1) # [M, K+1] + output = probs @ v # [M, H] + +The sm_100 ``attn_res_fwd`` fused kernel does exactly this in one launch +under the documented constraints: + + B == 1 + N = K + 1 in [1, 12] + T in [1, 16384] + H in {4096, 5120, 6144, 7168, 8192} (K3 uses H=7168) + layer_residual, block_residual: bf16 CUDA contiguous + res_weight: bf16 [H] or [H, 1] + rms_weight: bf16 [H] + +Interface differences vs HF: +* The kernel packs shapes as [T, B, H] and [K, T, B, H]; HF passes the + block residual as [M, K, H]. This module handles the reshape and + contiguity requirements before calling the kernel. +* The kernel returns ``rsigma`` and ``probs`` alongside ``output`` for + benchmarking / diagnostics; the module discards them on the module + path but exposes them on the standalone ``forward`` helper for tests. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch import nn + +from ._attn_res_kernels import ( + intree_attn_res_fwd, + is_attn_res_optimized_supported, + is_intree_attn_res_available, +) + + +class KimiK3AttnResidualKernelPath: + """Enum-like string tags for the selected kernel path.""" + + OPTIMIZED = "optimized" + REFERENCE = "reference" + + +# --------------------------------------------------------------------------- +# Reference implementations (chunked torch reference + HF direct). +# --------------------------------------------------------------------------- + + +def attn_res_fwd_chunked_reference( + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, + max_elements: int = 4 * 1024 * 1024, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Chunked torch reference for packed Attention Residual forward. + + Mirrors ``exisiting_optimization_work/Attention_residual/tests/util/attn_res_ref.py`` + so callers can produce byte-identical reference outputs without loading + the optimization tree. + + Args mirror :meth:`KimiK3AttnResidualOp.attn_res_fwd`: + + layer_residual: bf16 [T, B, H] + block_residual: bf16 [K, T, B, H] + res_weight: bf16 [H] or [H, 1] + rms_weight: bf16 [H] + + Returns ``(output, rsigma, probs, logits)`` with ``output`` in bf16 + ``[T, B, H]`` and the three saved tensors in fp32 ``[K+1, T, B]``. + """ + T, B, H = layer_residual.shape + K = int(block_residual.shape[0]) + N = K + 1 + M = T * B + layer_flat = layer_residual.reshape(M, H) + block_flat = block_residual.reshape(K, M, H) + res_w = res_weight.flatten().float() + rms_w = rms_weight.flatten().float() + + output = torch.empty_like(layer_residual) + rsigma = torch.empty((N, T, B), device=layer_residual.device, dtype=torch.float32) + probs = torch.empty_like(rsigma) + logits = torch.empty_like(rsigma) + + output_flat = output.reshape(M, H) + rsigma_flat = rsigma.reshape(N, M) + probs_flat = probs.reshape(N, M) + logits_flat = logits.reshape(N, M) + chunk_m = max(1, max_elements // max(N * H, 1)) + + for start in range(0, M, chunk_m): + end = min(start + chunk_m, M) + if K == 0: + values = layer_flat[start:end].unsqueeze(0).float() + else: + values = torch.cat( + [block_flat[:, start:end, :], layer_flat[start:end].unsqueeze(0)], + dim=0, + ).float() + rs = (values.square().mean(dim=-1) + rms_eps).rsqrt() + lg = (values * rs[..., None] * rms_w * res_w).sum(dim=-1) + pr = torch.softmax(lg, dim=0) + out = (pr[..., None] * values).sum(dim=0) + output_flat[start:end].copy_(out.to(layer_residual.dtype)) + rsigma_flat[:, start:end].copy_(rs) + probs_flat[:, start:end].copy_(pr) + logits_flat[:, start:end].copy_(lg) + + return output, rsigma, probs, logits + + +def apply_attn_res_reference( + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + proj_weight: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, +) -> torch.Tensor: + """HF ``_apply_attn_res`` mirror. + + ``prefix_sum`` is ``(num_tokens, hidden_size)`` and comes from the + running sub-block sum. + ``block_residual`` is ``(num_tokens, num_blocks, hidden_size)`` — the + HF layout, block-residual axis in the middle. + ``proj_weight`` is the ``proj.weight.squeeze(0)`` linear weight of + shape ``(hidden_size,)``. + ``rms_weight`` is the RMSNorm learnable weight of shape + ``(hidden_size,)``. + + Returns bf16 ``(num_tokens, hidden_size)`` matching the input dtype. + """ + v = torch.cat((block_residual, 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 + rms_eps) + score_weight = rms_weight.float() * proj_weight.squeeze().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) + + +# --------------------------------------------------------------------------- +# Fused-op module. +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _AttnResFwdInputs: + layer_residual: torch.Tensor # bf16 [T, B, H] + block_residual: torch.Tensor # bf16 [K, T, B, H] + res_weight: torch.Tensor # bf16 [H] or [H, 1] + rms_weight: torch.Tensor # bf16 [H] + + +class KimiK3AttnResidualOp(nn.Module): + """Decoder-layer Attention Residual fused op wrapper. + + The op owns the ``proj`` (``nn.Linear(hidden_size, 1, bias=False)``) + and ``norm`` (``KimiRMSNorm(hidden_size)`` — represented here by a raw + learnable ``weight`` and ``variance_epsilon``) parameters that + accompany each residual-selection site (self_attention_res, mlp_res, + output_attn_res). It exposes: + + * :meth:`forward_hf_layout` — the "HF-friendly" entry point that + accepts ``(prefix_sum: [M, H], block_residual: [M, K, H])`` and + returns ``[M, H]``. This is the shape the K3 decoder layer wants. + * :meth:`forward` — the raw kernel entry point mirroring + ``attn_res_fwd``: ``(layer_residual, block_residual_kthbh, ...)`` → + ``(output, rsigma, probs)``. + + Kernel path selection follows the same policy as the KDA module: + Blackwell sm_100 + a resolvable optimization root ⇒ ``OPTIMIZED``; + otherwise ``REFERENCE`` (pure-torch chunked reference). + """ + + def __init__( + self, + hidden_size: int, + rms_eps: float = 1e-6, + force_use_fallback_kernel: bool = False, + **kwargs, + ) -> None: + super().__init__() + self.hidden_size = int(hidden_size) + self.variance_epsilon = float(rms_eps) + self.force_use_fallback_kernel = bool(force_use_fallback_kernel) + + # Match HF: KimiRMSNorm.weight has shape [H]; the projection is + # ``nn.Linear(H, 1, bias=False)``. Storing as raw Parameters lets + # the module test copy weights directly from the HF module. + self.rms_weight = nn.Parameter(torch.ones(self.hidden_size)) + self.proj_weight = nn.Parameter(torch.zeros(1, self.hidden_size)) + + if ( + force_use_fallback_kernel + or not is_attn_res_optimized_supported() + or not is_intree_attn_res_available() + ): + self.kernel_path = KimiK3AttnResidualKernelPath.REFERENCE + else: + self.kernel_path = KimiK3AttnResidualKernelPath.OPTIMIZED + self._optimized_extension = None + + # ------------------------------------------------------------------ + # Weight loading helpers. + # ------------------------------------------------------------------ + + def copy_weights_from( + self, + hf_norm: nn.Module, + hf_proj: nn.Module, + ) -> None: + """Copy weights from an HF ``(KimiRMSNorm, nn.Linear(H,1))`` pair. + + HF wires the residual-selection site with a ``KimiRMSNorm`` and an + ``nn.Linear(H, 1, bias=False)`` — ``self_attention_res_{norm,proj}``, + ``mlp_res_{norm,proj}``, ``output_attn_res_{norm,proj}``. + """ + with torch.no_grad(): + self.rms_weight.data.copy_( + hf_norm.weight.detach().to( + dtype=self.rms_weight.dtype, device=self.rms_weight.device + ) + ) + # HF proj.weight has shape [1, H]; we store [1, H] so both + # layouts stay compatible without a squeeze at call time. + self.proj_weight.data.copy_( + hf_proj.weight.detach().to( + dtype=self.proj_weight.dtype, device=self.proj_weight.device + ) + ) + + # ------------------------------------------------------------------ + # Raw kernel entry. + # ------------------------------------------------------------------ + + def _validate_inputs(self, inputs: _AttnResFwdInputs) -> Tuple[int, int, int, int]: + lr = inputs.layer_residual + br = inputs.block_residual + rw = inputs.res_weight + nw = inputs.rms_weight + if lr.ndim != 3: + raise ValueError("layer_residual must have shape [T, B, H]") + T, B, H = lr.shape + if B != 1: + raise ValueError(f"attn_res_fwd requires B==1 (got B={B})") + if H != self.hidden_size: + raise ValueError( + f"layer_residual H={H} does not match module hidden_size={self.hidden_size}" + ) + if lr.dtype != torch.bfloat16 or not lr.is_cuda or not lr.is_contiguous(): + raise ValueError("layer_residual must be a CUDA bf16 contiguous tensor") + if br.ndim != 4: + raise ValueError("block_residual must have shape [K, T, B, H]") + K = int(br.shape[0]) + if tuple(br.shape[1:]) != (T, B, H): + raise ValueError( + f"block_residual shape {tuple(br.shape)} does not match [K, T, B, H] for T={T}, B={B}, H={H}" + ) + if br.dtype != torch.bfloat16 or not br.is_cuda or not br.is_contiguous(): + raise ValueError("block_residual must be a CUDA bf16 contiguous tensor") + if K + 1 > 12: + raise ValueError(f"attn_res_fwd requires N=K+1 in [1, 12] (got N={K + 1})") + if T > 16384: + raise ValueError(f"attn_res_fwd requires T in [1, 16384] (got T={T})") + if rw.numel() != H or rw.dtype != torch.bfloat16 or not rw.is_cuda: + raise ValueError("res_weight must be CUDA bf16 with H elements") + if nw.numel() != H or nw.dtype != torch.bfloat16 or not nw.is_cuda: + raise ValueError("rms_weight must be CUDA bf16 with H elements") + return T, B, H, K + + def _attn_res_fwd_optimized( + self, inputs: _AttnResFwdInputs + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self._validate_inputs(inputs) + output, rsigma, probs, _logits = intree_attn_res_fwd( + inputs.layer_residual, + inputs.block_residual, + inputs.res_weight, + inputs.rms_weight, + self.variance_epsilon, + ) + return output, rsigma, probs + + def _attn_res_fwd_reference( + self, inputs: _AttnResFwdInputs + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self._validate_inputs(inputs) + out, rsigma, probs, _logits = attn_res_fwd_chunked_reference( + inputs.layer_residual, + inputs.block_residual, + inputs.res_weight, + inputs.rms_weight, + self.variance_epsilon, + ) + return out, rsigma, probs + + @torch.no_grad() + def forward( + self, + layer_residual: torch.Tensor, + block_residual: torch.Tensor, + res_weight: Optional[torch.Tensor] = None, + rms_weight: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Kernel-shaped forward. + + Uses this module's own ``proj_weight`` / ``rms_weight`` as + weights unless overrides are supplied. Returns + ``(output, rsigma, probs)`` matching the fused kernel's contract. + """ + rw = res_weight if res_weight is not None else self.proj_weight + nw = rms_weight if rms_weight is not None else self.rms_weight + inputs = _AttnResFwdInputs( + layer_residual=layer_residual, + block_residual=block_residual, + res_weight=rw, + rms_weight=nw, + ) + if self.kernel_path == KimiK3AttnResidualKernelPath.OPTIMIZED: + return self._attn_res_fwd_optimized(inputs) + return self._attn_res_fwd_reference(inputs) + + # ------------------------------------------------------------------ + # HF-friendly entry point. + # ------------------------------------------------------------------ + + @torch.no_grad() + def forward_hf_layout( + self, + prefix_sum: torch.Tensor, + block_residual: torch.Tensor, + ) -> torch.Tensor: + """Apply the fused residual-selection op in the HF shape layout. + + ``prefix_sum`` is ``(num_tokens, hidden_size)`` bf16 CUDA. + ``block_residual`` is ``(num_tokens, num_blocks, hidden_size)`` bf16 + CUDA (HF layout with the block axis in the + middle). + + Returns ``(num_tokens, hidden_size)`` bf16. + + Internally reshapes to the kernel's ``[T, B=1, H]`` and + ``[K, T, B=1, H]`` layout, calls the fused op, and reshapes back. + """ + if prefix_sum.ndim != 2: + raise ValueError("prefix_sum must have shape [M, H]") + M, H = prefix_sum.shape + if H != self.hidden_size: + raise ValueError( + f"prefix_sum H={H} does not match module hidden_size={self.hidden_size}" + ) + if block_residual.ndim != 3 or block_residual.shape[0] != M or block_residual.shape[2] != H: + raise ValueError("block_residual must have shape [M, K, H] matching prefix_sum") + K = int(block_residual.shape[1]) + + # Kernel expects [T, B=1, H] and [K, T, B=1, H] bf16 contiguous. + # HF passes num_tokens as the flat batch axis, so we set T=M, B=1. + layer_kernel = prefix_sum.reshape(M, 1, H).contiguous() + # HF block_residual [M, K, H] → kernel [K, M, 1, H]. + block_kernel = block_residual.transpose(0, 1).reshape(K, M, 1, H).contiguous() + + output, _rsigma, _probs = self.forward(layer_kernel, block_kernel) + return output.reshape(M, H) + + # ------------------------------------------------------------------ + # Diagnostics. + # ------------------------------------------------------------------ + + def kernel_source(self) -> str: + """Return a stable string describing the kernel path in use.""" + if self.kernel_path == KimiK3AttnResidualKernelPath.OPTIMIZED: + return "" + return "" + + def precompile(self, verbose: bool = False) -> None: + """No-op: the in-tree ``trtllm::attn_res_fwd`` op is pre-compiled.""" + del verbose diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py new file mode 100644 index 000000000000..9f2bfd441219 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MLA in-tree module for TensorRT-LLM's PyTorch backend. + +K3 MLA is DeepSeek-V3-style multi-latent attention with three K3-specific +deltas that live at the module level (not the attention backend): + +* **NoPE.** ``mla_use_nope=True`` in K3 config disables the rotary + embedding; both the query and key rope slots pass through the backend + unchanged. +* **Output gate before ``o_proj``.** When ``mla_use_output_gate=True`` an + extra ``g_proj`` computes ``sigmoid(g_proj(hidden_states)) * attn_output`` + before the final projection. +* **Softmax scale.** ``(qk_nope + qk_rope) ** -0.5 = 192 ** -0.5`` for + real K3 dims — matches ``TrtllmAttention`` default MLA q_scaling. + +The module wraps the existing ``TrtllmAttention`` backend MLA path plus +``KVCacheManagerV2`` for both context and cached-decode. +""" + +from .kimi_k3_mla_attention import KimiK3MLAAttention + +__all__ = [ + "KimiK3MLAAttention", +] diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py new file mode 100644 index 000000000000..c374bad6dcdd --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 specialization of the shared PyTorch-backend MLA module. + +The base ``MLA`` class owns context attention, cached/chunked prefill, +absorbed generation, and paged-cache handling. This module only supplies the +K3 projection topology, NoPE identity table, KV-B checkpoint layout, and +gated output projection. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch import nn + +from ....functional import PositionEmbeddingType +from ....models.modeling_utils import QuantConfig +from ...attention_backend import AttentionMetadata, TrtllmAttention +from ...attention_backend.interface import PositionalEmbeddingParams, RopeParams +from ...model_config import ModelConfig +from ..mla import MLA + + +def _meta_safe_cast_dtype(module, dtype): + """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. + + ``Module.to`` dispatches ``aten._to_copy``, which MetaInitMode rejects + (it would silently fall back to full CPU construction of the model — + ~70 GB of host RAM per rank for Kimi K3). Under meta init the values + are garbage anyway, so a dtype-only re-allocation via ``empty_like`` + (an allowed init op) is equivalent; off meta this matches ``.to``. + """ + import torch as _torch + + def _cast(t): + if not t.is_floating_point(): + return t + if t.is_meta: + return _torch.empty_like(t, dtype=dtype) + return t.to(dtype=dtype) + + module._apply(_cast) + + +def _make_pos_embd_params( + *, + qk_rope_head_dim: int, + max_position_embeddings: int, +) -> PositionalEmbeddingParams: + """Build a valid rope config so the backend allocates a real cache. + + We use rope_gpt_neox with default theta=10000 and ``duplicate_data + =True`` (the same convention DeepSeek-V3-style MLA uses when + ``qk_rope_head_dim`` is present). The resulting ``rotary_cos_sin`` + has the exact shape the C++ MLA rope kernel indexes. Immediately + after backend construction we overwrite the tensor values with + ``(cos=1, sin=0)`` — an identity rotation, matching K3's NoPE. + """ + rope_params = RopeParams( + dim=qk_rope_head_dim, + theta=10000.0, + max_positions=max_position_embeddings, + original_max_positions=max_position_embeddings, + duplicate_data=True, + ) + return PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + # Match the working DeepSeek-V3-style MLA reference test + # (tests/unittest/_torch/attention/test_attention_mla.py) which + # sets ``is_neox=False``. The MLA fused rope kernel is GPT-J + # style regardless of this flag, but the C++ FMHA reads this bit + # elsewhere and stability under identity-cos-sin depends on the + # standard non-neox layout. + is_neox=False, + ) + + +def _write_identity_rope_values(cos_sin: torch.Tensor) -> None: + """Overwrite a rotary cos/sin table with identity values in place. + + Interleaved (cos, sin) pairs: index [::2] = cos, [1::2] = sin. + Setting cos=1 and sin=0 per position makes the rotation the + identity — a mathematical no-op — which preserves K3's NoPE + semantics without patching the backend. + """ + flat = cos_sin.reshape(-1) + with torch.no_grad(): + flat[0::2] = 1.0 + flat[1::2] = 0.0 + # Ensure the identity write reaches CUDA memory before any kernel + # launched from a different stream can read the table. + if cos_sin.is_cuda: + torch.cuda.synchronize(cos_sin.device) + + +def _install_identity_rope_table(backend: TrtllmAttention) -> None: + """Install an identity rotary cos/sin table on ``backend``. + + The C++ MLA rope kernels (``mla_rope_generation`` and the context + preprocess) read this table and apply the rotation; identity values + make that a copy, preserving K3's NoPE. + + The tensor SHAPE produced by ``create_rope_const_params`` is kept + intact so the C++ ``float2`` indexing stays valid. Only the values + are overwritten in place. ``_ensure_rope_table_size`` is replaced + with an identity-preserving resize: the table may GROW (so the + fused rope-generation op can never index out of bounds for long + sequences) but its values are always rewritten to identity right + after a regeneration, so the real sinusoids never leak in. + """ + cos_sin = backend.rotary_cos_sin + if cos_sin is None: + raise RuntimeError( + "backend.rotary_cos_sin is None after construction; check " + "pos_embd_params has a valid RopeParams with dim > 0." + ) + _write_identity_rope_values(cos_sin) + + orig_resize = backend._ensure_rope_table_size # bound method + + def _identity_preserving_resize(required_max_positions: int) -> None: + if required_max_positions <= backend.rope_params.max_positions: + return + orig_resize(required_max_positions) + _write_identity_rope_values(backend.rotary_cos_sin) + + backend._ensure_rope_table_size = _identity_preserving_resize + + +# --------------------------------------------------------------------------- +# KimiK3MLAAttention. +# --------------------------------------------------------------------------- + + +class KimiK3MLAAttention(MLA): + """Kimi K3 MLA implemented as a thin specialization of :class:`MLA`. + + K3 keeps the standard dense MLA attention/cache flow and only changes the + checkpoint projection topology, positional encoding, KV-B runtime layout, + and gated output projection. + """ + + def __init__( + self, + *, + hidden_size: int, + num_heads: int, + q_lora_rank: int, + kv_lora_rank: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + rms_norm_eps: Optional[float] = None, + dtype: Optional[torch.dtype] = None, + layer_idx: int = 0, + use_output_gate: bool = True, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantConfig] = None, + ) -> None: + pos_embd_params = _make_pos_embd_params( + qk_rope_head_dim=qk_rope_head_dim, + max_position_embeddings=max_position_embeddings, + ) + model_config = ModelConfig( + quant_config=quant_config if quant_config is not None else QuantConfig() + ) + super().__init__( + hidden_size=hidden_size, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + predicted_tokens_per_seq=1, + max_position_embeddings=max_position_embeddings, + bias=False, + pos_embd_params=pos_embd_params, + layer_idx=layer_idx, + dtype=dtype, + dense_bias=False, + config=model_config, + reduce_output=False, + fuse_qkv_a_proj=False, + rms_norm_eps=rms_norm_eps, + ) + # K3 calls forward_impl() directly to insert its output gate before + # the base row-parallel o_proj. The original executor metadata remains + # intact, so MLA performs its native mixed context/generation split. + self.register_to_config = False + + self.use_output_gate = use_output_gate + + if use_output_gate: + self.g_proj = nn.Linear( + hidden_size, + num_heads * v_head_dim, + bias=False, + ) + + # K3 is NoPE. The base MLA backends still require real RoPE tables, so + # retain their expected shape and replace every rotation with identity. + assert isinstance(self.mha, TrtllmAttention) + assert isinstance(self.mqa, TrtllmAttention) + _install_identity_rope_table(self.mha) + _install_identity_rope_table(self.mqa) + self.rotary_emb = None + self.apply_rotary_emb = False + + if dtype is not None: + _meta_safe_cast_dtype(self, dtype) + + def _apply_output_gate_and_o_proj( + self, + hidden_states: torch.Tensor, + attn_out: torch.Tensor, + ) -> torch.Tensor: + if self.use_output_gate: + attn_out = attn_out * self.g_proj(hidden_states).sigmoid() + return self.o_proj(attn_out) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + attn_out = self.create_output( + hidden_states, + attn_metadata.num_contexts, + ) + super().forward_impl( + None, + hidden_states, + attn_metadata, + output=attn_out, + ) + return self._apply_output_gate_and_o_proj(hidden_states, attn_out) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py new file mode 100644 index 000000000000..151a5c97f561 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 sparse MoE in-tree module. + +Ships the promoted ``KimiK3SparseMoeBlock`` and its supporting pieces +(``KimiK3MoEGate``, latent projections, shared experts, MXFP4-packed +routed expert bank, native TRTLLM-Gen SiTU dispatch). Structurally +mirrors HF ``KimiSparseMoeBlock`` at ``modeling_kimi.py:806-918``. + +Two kernel paths coexist under one module class: + +* ``use_fused_cubin=False`` — Python fallback with MXFP4 group-32 + routed expert weights, dequantized to canonical fp32 on demand. + Byte-exact HF parity under random weights. +* ``use_fused_cubin=True`` — native in-tree + ``torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner`` invocation + (``act_type=SiTu``) on checkpoint-derived MXFP4 weights shared with + the fallback bank. Routing goes through the op's + ``topk_weights``/``topk_ids`` bypass fed by the real K3 gate. +""" + +from .kimi_k3_moe_block import ( + KimiK3RoutedExpertBank, + KimiK3SparseMoeBlock, + MoEBlockProvenance, + copy_hf_moe_block_weights, +) +from .kimi_k3_moe_gate import KimiK3MoEGate, copy_hf_moe_gate_weights + +__all__ = [ + "KimiK3MoEGate", + "KimiK3RoutedExpertBank", + "KimiK3SparseMoeBlock", + "MoEBlockProvenance", + "copy_hf_moe_gate_weights", + "copy_hf_moe_block_weights", +] diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py new file mode 100644 index 000000000000..b0b0d7563c98 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dense MLP helper for the in-tree Kimi K3 MoE block. + +Kimi K3 uses the ``situ`` activation (not SiLU/SwiGLU). ``SituAndMul`` +computes ``beta * tanh(gate / beta) * sigmoid(gate)`` on the gate half +and optionally applies ``linear_beta * tanh(up / linear_beta)`` on the +up half, then multiplies. ``KimiK3MLP`` is the fused ``gate_up_proj + +down_proj`` layout used by the shared expert stack in HF +``KimiSparseMoeBlock`` — the same shape TRT-LLM's ``GatedMLP`` uses. + +Two activation paths coexist: + +* the eager fp32 ``SituAndMul`` module — the byte-exact HF reference, + used by the parity-test MoE block and as the fallback; +* the fused Triton ``trtllm::situ_and_mul`` custom op (same fp32 math + in a single kernel, modeled on ``modules/swiglu.py``'s + ``silu_and_mul_kernel``), enabled per ``KimiK3MLP`` instance via + ``use_fused_activation=True`` (the runtime model opts in). The op is + CUDA-graph-safe: no host synchronization and no data-dependent + control flow. +""" + +from __future__ import annotations + +import os +from typing import Mapping, Optional + +import torch +import triton # type: ignore[import] +import triton.language as tl # type: ignore[import] +import triton.language.extra.libdevice as tldevice # type: ignore[import] +from torch import nn + +from ...flashinfer_utils import IS_FLASHINFER_AVAILABLE + +# Route the RMSNorm forward through flashinfer's single-kernel fused RMSNorm +# instead of the eager pow/mean/rsqrt/mul/cast chain. Set to "0" to fall back +# to the eager reference (the exact-parity rollback lever). +_FUSED_RMSNORM = os.environ.get("KIMI_K3_FUSED_RMSNORM", "1") == "1" + + +class SituAndMul(nn.Module): + """K3 SiTU activation with gate/up multiplicative gating. + + Byte-identical to HF ``modeling_kimi.py``'s ``SituAndMul`` at + lines 41-59. Runs the math in fp32 for numerical stability + (matches HF), then casts back to the input's dtype. + """ + + def __init__( + self, + *, + beta: float = 1.0, + linear_beta: Optional[float] = None, + ) -> None: + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].to(torch.float32) + up = x[..., d:].to(torch.float32) + situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (situ_a * up).to(x.dtype) + + +@triton.jit +def situ_and_mul_kernel( + o_ptr, + o_stride, + x_ptr, + x_stride, + d, + beta, + linear_beta, + BLOCK_SIZE: tl.constexpr, + HAS_LINEAR_BETA: tl.constexpr, +) -> None: + """Fused :class:`SituAndMul` on a packed ``[gate | up]`` row layout. + + Loads ``gate = x[i, :d]`` and ``up = x[i, d:2d]``, computes (fp32) + ``beta * tanh(gate / beta) * sigmoid(gate) * up'`` with + ``up' = linear_beta * tanh(up / linear_beta)`` when + ``HAS_LINEAR_BETA`` else ``up``, and stores the product rounded to + ``o_ptr``'s element type. + """ + i = tl.program_id(axis=0).to(tl.int64) + j = tl.program_id(axis=1) + + o_row_ptr = o_ptr + o_stride * i + x_row_ptr = x_ptr + x_stride * i + + offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < d + + gate = tl.load(x_row_ptr + offsets, mask=mask).to(tl.float32) + up = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32) + + situ_a = beta * tldevice.tanh(gate / beta) * tl.sigmoid(gate) + if HAS_LINEAR_BETA: + up = linear_beta * tldevice.tanh(up / linear_beta) + result = situ_a * up + + tl.store(o_row_ptr + offsets, result, mask=mask) + + +@torch.library.custom_op("trtllm::situ_and_mul", mutates_args=()) +def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor: + """Fused SiTU activation (single Triton kernel, fp32 internal math). + + Args: + x: ``[num_tokens, 2 * d]`` packed ``[gate | up]`` GEMM output + (fp16/bf16/fp32; the last dim must be contiguous). + beta: SiTU gate ``beta`` (``activation_situ_beta``). + linear_beta: optional up-half ``linear_beta`` + (``activation_situ_linear_beta``); ``None`` keeps the up half + linear. + + Returns: + ``[num_tokens, d]`` tensor in ``x``'s dtype, numerically matching + the eager :class:`SituAndMul` reference. + """ + b, n = x.shape + + assert n % 2 == 0 + d = n // 2 + + o = torch.empty((b, d), dtype=x.dtype, device=x.device) + + def grid(meta: Mapping[str, int]) -> tuple[int, int]: + return (b, triton.cdiv(d, meta["BLOCK_SIZE"])) + + situ_and_mul_kernel[grid]( + o_ptr=o, + o_stride=o.stride(0), + x_ptr=x, + x_stride=x.stride(0), + d=d, + beta=float(beta), + linear_beta=float(linear_beta) if linear_beta is not None else 1.0, + BLOCK_SIZE=1024, + HAS_LINEAR_BETA=linear_beta is not None, + ) + + return o + + +@situ_and_mul.register_fake +def _(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor: + b, n = x.shape + + assert n % 2 == 0 + + return x.new_empty((b, n // 2)) + + +class NonSituActivation(nn.Module): + """SiLU/SwiGLU activation used as the non-SiTU mutation control. + + Splits the last dim into gate/up, applies SiLU to the gate, and + multiplies element-wise. Deliberately does NOT use the SiTU + ``beta * tanh(gate/beta) * sigmoid(gate)`` recipe. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d] + up = x[..., d:] + return torch.nn.functional.silu(gate) * up + + +class KimiK3MLP(nn.Module): + """K3 dense/shared-expert MLP module with TRT-LLM-style fused layout. + + Weight layout: + + * ``gate_up_proj``: ``nn.Linear(hidden_size, 2 * intermediate_size, bias=False)``. + Rows ``[:intermediate_size]`` correspond to HF's ``gate`` (KimiMLP.gate_proj + or KimiBlockSparseMLP.w1). Rows ``[intermediate_size:]`` correspond to + HF's ``up`` (KimiMLP.up_proj or KimiBlockSparseMLP.w3). + * ``down_proj``: ``nn.Linear(intermediate_size, hidden_size, bias=False)``. + Matches HF ``KimiMLP.down_proj`` or ``KimiBlockSparseMLP.w2``. + + Forward: ``down_proj( activation( gate_up_proj(x) ) )``. Default + ``activation`` is :class:`SituAndMul`; pass a different callable to + run mutation controls (e.g. :class:`NonSituActivation` for a + negative-control test). ``use_fused_activation=True`` routes CUDA + inputs through the fused Triton ``trtllm::situ_and_mul`` op instead + of the eager module (only valid with the default SiTU activation). + """ + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + situ_beta: float = 4.0, + situ_linear_beta: Optional[float] = 25.0, + activation: Optional[nn.Module] = None, + use_fused_activation: bool = False, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + if use_fused_activation and activation is not None: + raise ValueError( + "use_fused_activation only fuses the default SiTU activation; " + "drop the custom activation module or the flag" + ) + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.use_fused_activation = use_fused_activation + + self.gate_up_proj = nn.Linear( + hidden_size, + 2 * intermediate_size, + bias=False, + dtype=dtype, + device=device, + ) + self.down_proj = nn.Linear( + intermediate_size, + hidden_size, + bias=False, + dtype=dtype, + device=device, + ) + self.activation = ( + activation + if activation is not None + else SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h1 = self.gate_up_proj(x) + if self.use_fused_activation and h1.is_cuda: + act = self.activation + h2 = torch.ops.trtllm.situ_and_mul( + h1.reshape(-1, h1.shape[-1]), act.beta, act.linear_beta + ).reshape(*h1.shape[:-1], self.intermediate_size) + else: + h2 = self.activation(h1) + return self.down_proj(h2) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching HF ``KimiRMSNorm`` semantics exactly. + + HF ``KimiRMSNorm.forward``:: + + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + eps) + return self.weight * hidden_states.to(input_dtype) + + ``self.weight`` in HF is initialised in the module's ambient dtype + (bf16 or fp32). Callers pin the weight dtype here too so byte-exact + parity holds regardless of the ambient dtype. + """ + + 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: + # flashinfer's fused RMSNorm does the same fp32-accumulate + # normalization in one kernel, collapsing the eager + # pow/mean/rsqrt/mul/cast launch chain. Rounding differs by one final + # cast: flashinfer multiplies by ``weight`` in fp32 and casts once at + # the end, while the eager path casts the normalized value to the + # input dtype BEFORE the weight multiply — so outputs can differ by + # ~1 ulp and byte-exact HF parity requires the eager path. It is only + # valid for a CUDA fp16/bf16 input whose dtype matches the weight; + # CPU / fp32 parity paths, meta init, and the KIMI_K3_FUSED_RMSNORM=0 + # rollback keep the exact eager math below. + if ( + _FUSED_RMSNORM + and IS_FLASHINFER_AVAILABLE + and hidden_states.is_cuda + and hidden_states.dtype in (torch.float16, torch.bfloat16) + and self.weight.dtype == hidden_states.dtype + ): + from ...custom_ops import flashinfer_rmsnorm + + return flashinfer_rmsnorm(hidden_states.contiguous(), self.weight, self.eps) + input_dtype = hidden_states.dtype + h = hidden_states.to(torch.float32) + variance = h.pow(2).mean(-1, keepdim=True) + h = h * torch.rsqrt(variance + self.eps) + return self.weight * h.to(input_dtype) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py new file mode 100644 index 000000000000..e6873815f06b --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Native TRTLLM-Gen SiTU MoE dispatch for the in-tree Kimi K3 sparse MoE block. + +The K3 MoE block has two mutually exclusive kernel paths: + +1. **Python fallback** — MXFP4 group-32 routed expert weights are + dequantized on the fly, then fed through per-expert + ``gate_up_proj + activation + down_proj`` linears. Byte-exact HF + parity under random weights. + +2. **Native fused SiTU path** — routed compute goes through the in-tree + ``torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner`` custom op + with ``act_type=ActType_TrtllmGen.SiTu``. Activations are dynamically + quantized to MXFP8 (``mxfp8_quantize``); weights are the checkpoint's + MXFP4 group-32 tensors padded and shuffled with the same + pad/shuffle/interleave contract as + ``MXFP4WeightTRTLLMGenFusedMoEMethod`` (see + ``fused_moe/quantization.py``). No FlashInfer private-cubin + environment variable is involved. + +Routing always goes through the ``topk_weights``/``topk_ids`` bypass: +K3's gate semantics (sigmoid scoring, ``e_score_correction_bias`` +affecting selection only, renormalize, no groups) match none of the +built-in trtllm-gen routing methods, so the K3 gate computes top-k on +the host module and the op consumes the precomputed result verbatim. + +Weight/activation packing conventions (must stay in sync with the +generated ``GemmGatedActOptions.h`` SiTuGlu definition): + +* FC1 packs ``w3`` (up/linear) in the first half and ``w1`` (gate) in + the second half — the cubin evaluates ``x0`` (linear) from the first + half and ``x1`` (gate) from the second. This is the **opposite** of + the HF/Python reference layout where gate comes first. +* ``gemm1_alpha`` (cubin ``alpha``, gate side ``x1``) <- ``activation_situ_beta`` +* ``gemm1_beta`` (cubin ``beta``, linear side ``x0``) <- ``activation_situ_linear_beta`` +* ``gemm1_clamp_limit=None`` means +inf (no clamping) per the generated + ``KernelParamsDecl.h`` contract. +""" + +from __future__ import annotations + +from typing import Dict, Optional, Tuple + +import torch + +# TRTLLM-Gen backend alignment contract — keep in sync with +# MXFP4WeightTRTLLMGenFusedMoEMethod in fused_moe/quantization.py and the +# roundUp calls in blockScaleMoe/runner.cu. +INPUT_HIDDEN_ALIGNMENT = 512 +WEIGHT_ALIGNMENT = 128 +SCALING_VECTOR_SIZE = 32 +EPILOGUE_TILE_M = 128 + +# Local memo for the (expensive) shuffle permute-index computation, shared +# across experts/instances with identical shapes. +_CACHE_PERMUTE_INDICES: Dict[tuple, torch.Tensor] = {} + + +def _round_up(x: int, alignment: int) -> int: + return (x + alignment - 1) // alignment * alignment + + +def get_moe_sm_version() -> int: + """Return the runtime SM version used for kernel-support checks.""" + if not torch.cuda.is_available() or torch.cuda.device_count() == 0: + return -1 + try: + from tensorrt_llm._utils import get_sm_version as _tllm_get_sm_version + except ImportError: # pragma: no cover — source-loader stub path + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor + return int(_tllm_get_sm_version()) + + +def is_native_situ_supported() -> bool: + """Native SiTU cubins are Blackwell sm_100f (SM100/SM103) only.""" + return get_moe_sm_version() in (100, 103) + + +def assert_native_situ_supported( + *, + hidden_size: int, + intermediate_size: int, + group_size: int = SCALING_VECTOR_SIZE, +) -> None: + """Fail loudly (before any launch) when the fused SiTU path cannot run.""" + if not torch.cuda.is_available(): + raise RuntimeError("native SiTU MoE requires CUDA; no CUDA device is available") + sm = get_moe_sm_version() + if sm not in (100, 103): + raise RuntimeError(f"native SiTU MoE requires SM100/SM103 (Blackwell); running on SM{sm}") + if group_size != SCALING_VECTOR_SIZE: + raise RuntimeError( + f"native SiTU MoE requires MXFP4 group_size {SCALING_VECTOR_SIZE}, got {group_size}" + ) + if hidden_size % group_size != 0 or intermediate_size % group_size != 0: + raise RuntimeError( + f"hidden_size {hidden_size} and intermediate_size {intermediate_size} must be " + f"multiples of the MXFP4 group size {group_size}" + ) + + +def padded_fused_shapes(hidden_size: int, intermediate_size: int) -> Tuple[int, int, int]: + """Return (hidden_padded_fc1, hidden_padded_fc2, intermediate_padded). + + FC1 consumes activations along hidden (K dim, 512-aligned); FC2 produces + hidden (N dim, 128-aligned); intermediate is 128-aligned on both sides. + """ + return ( + _round_up(hidden_size, INPUT_HIDDEN_ALIGNMENT), + _round_up(hidden_size, WEIGHT_ALIGNMENT), + _round_up(intermediate_size, WEIGHT_ALIGNMENT), + ) + + +def pack_routed_expert_weights( + *, + w1_packed: torch.Tensor, + w1_scales: torch.Tensor, + w3_packed: torch.Tensor, + w3_scales: torch.Tensor, + w2_packed: torch.Tensor, + w2_scales: torch.Tensor, + device: torch.device, +) -> Dict[str, torch.Tensor]: + """Pad + shuffle checkpoint MXFP4 expert weights into the TRTLLM-Gen layout. + + Inputs are the per-expert MXFP4 tensors as stored by + :class:`KimiK3RoutedExpertBank` (HF layout, group_size=32): + + * ``w1_packed``/``w3_packed``: ``uint8 [E, I, H // 2]`` (w1 = gate, w3 = up) + * ``w1_scales``/``w3_scales``: ``uint8 [E, I, H // 32]`` (E8M0 biased exponents) + * ``w2_packed``: ``uint8 [E, H, I // 2]``, ``w2_scales``: ``uint8 [E, H, I // 32]`` + + Returns CUDA buffers matching ``MXFP4WeightTRTLLMGenFusedMoEMethod``'s + device layout: + + * ``gemm1_weights``: ``uint8 [E, 2 * I_pad, H_pad512 // 2]`` — w3 first, + w1 second, then row-shuffled for the gated-act GEMM. + * ``gemm1_weights_scale``: ``uint8 [E, 2 * I_pad, H_pad512 // 32]`` — + shuffled + block-scale interleaved. + * ``gemm2_weights``: ``uint8 [E, H_pad128, I_pad // 2]`` — row-shuffled. + * ``gemm2_weights_scale``: ``uint8 [E, H_pad128, I_pad // 32]`` — + shuffled + block-scale interleaved. + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + maybe_pad_for_mxfp4, + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, + ) + from tensorrt_llm.quantization.utils.fp4_utils import float4_sf_dtype + + for name, t in ( + ("w1_packed", w1_packed), + ("w1_scales", w1_scales), + ("w3_packed", w3_packed), + ("w3_scales", w3_scales), + ("w2_packed", w2_packed), + ("w2_scales", w2_scales), + ): + if t.dtype != torch.uint8: + raise RuntimeError(f"{name} must be uint8 MXFP4 data, got {t.dtype}") + + num_experts, intermediate_size, hidden_half = w1_packed.shape + hidden_size = hidden_half * 2 + h_pad_fc1, h_pad_fc2, i_pad = padded_fused_shapes(hidden_size, intermediate_size) + + gemm1_weights = torch.zeros( + (num_experts, 2 * i_pad, h_pad_fc1 // 2), dtype=torch.uint8, device=device + ) + gemm1_weights_scale = torch.zeros( + (num_experts, 2 * i_pad, h_pad_fc1 // SCALING_VECTOR_SIZE), + dtype=torch.uint8, + device=device, + ) + gemm2_weights = torch.zeros( + (num_experts, h_pad_fc2, i_pad // 2), dtype=torch.uint8, device=device + ) + gemm2_weights_scale = torch.zeros( + (num_experts, h_pad_fc2, i_pad // SCALING_VECTOR_SIZE), + dtype=torch.uint8, + device=device, + ) + + for e in range(num_experts): + # ---- FC1 weights: pad, place w3 (linear) first / w1 (gate) second, shuffle. + dst = gemm1_weights[e] + dst_w3, dst_w1 = dst.chunk(2, dim=0) + dst_w3.copy_(maybe_pad_for_mxfp4(w3_packed[e].to(device), h_pad_fc1 // 2, i_pad)) + dst_w1.copy_(maybe_pad_for_mxfp4(w1_packed[e].to(device), h_pad_fc1 // 2, i_pad)) + permute = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + dst, _CACHE_PERMUTE_INDICES, EPILOGUE_TILE_M + ) + dst.copy_(torch.ops.trtllm.shuffle_matrix(dst, permute.to(device))) + + # ---- FC1 scales: pad, place, shuffle with sf indices, interleave. + dst_sf = gemm1_weights_scale[e] + dst_sf_w3, dst_sf_w1 = dst_sf.chunk(2, dim=0) + dst_sf_w3.copy_( + maybe_pad_for_mxfp4(w3_scales[e].to(device), h_pad_fc1 // SCALING_VECTOR_SIZE, i_pad) + ) + dst_sf_w1.copy_( + maybe_pad_for_mxfp4(w1_scales[e].to(device), h_pad_fc1 // SCALING_VECTOR_SIZE, i_pad) + ) + permute_sf = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + dst_sf.view(float4_sf_dtype), + _CACHE_PERMUTE_INDICES, + EPILOGUE_TILE_M, + num_elts_per_sf=SCALING_VECTOR_SIZE, + ) + shuffled_sf = torch.ops.trtllm.shuffle_matrix( + dst_sf.view(float4_sf_dtype), permute_sf.to(device) + ) + dst_sf.copy_( + torch.ops.trtllm.block_scale_interleave( + shuffled_sf.view(float4_sf_dtype).reshape(dst_sf.shape) + ) + .view(torch.uint8) + .reshape(dst_sf.shape) + ) + + # ---- FC2 weights: pad + shuffle. + dst2 = gemm2_weights[e] + dst2.copy_(maybe_pad_for_mxfp4(w2_packed[e].to(device), i_pad // 2, h_pad_fc2)) + permute2 = trtllmgen_maybe_get_cached_w2_permute_indices( + dst2, _CACHE_PERMUTE_INDICES, EPILOGUE_TILE_M + ) + dst2.copy_(torch.ops.trtllm.shuffle_matrix(dst2, permute2.to(device))) + + # ---- FC2 scales: pad + shuffle with sf indices + interleave. + dst2_sf = gemm2_weights_scale[e] + dst2_sf.copy_( + maybe_pad_for_mxfp4(w2_scales[e].to(device), i_pad // SCALING_VECTOR_SIZE, h_pad_fc2) + ) + permute2_sf = trtllmgen_maybe_get_cached_w2_permute_indices( + dst2_sf.view(float4_sf_dtype), + _CACHE_PERMUTE_INDICES, + EPILOGUE_TILE_M, + num_elts_per_sf=SCALING_VECTOR_SIZE, + ) + shuffled2_sf = torch.ops.trtllm.shuffle_matrix( + dst2_sf.view(float4_sf_dtype), permute2_sf.to(device) + ) + dst2_sf.copy_( + torch.ops.trtllm.block_scale_interleave(shuffled2_sf.view(float4_sf_dtype)) + .view(torch.uint8) + .reshape(dst2_sf.shape) + ) + + return { + "gemm1_weights": gemm1_weights, + "gemm1_weights_scale": gemm1_weights_scale, + "gemm2_weights": gemm2_weights, + "gemm2_weights_scale": gemm2_weights_scale, + } + + +def make_situ_alpha_beta( + *, + local_num_experts: int, + situ_beta: float, + situ_linear_beta: float, + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build the per-expert CUDA float32 alpha/beta buffers for the op boundary. + + The kernel reads one alpha/beta pair per local expert; scalar broadcast is + not part of the contract. ``alpha`` is the gate-side (x1) parameter and maps + to Kimi's ``activation_situ_beta``; ``beta`` is the linear-side (x0) + parameter and maps to ``activation_situ_linear_beta``. + """ + if situ_beta <= 0.0 or situ_linear_beta <= 0.0: + raise RuntimeError( + f"SiTu alpha/beta must be > 0 (got alpha={situ_beta}, beta={situ_linear_beta})" + ) + gemm1_alpha = torch.full( + (local_num_experts,), float(situ_beta), dtype=torch.float32, device=device + ).contiguous() + gemm1_beta = torch.full( + (local_num_experts,), float(situ_linear_beta), dtype=torch.float32, device=device + ).contiguous() + return gemm1_alpha, gemm1_beta + + +def invoke_native_situ_moe( + *, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + gemm1_weights: torch.Tensor, + gemm1_weights_scale: torch.Tensor, + gemm2_weights: torch.Tensor, + gemm2_weights_scale: torch.Tensor, + gemm1_alpha: torch.Tensor, + gemm1_beta: torch.Tensor, + num_experts: int, + top_k: int, + valid_hidden_size: int, + valid_intermediate_size: int, + local_expert_offset: int = 0, + local_num_experts: Optional[int] = None, + act_type: Optional[int] = None, + tune_max_num_tokens: int = 8192, +) -> torch.Tensor: + """Run routed-expert compute through the in-tree SiTU fused MoE op. + + Parameters + ---------- + hidden_states + bf16 ``[num_tokens, valid_hidden_size]`` routed input (unpadded). + topk_ids / topk_weights + Precomputed K3 routing (any integer/float dtype; converted to the + op's int32/bfloat16 contract here). ``topk_weights`` must already + include renormalization and ``routed_scaling_factor`` — the bypass + feeds them into finalize verbatim. ``topk_ids`` are GLOBAL expert + ids; under EP the kernel skips ids outside + ``[local_expert_offset, local_expert_offset + local_num_experts)`` + and those tokens contribute zeros (caller allreduces partials). + gemm*_weights / gemm*_weights_scale + Output of :func:`pack_routed_expert_weights` for this rank's + local expert slice (leading dim = ``local_num_experts``). + gemm1_alpha / gemm1_beta + Output of :func:`make_situ_alpha_beta` (``[local_num_experts]``). + + Returns bf16 ``[num_tokens, valid_hidden_size]``. + """ + from tensorrt_llm._torch.utils import ActType_TrtllmGen + + if act_type is None: + act_type = int(ActType_TrtllmGen.SiTu) + if local_num_experts is None: + local_num_experts = num_experts + if gemm1_weights.shape[0] != local_num_experts: + raise RuntimeError( + f"gemm1_weights holds {gemm1_weights.shape[0]} experts but " + f"local_num_experts={local_num_experts}" + ) + + if hidden_states.dtype != torch.bfloat16: + raise RuntimeError(f"native SiTU MoE expects bf16 hidden_states, got {hidden_states.dtype}") + if not hidden_states.is_cuda: + raise RuntimeError("native SiTU MoE requires CUDA hidden_states") + num_tokens, hidden_size = hidden_states.shape + if hidden_size != valid_hidden_size: + raise RuntimeError( + f"hidden_states last dim {hidden_size} != valid_hidden_size {valid_hidden_size}" + ) + + intermediate_size_padded = gemm1_weights.shape[-2] // 2 + hidden_padded_fc1 = gemm1_weights.shape[-1] * 2 + expected_h_pad, _, expected_i_pad = padded_fused_shapes( + valid_hidden_size, valid_intermediate_size + ) + if hidden_padded_fc1 != expected_h_pad or intermediate_size_padded != expected_i_pad: + raise RuntimeError( + f"fused weight shapes do not match the padding contract: " + f"padded hidden {hidden_padded_fc1} (expected {expected_h_pad}, " + f"valid {valid_hidden_size}), padded intermediate " + f"{intermediate_size_padded} (expected {expected_i_pad}, " + f"valid {valid_intermediate_size})" + ) + + # topk contract: int32 / bfloat16, contiguous, [num_tokens, top_k]. + topk_ids = topk_ids.to(device=hidden_states.device, dtype=torch.int32).contiguous() + topk_weights = topk_weights.to(device=hidden_states.device, dtype=torch.bfloat16).contiguous() + if topk_ids.shape != (num_tokens, top_k) or topk_weights.shape != (num_tokens, top_k): + raise RuntimeError( + f"topk tensors must be [num_tokens={num_tokens}, top_k={top_k}]; got " + f"topk_ids {tuple(topk_ids.shape)}, topk_weights {tuple(topk_weights.shape)}" + ) + + # Dynamic MXFP8 activation quantization; pads hidden to the FC1 alignment. + x_fp8, x_sf = torch.ops.trtllm.mxfp8_quantize( + hidden_states.contiguous(), False, alignment=INPUT_HIDDEN_ALIGNMENT + ) + + output = torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner( + None, # routing_logits — unused with the topk bypass + None, # routing_bias + x_fp8, + x_sf.flatten(), + gemm1_weights, + gemm1_weights_scale, + None, # gemm1_bias + gemm1_alpha, + gemm1_beta, + None, # gemm1_clamp_limit — nullptr means +inf for clmp kernels + gemm2_weights, + gemm2_weights_scale, + None, # gemm2_bias + num_experts, + top_k, + None, # n_group — K3 has no expert groups + None, # topk_group + intermediate_size_padded, + valid_hidden_size, + valid_intermediate_size, + local_expert_offset, + local_num_experts, + None, # routed_scaling_factor — already folded into topk_weights + 1, # routing_method_type (Renormalize) — inert under the topk bypass + act_type, + topk_weights, + topk_ids, + tune_max_num_tokens=tune_max_num_tokens, + ) + + if output.shape[-1] > valid_hidden_size: + output = output[:, :valid_hidden_size].contiguous() + return output diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py new file mode 100644 index 000000000000..7d91294462a1 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP4 (E2M1 + E8M0) group-scaled quantization utilities. + +Kimi K3's routed expert Linear weights are stored in the +``mxfp4-pack-quantized`` format with ``group_size=32`` (see +``config.json::quantization_config``). Each 32-element group along the +input axis has a single ``uint8`` E8M0 scale (biased exponent, value +``2 ** (scale_u8 - 127)``) and each fp4 element is one of the eight +representable E2M1 magnitudes ``{0, 0.5, 1, 1.5, 2, 3, 4, 6}`` signed. +Two fp4 values pack into a single ``uint8``: low nibble at even element +index, high nibble at odd element index. + +The pack/unpack path is pure PyTorch so random-weight tests can +construct MXFP4 tensors on-the-fly without a full quantization +pipeline. Correctness properties tests depend on: + +* ``quantize_last_dim_mxfp4`` picks the E8M0 scale so the largest fp4 + magnitude ``6.0`` covers the group's peak absolute value, then + quantizes each element to the nearest representable magnitude. +* ``dequantize_last_dim_mxfp4`` returns exactly what was stored — no + rounding, no dtype loss beyond the E2M1 grid the values were + quantized onto. So ``dequantize(quantize(x))`` is idempotent. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +# E2M1 representable magnitudes at fp4 codes 0..7. Signed via the top +# bit: codes 8..15 are the negatives of codes 0..7. +_FP4_MAGNITUDES = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, +) +_FP4_VALUES = torch.cat([_FP4_MAGNITUDES, -_FP4_MAGNITUDES], dim=0) + +FP4_MAX_MAGNITUDE = 6.0 +E8M0_BIAS = 127 +E8M0_MIN = 0 +E8M0_MAX = 255 +DEFAULT_GROUP_SIZE = 32 + + +def _pack_two_nibbles(fp4_codes: torch.Tensor) -> torch.Tensor: + assert fp4_codes.dtype == torch.uint8, fp4_codes.dtype + assert fp4_codes.shape[-1] % 2 == 0, fp4_codes.shape + lo = fp4_codes[..., 0::2] & 0x0F + hi = fp4_codes[..., 1::2] & 0x0F + return (hi << 4) | lo + + +def _unpack_two_nibbles(packed: torch.Tensor) -> torch.Tensor: + assert packed.dtype == torch.uint8, packed.dtype + lo = packed & 0x0F + hi = (packed >> 4) & 0x0F + stacked = torch.stack([lo, hi], dim=-1) + return stacked.reshape(*packed.shape[:-1], packed.shape[-1] * 2) + + +def _quantize_group_e2m1( + group: torch.Tensor, atol_e8m0: float = 1e-30 +) -> Tuple[torch.Tensor, torch.Tensor]: + max_abs = group.abs().amax(dim=-1) + zero_mask = max_abs < atol_e8m0 + ratio = max_abs / FP4_MAX_MAGNITUDE + ratio = torch.where(zero_mask, torch.ones_like(ratio), ratio) + log2r = torch.log2(ratio) + e_signed = torch.ceil(log2r).to(torch.int64) + E8M0_BIAS + e_signed = e_signed.clamp(min=E8M0_MIN, max=E8M0_MAX) + e_signed = torch.where(zero_mask, torch.zeros_like(e_signed), e_signed) + scale_u8 = e_signed.to(torch.uint8) + + scale_val = torch.pow( + torch.tensor(2.0, dtype=torch.float32, device=group.device), + (scale_u8.to(torch.float32) - E8M0_BIAS), + ) + scaled = group / scale_val.unsqueeze(-1) + + values = _FP4_VALUES.to(device=group.device) + diffs = (scaled.unsqueeze(-1) - values.view(*(1,) * scaled.ndim, 16)).abs() + codes = diffs.argmin(dim=-1).to(torch.uint8) + return codes, scale_u8 + + +def _dequantize_group_e2m1(codes: torch.Tensor, scale_u8: torch.Tensor) -> torch.Tensor: + values = _FP4_VALUES.to(device=codes.device) + magnitudes = values[codes.to(torch.long)] + scale_val = torch.pow( + torch.tensor(2.0, dtype=torch.float32, device=codes.device), + (scale_u8.to(torch.float32) - E8M0_BIAS), + ) + return magnitudes * scale_val.unsqueeze(-1) + + +def quantize_last_dim_mxfp4( + x: torch.Tensor, group_size: int = DEFAULT_GROUP_SIZE +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize an arbitrary-shape fp32 tensor along the last dim to MXFP4. + + Returns ``(packed_u8, scales_u8)`` where: + + * ``packed_u8`` has shape ``x.shape[:-1] + (x.shape[-1] // 2,)``, + dtype uint8; two fp4 codes packed per byte. + * ``scales_u8`` has shape ``x.shape[:-1] + (x.shape[-1] // group_size,)``, + dtype uint8; one E8M0 scale byte per group of ``group_size`` elements + along the last dim. + + Constraints: ``x.shape[-1] % group_size == 0`` and + ``group_size % 2 == 0``. + """ + assert x.dtype == torch.float32, f"expected fp32 input, got {x.dtype}" + assert x.shape[-1] % group_size == 0, ( + f"last dim {x.shape[-1]} not divisible by group_size {group_size}" + ) + assert group_size % 2 == 0, group_size + + last = x.shape[-1] + num_groups = last // group_size + lead = x.shape[:-1] + + grouped = x.reshape(*lead, num_groups, group_size) + codes, scales = _quantize_group_e2m1(grouped) + + codes_flat = codes.reshape(*lead, num_groups * group_size) + packed = _pack_two_nibbles(codes_flat) + return packed, scales + + +def dequantize_last_dim_mxfp4( + packed_u8: torch.Tensor, + scales_u8: torch.Tensor, + group_size: int = DEFAULT_GROUP_SIZE, +) -> torch.Tensor: + """Inverse of :func:`quantize_last_dim_mxfp4`. Returns fp32.""" + assert packed_u8.dtype == torch.uint8, packed_u8.dtype + assert scales_u8.dtype == torch.uint8, scales_u8.dtype + lead = packed_u8.shape[:-1] + n_over_2 = packed_u8.shape[-1] + n = n_over_2 * 2 + num_groups = scales_u8.shape[-1] + assert num_groups * group_size == n, ( + f"num_groups {num_groups} * group_size {group_size} != last_dim {n}" + ) + + codes_flat = _unpack_two_nibbles(packed_u8) + codes_grouped = codes_flat.reshape(*lead, num_groups, group_size) + deq_grouped = _dequantize_group_e2m1(codes_grouped, scales_u8) + return deq_grouped.reshape(*lead, n) + + +def canonical_mxfp4_fp32(x: torch.Tensor, group_size: int = DEFAULT_GROUP_SIZE) -> torch.Tensor: + """Round-trip ``x`` (fp32) through MXFP4 pack and unpack. + + Convenience for tests: yields the fp32 value that a stored MXFP4 + weight actually decodes to. Callers use this to initialize a + reference (unquantized) module with the same values a K3 MXFP4 + weight produces, giving byte-exact parity. + """ + packed, scales = quantize_last_dim_mxfp4(x, group_size=group_size) + return dequantize_last_dim_mxfp4(packed, scales, group_size=group_size) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py new file mode 100644 index 000000000000..3acc33e5b24a --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py @@ -0,0 +1,677 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KimiK3SparseMoeBlock — in-tree Kimi K3 sparse MoE module. + +Structural mirror of HF ``KimiSparseMoeBlock`` at +``modeling_kimi.py:806-918`` end to end: + +* :class:`KimiK3MoEGate` for routing (see :mod:`kimi_k3_moe_gate`). +* :class:`KimiK3RoutedExpertBank` — per-expert MXFP4-packed + ``w1 / w2 / w3`` linear weights (group_size=32), dequantized on the + fly during the Python fallback path. +* :class:`KimiK3MLP` shared expert stack (``num_shared_experts`` fused + into one KimiMLP with ``intermediate_size = moe_intermediate_size * + num_shared_experts``). +* Latent projections ``routed_expert_down_proj`` / + ``routed_expert_up_proj`` around the routed compute when + ``routed_expert_hidden_size`` is set, plus optional + ``routed_expert_norm`` (:class:`KimiK3RMSNorm`). + +Two mutually exclusive kernel paths coexist: + +* ``use_fused_cubin=False`` (default) — Python fallback with MXFP4 + bank + activation. Byte-exact HF parity under random weights when + weights are canonicalized via :func:`copy_hf_moe_block_weights`. +* ``use_fused_cubin=True`` — native in-tree SiTU path through + ``torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner`` + (``act_type=SiTu``) on checkpoint-derived MXFP4 weights. The same + MXFP4 bank is the quantization source of truth for both paths, so + fused-vs-fallback comparisons differ only by activation quantization + (MXFP8) and kernel arithmetic. Routing uses the K3 gate's real + top-k via the op's ``topk_weights``/``topk_ids`` bypass; weights must + be loaded through :func:`copy_hf_moe_block_weights` (or + :meth:`KimiK3SparseMoeBlock.build_fused_weights`) before forward. + +Two mutation flags cover the negative controls required by AC4: + +* ``missing_shared_experts_mutation`` — skip the + ``+ shared_experts(identity)`` addition. +* ``non_situ_activation_mutation`` — replace SiTU with a SiLU-based + activation in both routed and shared experts. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, List, Optional, Tuple + +import torch +from torch import nn + +from ._mlp import KimiK3MLP, KimiK3RMSNorm, NonSituActivation, SituAndMul +from ._moe_kernels import ( + assert_native_situ_supported, + invoke_native_situ_moe, + make_situ_alpha_beta, + pack_routed_expert_weights, +) +from ._mxfp4 import DEFAULT_GROUP_SIZE, dequantize_last_dim_mxfp4, quantize_last_dim_mxfp4 +from .kimi_k3_moe_gate import KimiK3MoEGate + + +class KimiK3RoutedExpertBank(nn.Module): + """MXFP4-packed routed expert weight bank. + + Stores per-expert ``w1``, ``w2``, ``w3`` (matching HF + ``KimiBlockSparseMLP``'s naming) in ``mxfp4-pack-quantized`` form + with ``group_size=32``. + + Shapes (all ``num_experts`` on the leading dim): + + * ``w1_packed`` / ``w3_packed``: ``uint8 [E, intermediate, hidden // 2]`` + * ``w2_packed``: ``uint8 [E, hidden, intermediate // 2]`` + * ``w1_scales`` / ``w3_scales``: ``uint8 [E, intermediate, hidden // group_size]`` + * ``w2_scales``: ``uint8 [E, hidden, intermediate // group_size]`` + + ``hidden`` is the effective routed-expert input size (i.e. + ``moe_hidden_size = routed_expert_hidden_size`` when the latent + path is active, else ``config.hidden_size``) and ``intermediate`` is + ``moe_intermediate_size``. + + All tensors are registered as buffers (not Parameters) since MXFP4 + packed weights are integer-typed and not differentiable inputs for + the module tests. This module never trains. + """ + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + group_size: int = DEFAULT_GROUP_SIZE, + activation: Optional[nn.Module] = None, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + assert hidden_size % group_size == 0, ( + f"hidden_size {hidden_size} not divisible by group_size {group_size}" + ) + assert intermediate_size % group_size == 0, ( + f"intermediate_size {intermediate_size} not divisible by group_size {group_size}" + ) + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.group_size = group_size + self.activation = ( + activation if activation is not None else SituAndMul(beta=4.0, linear_beta=25.0) + ) + + h = hidden_size + i = intermediate_size + g = group_size + dev = device + + # w1, w3: [I, H]. w2: [H, I]. + # NOTE: allocated with torch.empty (not zeros) so the buffers stay on + # the meta device under the PyTorch backend's MetaInitMode and get + # materialized directly on CUDA; real checkpoints overwrite every + # element, and tests populate via store_expert before any forward. + self.register_buffer( + "w1_packed", + torch.empty(num_experts, i, h // 2, dtype=torch.uint8, device=dev), + ) + self.register_buffer( + "w1_scales", + torch.empty(num_experts, i, h // g, dtype=torch.uint8, device=dev), + ) + self.register_buffer( + "w3_packed", + torch.empty(num_experts, i, h // 2, dtype=torch.uint8, device=dev), + ) + self.register_buffer( + "w3_scales", + torch.empty(num_experts, i, h // g, dtype=torch.uint8, device=dev), + ) + self.register_buffer( + "w2_packed", + torch.empty(num_experts, h, i // 2, dtype=torch.uint8, device=dev), + ) + self.register_buffer( + "w2_scales", + torch.empty(num_experts, h, i // g, dtype=torch.uint8, device=dev), + ) + + def store_expert( + self, + expert_idx: int, + w1_fp32: torch.Tensor, + w2_fp32: torch.Tensor, + w3_fp32: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize fp32 ``[out, in]`` weights and store them. + + Returns the fp32 round-tripped (canonical) values for each + weight so callers can initialise a reference module with the + exact numerical values the bank now holds. + """ + assert w1_fp32.shape == (self.intermediate_size, self.hidden_size), ( + w1_fp32.shape, + (self.intermediate_size, self.hidden_size), + ) + assert w2_fp32.shape == (self.hidden_size, self.intermediate_size), ( + w2_fp32.shape, + (self.hidden_size, self.intermediate_size), + ) + assert w3_fp32.shape == w1_fp32.shape + + w1_packed, w1_scales = quantize_last_dim_mxfp4(w1_fp32, self.group_size) + w2_packed, w2_scales = quantize_last_dim_mxfp4(w2_fp32, self.group_size) + w3_packed, w3_scales = quantize_last_dim_mxfp4(w3_fp32, self.group_size) + self.w1_packed[expert_idx].copy_(w1_packed) + self.w1_scales[expert_idx].copy_(w1_scales) + self.w3_packed[expert_idx].copy_(w3_packed) + self.w3_scales[expert_idx].copy_(w3_scales) + self.w2_packed[expert_idx].copy_(w2_packed) + self.w2_scales[expert_idx].copy_(w2_scales) + + w1_canon = dequantize_last_dim_mxfp4(w1_packed, w1_scales, self.group_size) + w2_canon = dequantize_last_dim_mxfp4(w2_packed, w2_scales, self.group_size) + w3_canon = dequantize_last_dim_mxfp4(w3_packed, w3_scales, self.group_size) + return w1_canon, w2_canon, w3_canon + + def dequantize_expert(self, expert_idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + w1 = dequantize_last_dim_mxfp4( + self.w1_packed[expert_idx], self.w1_scales[expert_idx], self.group_size + ) + w2 = dequantize_last_dim_mxfp4( + self.w2_packed[expert_idx], self.w2_scales[expert_idx], self.group_size + ) + w3 = dequantize_last_dim_mxfp4( + self.w3_packed[expert_idx], self.w3_scales[expert_idx], self.group_size + ) + return w1, w2, w3 + + def forward_expert(self, expert_idx: int, tokens: torch.Tensor) -> torch.Tensor: + """Run a single expert on ``tokens`` ``[N, H]``. + + Matches HF ``KimiBlockSparseMLP.forward`` with the ``situ`` + branch: ``cat[w1(x), w3(x)] → SituAndMul → w2``. Compute is in + ``tokens.dtype``; weights are dequantized to canonical fp32 then + cast down. Because canonical MXFP4 magnitudes are exactly + representable in bf16, the cast is lossless. + """ + w1_f, w2_f, w3_f = self.dequantize_expert(expert_idx) + dt = tokens.dtype + w1 = w1_f.to(dt) if w1_f.dtype != dt else w1_f + w2 = w2_f.to(dt) if w2_f.dtype != dt else w2_f + w3 = w3_f.to(dt) if w3_f.dtype != dt else w3_f + gate = tokens @ w1.t() + up = tokens @ w3.t() + gate_up = torch.cat([gate, up], dim=-1) + act = self.activation(gate_up) + y = act @ w2.t() + return y + + +@dataclass +class MoEBlockProvenance: + """Provenance record returned by :func:`copy_hf_moe_block_weights`.""" + + n_experts: int + shared_expert_names: List[str] + routed_expert_layout: Tuple[int, int] + latent: bool + latent_use_norm: bool + canonicalized: bool + + +class KimiK3SparseMoeBlock(nn.Module): + """Kimi K3 sparse MoE block — mirrors HF ``KimiSparseMoeBlock``. + + Parameters + ---------- + config + ``KimiLinearConfig``-like — provides ``hidden_size``, + ``num_experts``, ``num_experts_per_token``, ``moe_intermediate_size``, + ``moe_renormalize``, ``routed_expert_hidden_size`` (optional; + enables the latent path), ``latent_moe_use_norm``, + ``num_shared_experts``, ``activation_situ_beta``, + ``activation_situ_linear_beta``, ``rms_norm_eps``. + missing_shared_experts_mutation + Mutation control — skip the ``+ shared_experts(identity)`` add. + non_situ_activation_mutation + Mutation control — replace SiTU with SiLU-based activation in + both routed and shared experts. + use_fused_cubin + When True, ``forward()`` dispatches routed compute through the + in-tree ``torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner`` + custom op with ``act_type=SiTu``. The MXFP4 ``expert_bank`` is + still allocated (it is the checkpoint-quantization source of + truth); the fused device buffers are derived from it by + :meth:`build_fused_weights`, which + :func:`copy_hf_moe_block_weights` calls automatically. Forward + raises if the fused weights have not been built — there is no + random-weight fallback. + dtype + Dtype used for latent projections, shared experts, RMSNorm + weight. Default fp32. + Set to bf16 to match HF's bf16 forward at real K3 dims. + """ + + def __init__( + self, + config: Any, + *, + missing_shared_experts_mutation: bool = False, + non_situ_activation_mutation: bool = False, + use_fused_cubin: bool = False, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_token + self.moe_renormalize = config.moe_renormalize + + self.use_latent_moe = getattr(config, "routed_expert_hidden_size", None) is not None + self.moe_hidden_size = ( + config.routed_expert_hidden_size if self.use_latent_moe else config.hidden_size + ) + self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False) + + # EP sharding trivialized — module tests run on one GPU. + self.ep_size = 1 + self.experts_per_rank = config.num_experts + self.ep_rank = 0 + + self.missing_shared_experts_mutation = missing_shared_experts_mutation + self.non_situ_activation_mutation = non_situ_activation_mutation + self.use_fused_cubin = use_fused_cubin + self._proj_dtype = dtype + self._cubin_call_count = 0 + + situ_beta = getattr(config, "activation_situ_beta", 4.0) + situ_linear_beta = getattr(config, "activation_situ_linear_beta", 25.0) + self._situ_beta = situ_beta + self._situ_linear_beta = situ_linear_beta + + if non_situ_activation_mutation: + routed_activation: nn.Module = NonSituActivation() + else: + routed_activation = SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta) + + self.gate = KimiK3MoEGate(config, device=device) + + # Routed expert storage — the MXFP4 bank is always the checkpoint + # quantization source of truth. The fused path derives its packed + # and shuffled device buffers from the same bank so both paths see + # identical canonical weights. + self.expert_bank = KimiK3RoutedExpertBank( + num_experts=config.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + activation=routed_activation, + device=device, + ) + self._fused_bank_ready = False + self.gemm1_weights: Optional[torch.Tensor] = None + self.gemm1_weights_scale: Optional[torch.Tensor] = None + self.gemm2_weights: Optional[torch.Tensor] = None + self.gemm2_weights_scale: Optional[torch.Tensor] = None + self._gemm1_alpha: Optional[torch.Tensor] = None + self._gemm1_beta: Optional[torch.Tensor] = None + if use_fused_cubin: + # Fail before any weight processing when the platform cannot run + # the fused path at all. + assert_native_situ_supported( + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + ) + if non_situ_activation_mutation: + raise RuntimeError( + "non_situ_activation_mutation is a Python-reference mutation " + "control; it cannot be combined with use_fused_cubin=True" + ) + + # Shared experts — HF fuses ``num_shared_experts`` KimiMLPs into + # one, with ``intermediate_size = moe_intermediate_size * + # num_shared_experts``. Unquantized Linear in ``_proj_dtype``. + self.num_shared_experts = getattr(config, "num_shared_experts", None) + if self.num_shared_experts is not None: + shared_activation: nn.Module = ( + NonSituActivation() + if non_situ_activation_mutation + else SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta) + ) + self.shared_experts = KimiK3MLP( + hidden_size=config.hidden_size, + intermediate_size=(config.moe_intermediate_size * self.num_shared_experts), + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + activation=shared_activation, + dtype=dtype, + device=device, + ) + else: + self.shared_experts = None + + # Latent projections around routed compute. + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear( + config.hidden_size, + self.moe_hidden_size, + bias=False, + dtype=dtype, + device=device, + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, + config.hidden_size, + bias=False, + dtype=dtype, + device=device, + ) + if self.latent_moe_use_norm: + self.routed_expert_norm = KimiK3RMSNorm( + self.moe_hidden_size, + eps=getattr(config, "rms_norm_eps", 1e-5), + dtype=dtype, + device=device, + ) + else: + self.routed_expert_norm = None + else: + self.routed_expert_down_proj = None + self.routed_expert_up_proj = None + self.routed_expert_norm = None + + # ------------------------------------------------------------------ + # Fused path — native in-tree TRTLLM-Gen SiTU dispatch. + # ------------------------------------------------------------------ + + def build_fused_weights(self) -> None: + """Derive the fused TRTLLM-Gen device buffers from ``expert_bank``. + + Packs the bank's per-expert MXFP4 tensors into the padded and + shuffled ``gemm1_*``/``gemm2_*`` layout expected by + ``mxe4m3_mxe2m1_block_scale_moe_runner`` (w3 first / w1 second, + opposite of the HF gate-first order) and materializes the + per-expert alpha/beta CUDA buffers. Must be called after the bank + holds checkpoint weights (``copy_hf_moe_block_weights`` does this + automatically for fused-mode blocks). + """ + assert self.use_fused_cubin, "build_fused_weights requires use_fused_cubin=True" + device = torch.device(f"cuda:{torch.cuda.current_device()}") + packed = pack_routed_expert_weights( + w1_packed=self.expert_bank.w1_packed, + w1_scales=self.expert_bank.w1_scales, + w3_packed=self.expert_bank.w3_packed, + w3_scales=self.expert_bank.w3_scales, + w2_packed=self.expert_bank.w2_packed, + w2_scales=self.expert_bank.w2_scales, + device=device, + ) + self.gemm1_weights = packed["gemm1_weights"] + self.gemm1_weights_scale = packed["gemm1_weights_scale"] + self.gemm2_weights = packed["gemm2_weights"] + self.gemm2_weights_scale = packed["gemm2_weights_scale"] + self._gemm1_alpha, self._gemm1_beta = make_situ_alpha_beta( + local_num_experts=self.num_experts, + situ_beta=self._situ_beta, + situ_linear_beta=self._situ_linear_beta, + device=device, + ) + self._fused_bank_ready = True + + def _moe_infer_fused( + self, + routed_in: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + """Fused variant of :meth:`_moe_infer` via the in-tree SiTU op. + + ``topk_ids``/``topk_weights`` come from the K3 gate (bypass + contract; weights already renormalized and scaled). + """ + if not (self.use_fused_cubin and self._fused_bank_ready): + raise RuntimeError( + "fused SiTU path invoked, but fused weights were never built. " + "Load checkpoint weights via copy_hf_moe_block_weights() or call " + "build_fused_weights(); there is no random-weight fallback." + ) + result = invoke_native_situ_moe( + hidden_states=routed_in, + topk_ids=topk_ids, + topk_weights=topk_weights, + gemm1_weights=self.gemm1_weights, + gemm1_weights_scale=self.gemm1_weights_scale, + gemm2_weights=self.gemm2_weights, + gemm2_weights_scale=self.gemm2_weights_scale, + gemm1_alpha=self._gemm1_alpha, + gemm1_beta=self._gemm1_beta, + num_experts=self.num_experts, + top_k=self.top_k, + valid_hidden_size=self.moe_hidden_size, + valid_intermediate_size=self.config.moe_intermediate_size, + ) + self._cubin_call_count += 1 + return result + + # ------------------------------------------------------------------ + # Forward + Python fallback MoE inference. + # ------------------------------------------------------------------ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Reproduce HF ``KimiSparseMoeBlock.forward``. + + Python-fallback path: + + 1. ``identity = hidden_states`` (used for shared expert input). + 2. Gate → ``topk_idx, topk_weight`` from raw ``hidden_states``. + 3. Reshape hidden to ``[T, hidden_size]``. + 4. If latent, ``hidden = routed_expert_down_proj(hidden)``. + 5. :meth:`_moe_infer` — dispatch tokens across selected experts, + weight-mix outputs, reshape back to ``[..., moe_hidden]``. + 6. If latent, apply ``routed_expert_norm`` (optional) and + ``routed_expert_up_proj``. + 7. Reshape back to ``hidden_states`` shape. + 8. Add ``shared_experts(identity)`` if configured (mutation: + skip). + + Fused path (``use_fused_cubin=True``): same skeleton, but step 5 + is :meth:`_moe_infer_fused`, which invokes the in-tree + ``mxe4m3_mxe2m1_block_scale_moe_runner`` op with + ``act_type=SiTu``. The gate at step 2 runs in both modes — the + fused path feeds its real top-k through the op's + ``topk_weights``/``topk_ids`` bypass. Steps 4/6/8 still run + (outside the fused kernel's scope). + """ + identity = hidden_states + orig_shape = hidden_states.shape + + topk_idx, topk_weight = self.gate(hidden_states) + + flat = hidden_states.view(-1, self.hidden_size) + + if self.use_latent_moe: + routed_in = self.routed_expert_down_proj(flat.to(self._proj_dtype)) + else: + routed_in = flat.to(self._proj_dtype) + + if self.use_fused_cubin: + routed_in_bf16 = ( + routed_in if routed_in.dtype == torch.bfloat16 else routed_in.to(torch.bfloat16) + ) + y = self._moe_infer_fused(routed_in_bf16, topk_idx, topk_weight) + if y.dtype != self._proj_dtype: + y = y.to(self._proj_dtype) + else: + y = self._moe_infer(routed_in, topk_idx, topk_weight) + + if self.use_latent_moe: + if self.routed_expert_norm is not None: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + + y = y.view(*orig_shape) + + if self.shared_experts is not None and not self.missing_shared_experts_mutation: + shared = self.shared_experts(identity.to(self._proj_dtype)) + y = y + shared.to(y.dtype) + return y.to(hidden_states.dtype) + + def _moe_infer( + self, + x: torch.Tensor, + topk_ids: torch.Tensor, + topk_weight: torch.Tensor, + ) -> torch.Tensor: + """Byte-for-byte port of HF ``KimiSparseMoeBlock.moe_infer``. + + HF does per-expert compute in the token dtype (the expert Linear + layers inherit HF's module dtype), and casts to + ``topk_weight.dtype`` for the weighted sum. We mirror that so + both fp32 and bf16 HF modes stay byte-exact. + """ + cnts = topk_ids.new_zeros((topk_ids.shape[0], self.num_experts)) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + + tokens_per_expert_cpu = tokens_per_expert.cpu().tolist() + + outputs: List[torch.Tensor] = [] + start = 0 + for i, n_tokens in enumerate(tokens_per_expert_cpu): + end = start + int(n_tokens) + if n_tokens == 0: + continue + tokens_for_i = sorted_tokens[start:end] + expert_out = self.expert_bank.forward_expert(i, tokens_for_i) + outputs.append(expert_out) + start = end + + outs = torch.cat(outputs, dim=0) if outputs else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +# --------------------------------------------------------------------------- +# HF → K3 copy helper. +# --------------------------------------------------------------------------- + + +def copy_hf_moe_block_weights( + hf: nn.Module, + k3: KimiK3SparseMoeBlock, +) -> MoEBlockProvenance: + """Copy weights from HF ``KimiSparseMoeBlock`` to a K3 block. + + Steps: + + 1. Copy gate params (identity name mapping). + 2. Copy latent down/up projections and optional RMSNorm. + 3. Copy ``shared_experts`` (fused KimiMLP) parameters. + 4. For each routed expert: + a. Read HF ``experts[i].{w1,w2,w3}.weight`` fp32. + b. Quantize and store in ``k3.expert_bank``. + c. Retrieve canonical fp32 (quantize→dequantize) values. + d. Overwrite HF's Linear weights with the canonical values so + both modules see byte-identical numbers on forward. + 5. When ``k3.use_fused_cubin=True``, additionally derive the fused + TRTLLM-Gen device buffers from the freshly loaded bank via + :meth:`KimiK3SparseMoeBlock.build_fused_weights`. Both paths then + hold the same canonical checkpoint weights. + + Returns provenance metadata for logging. + """ + with torch.no_grad(): + k3.gate.weight.data.copy_(hf.gate.weight.data.to(k3.gate.weight.dtype)) + k3.gate.e_score_correction_bias.data.copy_( + hf.gate.e_score_correction_bias.data.to(k3.gate.e_score_correction_bias.dtype) + ) + + if k3.use_latent_moe: + with torch.no_grad(): + k3.routed_expert_down_proj.weight.data.copy_( + hf.routed_expert_down_proj.weight.data.to(k3.routed_expert_down_proj.weight.dtype) + ) + k3.routed_expert_up_proj.weight.data.copy_( + hf.routed_expert_up_proj.weight.data.to(k3.routed_expert_up_proj.weight.dtype) + ) + if k3.routed_expert_norm is not None: + assert hasattr(hf, "routed_expert_norm"), ( + "latent_moe_use_norm=True but HF has no routed_expert_norm" + ) + with torch.no_grad(): + k3.routed_expert_norm.weight.data.copy_( + hf.routed_expert_norm.weight.data.to(k3.routed_expert_norm.weight.dtype) + ) + + shared_names: List[str] = [] + hf_shared_experts = getattr(hf, "shared_experts", None) + if k3.shared_experts is None: + assert hf_shared_experts is None, ( + "HF block has shared_experts but K3 block has none; " + "shared-expert weights would be silently dropped" + ) + else: + assert hf_shared_experts is not None, ( + "K3 block has shared_experts but HF block has none; " + "shared-expert weights would stay randomly initialized" + ) + with torch.no_grad(): + gate_up_fused = torch.cat( + [ + hf.shared_experts.gate_proj.weight.data, + hf.shared_experts.up_proj.weight.data, + ], + dim=0, + ).to(k3.shared_experts.gate_up_proj.weight.dtype) + k3.shared_experts.gate_up_proj.weight.data.copy_(gate_up_fused) + k3.shared_experts.down_proj.weight.data.copy_( + hf.shared_experts.down_proj.weight.data.to(k3.shared_experts.down_proj.weight.dtype) + ) + shared_names = ["gate_up_proj.weight", "down_proj.weight"] + + for i in range(k3.num_experts): + expert = hf.experts[i] + w1 = expert.w1.weight.data.to(torch.float32) + w2 = expert.w2.weight.data.to(torch.float32) + w3 = expert.w3.weight.data.to(torch.float32) + w1c, w2c, w3c = k3.expert_bank.store_expert(i, w1, w2, w3) + # bf16 exactly represents every MXFP4 magnitude + # {0, 0.5, 1, 1.5, 2, 3, 4, 6} scaled by 2^e, so the cast to HF's + # own weight dtype is lossless and byte-parity holds under bf16 + # HF too. + with torch.no_grad(): + expert.w1.weight.data.copy_(w1c.to(expert.w1.weight.dtype)) + expert.w2.weight.data.copy_(w2c.to(expert.w2.weight.dtype)) + expert.w3.weight.data.copy_(w3c.to(expert.w3.weight.dtype)) + + if k3.use_fused_cubin: + k3.build_fused_weights() + + return MoEBlockProvenance( + n_experts=k3.num_experts, + shared_expert_names=shared_names, + routed_expert_layout=(k3.moe_hidden_size, k3.config.moe_intermediate_size), + latent=k3.use_latent_moe, + latent_use_norm=k3.routed_expert_norm is not None, + canonicalized=True, + ) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py new file mode 100644 index 000000000000..003175f2854f --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 MoE routing module. + +Structural mirror of HF ``KimiMoEGate`` (see +``model_config/modeling_kimi.py:710``). K3 routing inherits DeepSeek-V3's +``noaux_tc`` topology but pins the following K3 config choices (see +``configuration_kimi_k3.py``): + +* ``moe_router_activation_func = "sigmoid"`` — per-expert sigmoid, not softmax. +* ``e_score_correction_bias`` (per-expert) — bias added *only* for + top-k *selection*; the returned ``topk_weight`` samples the *raw* + sigmoid ``scores``, not the bias-adjusted ``scores_for_choice``. +* ``moe_renormalize = True`` — for ``top_k > 1``, divide ``topk_weight`` + by ``sum + 1e-20`` before scaling. +* ``routed_scaling_factor`` — final multiplicative scale. +* ``num_expert_group = 1``, ``topk_group = 1`` — K3 config disables the + grouped top-k branch used by DeepSeek. The gate still handles the + grouped branch when a caller flips those config knobs. + +Parameter names / shapes match HF ``KimiMoEGate`` so +:func:`copy_hf_moe_gate_weights` is identity name mapping. +""" + +from __future__ import annotations + +from typing import Any, Tuple + +import torch +from torch import nn + +from ..fused_moe.routing import DeepSeekV3MoeRoutingMethod, Deepseekv3RoutingImpl + + +class KimiK3MoEGate(nn.Module): + """K3 MoE routing — structural mirror of HF ``KimiMoEGate``. + + Positive path reproduces HF ``KimiMoEGate.forward`` at + ``modeling_kimi.py:747-803`` byte-identically under K3's + ``sigmoid`` scoring, top-k over the full expert set, raw sigmoid + weights, renormalization + scaling profile. + + Three mutation flags gate the negative controls required by AC6: + + * ``softmax_routing_mutation`` — softmax over experts instead of + per-expert sigmoid. + * ``biased_weights_mutation`` — gather ``topk_weight`` from the + bias-adjusted scores rather than the raw sigmoid scores. + * ``omit_renormalize_mutation`` — skip the renormalize step even + when the config asks for it. + """ + + def __init__( + self, + config: Any, + *, + softmax_routing_mutation: bool = False, + biased_weights_mutation: bool = False, + omit_renormalize_mutation: bool = False, + 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" + ) + + # Same parameter shapes / names as HF ``KimiMoEGate``. + # + # ``logits_gemm_dtype=torch.bfloat16`` stores the gate weight in + # bf16 and runs the logits GEMM as a single bf16xbf16 kernel with + # fp32 accumulate/output (``trtllm::dsv3_router_gemm_op``). The K3 + # checkpoint stores this weight in bf16, so the fp32 master was an + # exact upcast and bf16 storage is lossless; this removes the + # per-layer bf16->fp32 input cast + fp32 splitK-reduce that ran + # inside the decode CUDA graph (~5 us x 92 layers per step). + # Default ``None`` keeps the legacy fp32 GEMM (module parity tests). + 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, device=device)) + + self.softmax_routing_mutation = softmax_routing_mutation + self.biased_weights_mutation = biased_weights_mutation + self.omit_renormalize_mutation = omit_renormalize_mutation + + # Fast path: the fused ``noaux_tc`` routing kernel computes exactly K3's + # production routing contract in one launch -- per-expert sigmoid, + # ``e_score_correction_bias`` added for *selection* only, top-k weights + # sampled from the raw sigmoid scores, renormalized by ``sum + 1e-20``, + # then scaled by ``routed_scaling_factor``. Route through the shared + # ``Deepseekv3RoutingImpl`` (same op DeepSeek-V3 uses) when the config is + # eligible and none of the parity-breaking mutation controls are active. + # The eager path below stays the reference for those controls, for + # softmax scoring, for ``moe_renormalize=False``, and for grouped / + # oversized configs the kernel does not support. + self._routing_impl = Deepseekv3RoutingImpl( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + is_fused=True, + ) + # Bounds mirror the n_group == 1 branch of + # ``Deepseekv3RoutingImpl.noaux_tc`` (num_experts <= 1024, top_k <= 32); + # staying inside them guarantees the fused kernel branch is taken (never + # the impl's own PyTorch fallback, whose grouped path differs from K3's). + self._use_fused_routing = ( + self.moe_router_activation_func == "sigmoid" + and self.num_expert_group == 1 + and self.moe_renormalize + and self.top_k > 1 + and self.num_experts <= 1024 + and self.top_k <= 32 + and not softmax_routing_mutation + and not biased_weights_mutation + and not omit_renormalize_mutation + ) + + def _score(self, logits: torch.Tensor) -> torch.Tensor: + if self.softmax_routing_mutation: + return logits.softmax(dim=1) + if self.moe_router_activation_func == "sigmoid": + return logits.sigmoid() + return logits.softmax(dim=1) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Routing logits ``[num_tokens, num_experts]``, fp32, pre-sigmoid. + + Used when the MoE block is hosted under ``ConfigurableMoE``: the + post-linear gate math (sigmoid, bias-for-selection, renormalize, + ``routed_scaling_factor``) runs inside the wrapper's routing method + per chunk; only the gate GEMM stays here, keeping the checkpoint + parameter mapping identity. + """ + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + # Single bf16xbf16 -> fp32 GEMM (fp32 accumulate); no input + # upcast kernel, no fp32 splitK-reduce. K3's 896 experts miss + # the op's specialized 256-expert kernels and take its cublas + # path, which is the point here (one fused kernel). + 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) -> DeepSeekV3MoeRoutingMethod: + """Return the shared DeepSeekV3 router used by ConfigurableMoE.""" + 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." + ) + if ( + self.softmax_routing_mutation + or self.biased_weights_mutation + or self.omit_renormalize_mutation + ): + raise ValueError( + "Kimi K3 routing mutation flags are reference-test controls " + "and cannot be used by ConfigurableMoE." + ) + return DeepSeekV3MoeRoutingMethod( + 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, + ) + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + logits = self.compute_logits(hidden_states) + # ``compute_logits`` flattens to [num_tokens, num_experts]; derive + # the token count from it so any input rank works. + num_tokens = logits.shape[0] + + # ``trtllm::noaux_tc_op`` is a CUDA-only custom op; CPU inputs + # (reference / parity tests) fall through to the eager path below, + # which stays the routing-contract reference on every device. + if self._use_fused_routing and logits.is_cuda: + # One fused kernel replaces the sigmoid -> (+bias) -> top-k -> + # gather -> renormalize -> scale chain below. ``noaux_tc`` returns + # (weights, indices); return the eager dtype contract -- int64 + # indices (as ``torch.topk`` yields) and fp32 weights -- so every + # downstream consumer is byte-for-byte unaffected by the swap. + topk_weight, topk_idx = self._routing_impl.noaux_tc( + logits, self.e_score_correction_bias.float() + ) + return topk_idx.to(torch.int64), topk_weight.to(torch.float32) + + scores = self._score(logits) + scores = scores.view(num_tokens, -1) + + # Bias is applied for *selection*, not for the returned weight. + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + + if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: + group_scores = ( + scores_for_choice.view(num_tokens, self.num_expert_group, -1) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + num_tokens, + self.num_expert_group, + self.num_experts // self.num_expert_group, + ) + .reshape(num_tokens, -1) + ) + tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) + else: + tmp_scores = scores_for_choice + + _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False) + + # Positive contract: gather from raw ``scores`` (not bias-adjusted). + weight_source = scores_for_choice if self.biased_weights_mutation else scores + topk_weight = weight_source.gather(1, topk_idx) + + if self.top_k > 1 and self.moe_renormalize and not self.omit_renormalize_mutation: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + + topk_weight = topk_weight * self.routed_scaling_factor + return topk_idx, topk_weight + + +def copy_hf_moe_gate_weights( + hf: nn.Module, + k3: KimiK3MoEGate, +) -> dict[str, tuple[tuple[int, ...], str]]: + """Copy parameters from HF ``KimiMoEGate`` into ``k3``. + + Identity name mapping (``weight`` + ``e_score_correction_bias``). + Returns a ``{name: (shape, dtype)}`` provenance dict. + """ + src_params = dict(hf.named_parameters()) + dst_params = dict(k3.named_parameters()) + missing_on_k3 = sorted(set(src_params) - set(dst_params)) + missing_on_hf = sorted(set(dst_params) - set(src_params)) + if missing_on_k3: + raise KeyError(f"copy_hf_moe_gate_weights: HF params missing on K3: {missing_on_k3}") + if missing_on_hf: + raise KeyError(f"copy_hf_moe_gate_weights: K3 params missing on HF: {missing_on_hf}") + provenance = {} + for name, src in src_params.items(): + dst = dst_params[name] + if src.shape != dst.shape: + raise ValueError( + f"shape mismatch for {name}: HF {tuple(src.shape)} vs K3 {tuple(dst.shape)}" + ) + with torch.no_grad(): + dst.data.copy_(src.data.to(dtype=dst.dtype, device=dst.device)) + provenance[name] = (tuple(src.shape), str(src.dtype)) + return provenance diff --git a/tensorrt_llm/_torch/modules/kimi_kda/__init__.py b/tensorrt_llm/_torch/modules/kimi_kda/__init__.py new file mode 100644 index 000000000000..c457a986744c --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi Delta Attention (KDA) in-tree module for TensorRT-LLM's PyTorch backend. + +KDA is the linear-attention block used at ``linear_attn_config.kda_layers`` +positions in the Kimi K3 text-core. It carries a short-convolution state and a +delta-rule recurrent state per layer, so it follows the hybrid-cache / +mamba ownership pattern rather than the paged-KV FMHA attention-backend +interface. +""" + +from .kimi_kda_mixer import KimiKDAKernelPath, KimiKDALinearAttention + +__all__ = [ + "KimiKDAKernelPath", + "KimiKDALinearAttention", +] diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py new file mode 100644 index 000000000000..165912f219be --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py @@ -0,0 +1,221 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""PyTorch wrapper for KDA decode fusion.""" + +from __future__ import annotations + +import torch + +_DUMMY_CACHE: dict = {} + + +def _require_cuda_bf16(name: str, tensor: torch.Tensor) -> None: + """Validate that a tensor is CUDA bf16.""" + if not tensor.is_cuda or tensor.dtype is not torch.bfloat16: + raise TypeError(f"{name} must be a CUDA bfloat16 tensor") + + +def _require_cuda_fp32(name: str, tensor: torch.Tensor) -> None: + """Validate that a tensor is CUDA fp32.""" + if not tensor.is_cuda or tensor.dtype is not torch.float32: + raise TypeError(f"{name} must be a CUDA float32 tensor") + + +def _dummy_tensor( + tag: str, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + fill: float = 0.0, +) -> torch.Tensor: + """Return a cached dummy tensor for optional CUDA arguments.""" + device = torch.device(device) + key = (tag, shape, dtype, device.type, device.index, fill) + tensor = _DUMMY_CACHE.get(key) + if tensor is None: + if fill == 0.0: + tensor = torch.zeros(shape, dtype=dtype, device=device) + elif fill == 1.0: + tensor = torch.ones(shape, dtype=dtype, device=device) + else: + tensor = torch.full(shape, fill, dtype=dtype, device=device) + _DUMMY_CACHE[key] = tensor + return tensor + + +def run_kda_decode_fusion_cuda( + *, + x_q: torch.Tensor, + x_k: torch.Tensor, + x_v: torch.Tensor, + w_q_t: torch.Tensor, + w_k_t: torch.Tensor, + w_v_t: torch.Tensor, + bias_q: torch.Tensor | None, + bias_k: torch.Tensor | None, + bias_v: torch.Tensor | None, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + A_log: torch.Tensor, + g: torch.Tensor, + dt_bias: torch.Tensor | None, + beta: torch.Tensor, + state: torch.Tensor, + onorm_g: torch.Tensor | None = None, + onorm_weight: torch.Tensor | None = None, + out: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + scale: float = 128**-0.5, + onorm_eps: float = 1e-5, + lower_bound: float | None = None, + use_beta_sigmoid_in_kernel: bool = True, + verbose: bool = False, + update_conv_cache: bool = False, +) -> torch.Tensor: + """Run CUDA KDA decode fusion for the tuned decode shapes. + + ``ssm_state_indices=None`` selects the tuned batch-local static layout, + while a tensor selects the indexed state-pool layout. + """ + for name, tensor in ( + ("x_q", x_q), + ("x_k", x_k), + ("x_v", x_v), + ("w_q_t", w_q_t), + ("w_k_t", w_k_t), + ("w_v_t", w_v_t), + ("cs_q", cs_q), + ("cs_k", cs_k), + ("cs_v", cs_v), + ("g", g), + ("beta", beta), + ): + _require_cuda_bf16(name, tensor) + for name, tensor in (("A_log", A_log), ("state", state)): + _require_cuda_fp32(name, tensor) + + if x_q.ndim != 4 or x_k.ndim != 4 or x_v.ndim != 4: + raise ValueError("x_q, x_k, and x_v must be rank-4 decode tensors") + if x_q.shape[0] != 1 or x_k.shape[0] != 1 or x_v.shape[0] != 1: + raise ValueError("only T=1 decode inputs are supported") + if x_q.shape[-1] != 128 or x_k.shape[-1] != 128 or x_v.shape[-1] != 128: + raise ValueError("only K=128 and V=128 are supported") + + B = x_q.shape[1] + H = x_q.shape[2] + HV = x_v.shape[2] + if x_k.shape[1:3] != (B, H) or x_v.shape[1] != B: + raise ValueError("x_q, x_k, and x_v batch/head dimensions are inconsistent") + if H != HV or H not in (1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 96): + raise ValueError( + "CUDA KDA decode fusion supports H == HV in {1,2,3,4,6,8,12,16,24,32,48,96}" + ) + if ssm_state_indices is None and not state.is_contiguous(): + raise ValueError("state must be contiguous because it is updated in place") + if out is not None: + _require_cuda_bf16("out", out) + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + if tuple(out.shape) != (B, 1, HV, 128): + raise ValueError("out must have shape [B, 1, HV, 128]") + + device = x_q.device + apply_onorm = onorm_g is not None + if bias_q is None: + bias_q = _dummy_tensor("bias_q", (H * 128,), torch.bfloat16, device) + if bias_k is None: + bias_k = _dummy_tensor("bias_k", (H * 128,), torch.bfloat16, device) + if bias_v is None: + bias_v = _dummy_tensor("bias_v", (HV * 128,), torch.bfloat16, device) + if dt_bias is None: + dt_bias = _dummy_tensor("dt_bias", (H * 128,), torch.float32, device) + if onorm_g is None: + onorm_g = _dummy_tensor("onorm_g", (1, B, HV, 128), torch.bfloat16, device) + if onorm_weight is None: + onorm_weight = _dummy_tensor("onorm_weight", (128,), torch.float32, device, fill=1.0) + + for name, tensor in ( + ("bias_q", bias_q), + ("bias_k", bias_k), + ("bias_v", bias_v), + ("onorm_g", onorm_g), + ): + _require_cuda_bf16(name, tensor) + for name, tensor in (("dt_bias", dt_bias), ("onorm_weight", onorm_weight)): + _require_cuda_fp32(name, tensor) + + if update_conv_cache: + q_stride = H * 128 + v_stride = HV * 128 + if not ( + cs_q.stride(1) == 1 + and cs_k.stride(1) == 1 + and cs_v.stride(1) == 1 + and cs_q.stride(2) == q_stride + and cs_k.stride(2) == q_stride + and cs_v.stride(2) == v_stride + ): + raise ValueError( + "update_conv_cache expects transposed conv-state layout: " + "shape [B, dim, 3], stride(1)=1, stride(2)=dim" + ) + + if cu_seqlens is None: + cu_seqlens = torch.arange(B + 1, dtype=torch.int32, device=device) + else: + if not cu_seqlens.is_cuda or cu_seqlens.dtype is not torch.int32: + raise TypeError("cu_seqlens must be a CUDA int32 tensor") + if tuple(cu_seqlens.shape) != (B + 1,): + raise ValueError("cu_seqlens must have shape [B + 1]") + cu_seqlens = cu_seqlens.contiguous() + + args = ( + x_q.contiguous(), + x_k.contiguous(), + x_v.contiguous(), + w_q_t.contiguous(), + w_k_t.contiguous(), + w_v_t.contiguous(), + bias_q.contiguous(), + bias_k.contiguous(), + bias_v.contiguous(), + cs_q if update_conv_cache else cs_q.contiguous(), + cs_k if update_conv_cache else cs_k.contiguous(), + cs_v if update_conv_cache else cs_v.contiguous(), + A_log.contiguous(), + g.contiguous(), + dt_bias.contiguous(), + beta.contiguous(), + onorm_g.contiguous(), + onorm_weight.contiguous(), + ssm_state_indices, + cu_seqlens, + state, + ) + + use_lower_bound = lower_bound is not None + lower_bound_value = 0.0 if lower_bound is None else float(lower_bound) + launch_args = ( + bool(apply_onorm), + bool(update_conv_cache), + bool(use_lower_bound), + bool(use_beta_sigmoid_in_kernel), + lower_bound_value, + float(scale), + float(onorm_eps), + ) + return torch.ops.trtllm.kda_decode(*args, *launch_args, output=out) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py new file mode 100644 index 000000000000..283c721b1c1b --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -0,0 +1,390 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kernel dispatch for the in-tree KDA module. + +Both optimized KDA kernels are source-integrated into TensorRT-LLM: the +chunked prefill (CuTe DSL ``trtllm::kda_prefill``, see +``tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py``) and the +fused CUDA C++ single-token decode (``trtllm::kda_decode`` thop op wrapped +by ``_kda_decode``). Neither requires the external +``exisiting_optimization_work`` collection at runtime. + +On Blackwell (sm_100/sm_103) with the CuTe DSL toolchain available the +dispatch runs the optimized kernels. On any other arch (or when the in-tree +prefill op is unavailable) the dispatch falls back to FLA's ``chunk_kda`` / +``fused_recurrent_kda`` on-device references. + +Callers are the ``KimiKDALinearAttention`` module. The module owns the +HF-parity semantics (Q/K/V projections, conv, gating, ``o_norm``, +``o_proj``); the dispatch here only wraps the kernel-level computation of +the delta-rule inner loop plus its state update. +""" + +from __future__ import annotations + +import importlib +from types import ModuleType +from typing import Optional, Tuple + +import torch + +from . import _kda_decode + +try: + from tensorrt_llm._utils import get_sm_version as _tllm_get_sm_version +except ImportError: # pragma: no cover — source-loader stub path + _tllm_get_sm_version = None + + +def _default_get_sm_version() -> int: + if not torch.cuda.is_available() or torch.cuda.device_count() == 0: + return -1 + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor + + +def get_kda_sm_version() -> int: + """Return the runtime SM version used for KDA kernel selection. + + Prefers ``tensorrt_llm._utils.get_sm_version`` when the real package is + importable so environment-side overrides propagate. Falls back to a + plain CUDA-property probe when we are executing under the source-loader + stub subtree. + """ + if _tllm_get_sm_version is not None: + try: + return int(_tllm_get_sm_version()) + except RuntimeError: + # torch raises RuntimeError when no CUDA device is usable; + # the property probe below handles that case itself. + return _default_get_sm_version() + return _default_get_sm_version() + + +def is_kda_optimized_supported() -> bool: + """The optimized prefill/decode kernels are Blackwell sm_100 only.""" + return get_kda_sm_version() in (100, 103) + + +# --------------------------------------------------------------------------- +# In-tree KDA prefill op (CuTe DSL, trtllm::kda_prefill). +# --------------------------------------------------------------------------- + +_PREFILL_MODULE: Optional[ModuleType] = None +_PREFILL_IMPORT_ERROR: Optional[Exception] = None + + +def _load_prefill_module() -> ModuleType: + """Import the in-tree prefill custom-op module (registers the op).""" + global _PREFILL_MODULE, _PREFILL_IMPORT_ERROR + if _PREFILL_MODULE is not None: + return _PREFILL_MODULE + if _PREFILL_IMPORT_ERROR is not None: + raise _PREFILL_IMPORT_ERROR + try: + module = importlib.import_module( + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops" + ) + except Exception as exc: # typically ImportError when CuTe DSL is unavailable + _PREFILL_IMPORT_ERROR = exc + raise + _PREFILL_MODULE = module + return module + + +def is_intree_prefill_available() -> bool: + """True when the in-tree CuTe DSL prefill op can be imported.""" + try: + _load_prefill_module() + return True + except Exception: + return False + + +def _load_fla_chunk_kda() -> ModuleType: + return importlib.import_module("fla.ops.kda") + + +# --------------------------------------------------------------------------- +# In-tree KDA multi-token verify op (CuTe DSL, trtllm::kda_mtp_decode). +# --------------------------------------------------------------------------- + +_MTP_MODULE: Optional[ModuleType] = None +_MTP_IMPORT_ERROR: Optional[Exception] = None + + +def _load_mtp_module() -> ModuleType: + """Import the in-tree MTP verify custom-op module (registers the op).""" + global _MTP_MODULE, _MTP_IMPORT_ERROR + if _MTP_MODULE is not None: + return _MTP_MODULE + if _MTP_IMPORT_ERROR is not None: + raise _MTP_IMPORT_ERROR + try: + module = importlib.import_module( + "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops" + ) + except Exception as exc: # typically ImportError when CuTe DSL is unavailable + _MTP_IMPORT_ERROR = exc + raise + _MTP_MODULE = module + return module + + +def is_intree_mtp_available() -> bool: + """True when the in-tree CuTe DSL MTP verify op can be imported.""" + try: + _load_mtp_module() + return True + except Exception: + return False + + +def is_kda_mtp_verify_available() -> bool: + """True when the fused multi-token verify kernel can run here. + + Used by the executor (cache-manager sizing) to decide whether to + allocate the KDA replay caches instead of the legacy intermediate + verification buffers. + """ + return is_kda_optimized_supported() and is_intree_mtp_available() + + +# --------------------------------------------------------------------------- +# Dispatch API used by the KimiKDALinearAttention module. +# --------------------------------------------------------------------------- + + +class KDAKernelDispatch: + """Kernel dispatch state for one ``KimiKDALinearAttention`` instance. + + Attributes + ---------- + prefill_kernel_path : str + Selected prefill path: ``"optimized"`` or ``"fla"``. + decode_kernel_path : str + Selected decode path: ``"optimized"`` or ``"fla"``. + verify_kernel_path : str + Selected multi-token verify path: ``"optimized"`` (fused + ``trtllm::kda_mtp_decode`` replay kernel) or ``"fla"`` (sequential + per-step ``fused_recurrent_kda`` with intermediate-buffer state + promotion). + Notes + ----- + Prefill, decode, and verify dispatch are decided independently. All + require a supported GPU; the in-tree CuTe DSL ops additionally require + their modules to be importable. + """ + + _selection_logged = False + + def __init__( + self, + use_optimized_prefill: bool = True, + use_optimized_decode: bool = True, + use_optimized_verify: bool = True, + ) -> None: + optimized_supported = is_kda_optimized_supported() + self.prefill_kernel_path = "fla" + if use_optimized_prefill and optimized_supported and is_intree_prefill_available(): + self.prefill_kernel_path = "optimized" + self.decode_kernel_path = ( + "optimized" if use_optimized_decode and optimized_supported else "fla" + ) + self.verify_kernel_path = "fla" + if use_optimized_verify and optimized_supported and is_intree_mtp_available(): + self.verify_kernel_path = "optimized" + # One line per process so runs record which paths actually executed + # (the fallback is otherwise silent). + if not KDAKernelDispatch._selection_logged: + KDAKernelDispatch._selection_logged = True + try: + from tensorrt_llm.logger import logger + except ImportError: # pragma: no cover — source-loader stub path + pass + else: + logger.info( + f"KDA kernel dispatch: prefill={self.prefill_kernel_path} " + f"decode={self.decode_kernel_path} " + f"verify={self.verify_kernel_path}" + ) + + def get_prefill_source(self) -> str: + if self.prefill_kernel_path == "optimized": + return _load_prefill_module().__file__ or "" + return _load_fla_chunk_kda().__file__ or "" + + def get_decode_source(self) -> str: + if self.decode_kernel_path == "optimized": + return _kda_decode.__file__ or "" + return _load_fla_chunk_kda().__file__ or "" + + def mtp_verify(self, **kwargs) -> torch.Tensor: + """Run the fused KDA multi-token verify kernel. + + Thin passthrough to ``trtllm::kda_mtp_decode`` (see + ``custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py`` for the full + argument and state-management contract). Only defined on the + optimized path; the FLA fallback is the module's sequential + per-step loop with intermediate-buffer promotion. + """ + if self.verify_kernel_path != "optimized": + raise RuntimeError( + "mtp_verify called on non-optimized path; use the module's " + "sequential FLA verify fallback instead." + ) + _load_mtp_module() # registers trtllm::kda_mtp_decode + return torch.ops.trtllm.kda_mtp_decode(**kwargs) + + def prefill_chunk_kda( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: Optional[torch.Tensor], + safe_gate: bool, + lower_bound: Optional[float], + cu_seqlens: Optional[torch.Tensor], + chunk_size: int = 64, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run KDA chunked prefill. + + On the optimized path this replays the preprocessing that FLA's + ``ChunkKDAFunction.forward`` performs when the caller enables + ``use_qk_l2norm_in_kernel``, ``use_beta_sigmoid_in_kernel``, + ``use_gate_in_kernel``, and ``state_v_first``; then dispatches to + the in-tree ``trtllm::kda_prefill`` CuTe DSL op. + + On the FLA path it calls ``fla.ops.kda.chunk_kda`` directly with the + matching flags so the semantics are byte-equivalent. + + State layout contract (both paths): ``initial_state`` is consumed + and ``final_state`` returned in the V-first ``[N, H, V, K]`` layout — + the layout of the executor's ssm pool and of the fused decode + kernel. The in-tree prefill op natively uses the FLA-default K-first + ``[N, H, K, V]`` layout, so the optimized path transposes at both + boundaries (K == V == 128 for Kimi K3, so shapes alone cannot catch + a mix-up — the transpose is semantic). + """ + use_optimized = self.prefill_kernel_path == "optimized" + chunk_indices = None + if use_optimized and cu_seqlens is not None: + from fla.ops.utils.index import prepare_chunk_indices + + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + # The persistent K123 scheduler needs at least 4 total chunks + # (cgs_per_head = NT // 4 cooperative groups per head). The + # eqlen path guarantees this by padding to a 256-token multiple + # inside the op; varlen has no such pad, so small varlen + # batches (short-prompt contexts, NT < 4) launch with a + # zero-size grid -> DSLCudaRuntimeError. Route them to the FLA + # reference path (negligible perf impact at these sizes). The + # check must happen HERE, before the l2norm/beta-sigmoid + # pre-transforms below: the FLA path applies both in-kernel. + if chunk_indices.shape[0] < 4: + use_optimized = False + + if use_optimized: + import torch.nn.functional as F + from fla.modules.l2norm import l2norm_fwd + from fla.ops.common.gate import fused_beta_sigmoid + + q, _ = l2norm_fwd(q) + k, _ = l2norm_fwd(k) + beta = fused_beta_sigmoid(beta, scale=1.0).to(torch.bfloat16) + + real_T = q.shape[1] + if cu_seqlens is not None: + # The op's varlen single-seq path (Phase 2.1) expects the + # caller to zero-pad the packed tensors to a chunk multiple + # (FLA convention) while cu_seqlens keeps the real length; + # the op re-sentinels g's tail itself. Without this the op + # would run the mask-free kernel on a partial final chunk. + # Multi-seq varlen runs the masked path and needs no pad. + if cu_seqlens.shape[0] == 2 and real_T % chunk_size != 0: + pad = chunk_size - real_T % chunk_size + q = F.pad(q, (0, 0, 0, 0, 0, pad)) + k = F.pad(k, (0, 0, 0, 0, 0, pad)) + v = F.pad(v, (0, 0, 0, 0, 0, pad)) + g = F.pad(g, (0, 0, 0, 0, 0, pad)) + beta = F.pad(beta, (0, 0, 0, pad)) + + A_log_kernel = A_log.detach() if A_log is not None else None + dt_bias_kernel = dt_bias.detach() if dt_bias is not None else None + + _load_prefill_module() # registers trtllm::kda_prefill + + if initial_state is not None: + # Pool V-first [N, H, V, K] -> op K-first [N, H, K, V]. + initial_state = initial_state.transpose(-1, -2).contiguous() + + out, final_state = torch.ops.trtllm.kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + use_gate_in_kernel=True, + A_log=A_log_kernel, + dt_bias=dt_bias_kernel, + ) + if out.shape[1] != real_T: + out = out[:, :real_T] + if final_state is not None and final_state.numel() > 0: + # Op K-first [N, H, K, V] -> pool V-first [N, H, V, K]. + # .contiguous() also detaches the result from the op's + # shared per-shape S_out scratch, which the next same-shape + # call overwrites. + final_state = final_state.transpose(-1, -2).contiguous() + return out, final_state + + from fla.ops.kda import chunk_kda + + o, final_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=safe_gate, + lower_bound=lower_bound, + state_v_first=True, + cu_seqlens=cu_seqlens, + ) + return o, final_state + + def decode_kda(self, **kwargs) -> torch.Tensor: + """Run the fused KDA single-token decode kernel. + + This is only defined on the optimized path; the FLA fallback runs + ``fla.ops.kda.fused_recurrent_kda`` directly and does not use this + wrapper (see ``_decode_via_fla`` on the module). + """ + if self.decode_kernel_path != "optimized": + raise RuntimeError( + "decode_kda called on non-optimized path; use FLA path via " + "the module's fallback handling instead." + ) + return _kda_decode.run_kda_decode_fusion_cuda(**kwargs) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py new file mode 100644 index 000000000000..e306feb2781d --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -0,0 +1,723 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KimiKDALinearAttention — Kimi K3 linear-attention module for the PyTorch backend. + +Structural mirror of the HF reference ``KimiDeltaAttention`` in +``modeling_kimi.py``. Same parameter names, same layer shapes, same short +convolution + FLA gating + FusedRMSNormGated output-gate stack. The +delta-rule inner loop is routed through :mod:`_kda_kernels`, which selects +the optimized sm_100 CuTe/Triton chunked prefill and fused CUDA decode +kernels on Blackwell and falls back to the FLA references elsewhere. + +Cache ownership +--------------- +KDA carries three short-convolution states (``conv_state_{q,k,v}``, HF +layout ``[B, D, W]`` bf16) and one delta-rule recurrent state +(``recurrent_state``, layout ``[B, HV, V, K]`` fp32, matching the optimized +kernel's transposed convention). These match the hybrid-cache ownership +pattern used by the mamba modules; the runtime cache-manager plumbing +(``AttentionMetadata`` split, cache indices, spec/verify path) is deferred +to the model-assembly wiring goal. This module exposes parity entry points +that consume and return the state tensors directly so module-level tests +can prove state roundtrip without the runtime plumbing. + +Kernel mutations for negative controls +-------------------------------------- +Two invariants have their own construction switches so parity tests can +prove they are actually being enforced: + +* ``gate_lower_bound_override`` — replace the ``linear_attn_config`` + gate lower bound at forward time. A value that disagrees with the HF + reference must fail parity. +* ``wrong_state_layout`` — permute the recurrent state's V/K axes before + and after the decode kernel call so read/write hit mislabeled slots. + Because K == V the shape check still passes but the numerics break. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from einops import rearrange +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.ops.kda import fused_recurrent_kda +from torch import nn + +from ._kda_kernels import KDAKernelDispatch, is_kda_optimized_supported + + +def _meta_safe_cast_dtype(module, dtype): + """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. + + ``Module.to`` dispatches ``aten._to_copy``, which MetaInitMode rejects + (it would silently fall back to full CPU construction of the model — + ~70 GB of host RAM per rank for Kimi K3). Under meta init the values + are garbage anyway, so a dtype-only re-allocation via ``empty_like`` + (an allowed init op) is equivalent; off meta this matches ``.to``. + """ + import torch as _torch + + def _cast(t): + if not t.is_floating_point(): + return t + if t.is_meta: + return _torch.empty_like(t, dtype=dtype) + return t.to(dtype=dtype) + + module._apply(_cast) + + +class _MetaSafeFusedRMSNormGated(FusedRMSNormGated): + """FusedRMSNormGated whose init survives the model loader's MetaInitMode. + + ``FusedRMSNormGated.reset_parameters`` uses ``nn.init.ones_`` (a plain + ``fill_``), which MetaInitMode rejects on meta tensors and which would + force the whole model construction to fall back to eager CPU init. + ``uniform_(1, 1)`` produces identical values and is on MetaInitMode's + random-init allowlist. + """ + + def reset_parameters(self) -> None: + if self.elementwise_affine: + with torch.no_grad(): + self.weight.uniform_(1.0, 1.0) + + +@dataclass +class KimiKDACachedState: + """Per-layer KDA cache tensors in HF layout. + + ``conv_state_*`` — shape ``[B, D, W]`` bf16 where ``W`` is + ``short_conv_kernel_size`` and the newest processed token sits at + position ``W-1``. ``D`` is ``num_heads * head_dim`` for q/k and + ``num_heads * head_dim`` for v (K3 uses HV == H). + + ``recurrent_state`` — shape ``[B, HV, V, K]`` fp32. This is the + transposed layout the optimized KDA kernels expect, and it is the same + layout HF stores when running with ``transpose_state_layout=True``. + ``None`` fields are treated as zero. + """ + + conv_state_q: Optional[torch.Tensor] + conv_state_k: Optional[torch.Tensor] + conv_state_v: Optional[torch.Tensor] + recurrent_state: Optional[torch.Tensor] + + +class KimiKDAKernelPath: + """Enum-like string tags for the selected KDA kernel path.""" + + OPTIMIZED = "optimized" + FLA = "fla" + + +def _hf_conv_to_kernel_conv( + hf_cache: Optional[torch.Tensor], + b: int, + d: int, + w: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """HF ``[B, D, W]`` conv cache -> optimized kernel's ``[B, D, W-1]``. + + HF stores W positions with the newest processed token last. The + optimized decode kernel's ``cs_*`` argument stores the ``W-1`` + historical positions before the incoming token; drop the oldest column. + """ + if hf_cache is None: + return torch.zeros(b, d, w - 1, device=device, dtype=dtype) + return hf_cache[:, :, 1:].contiguous() + + +def _roll_hf_conv( + prev_hf: Optional[torch.Tensor], + x_new_col: torch.Tensor, + b: int, + d: int, + w: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Roll an HF-layout conv cache by one token. + + HF ``ShortConvolution.step`` does + ``cache.copy_(cache.roll(shifts=-1, dims=-1)); cache[:, :, -1] = x``. + We implement the same semantics via ``torch.cat`` so the update is + independent of the kernel's internal cs handling. + """ + if prev_hf is None: + prev = torch.zeros(b, d, w, device=device, dtype=dtype) + else: + prev = prev_hf.to(dtype=dtype) + return torch.cat([prev[:, :, 1:], x_new_col.to(dtype)], dim=-1).contiguous() + + +class KimiKDALinearAttention(nn.Module): + """Kimi K3 linear-attention module — in-tree production version. + + Parameters + ---------- + hidden_size : int + num_heads : int + head_dim : int + conv_kernel_size : int + use_full_rank_gate : bool + gate_lower_bound : Optional[float] + rms_norm_eps : float + dtype : Optional[torch.dtype] + layer_idx : int + use_optimized_prefill : bool + Enable the optimized prefill path when supported. + use_optimized_decode : bool + Enable the optimized decode path when supported. + gate_lower_bound_override : Optional[float] + Override the ``linear_attn_config`` gate lower bound. Test knob for + the "wrong gate lower bound" mutation control. + wrong_state_layout : bool + Swap the V and K axes of the recurrent state around the decode + kernel call. Test knob for the "wrong state layout" mutation + control on the decode path. + """ + + def __init__( + self, + *, + hidden_size: int, + num_heads: int, + head_dim: int, + conv_kernel_size: int, + use_full_rank_gate: bool, + gate_lower_bound: Optional[float], + rms_norm_eps: float = 1e-5, + dtype: Optional[torch.dtype] = None, + layer_idx: int = 0, + use_optimized_prefill: bool = True, + use_optimized_decode: bool = True, + gate_lower_bound_override: Optional[float] = None, + wrong_state_layout: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = head_dim + self.head_k_dim = head_dim + self.num_k_heads = num_heads + self.conv_size = conv_kernel_size + self.use_full_rank_gate = use_full_rank_gate + self.gate_lower_bound = gate_lower_bound + self.rms_norm_eps = rms_norm_eps + self.layer_idx = layer_idx + self.gate_lower_bound_override = gate_lower_bound_override + self.wrong_state_layout = wrong_state_layout + + projection_k_size = self.head_k_dim * self.num_k_heads + projection_size = self.head_dim * self.num_heads + + self.q_proj = nn.Linear(hidden_size, projection_k_size, bias=False) + self.k_proj = nn.Linear(hidden_size, projection_k_size, bias=False) + self.v_proj = nn.Linear(hidden_size, projection_size, bias=False) + + self.q_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=conv_kernel_size, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=conv_kernel_size, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=projection_size, + kernel_size=conv_kernel_size, + activation="silu", + ) + + self.A_log = nn.Parameter( + torch.log(torch.empty(num_heads, dtype=torch.float32).uniform_(1, 16)) + ) + self.f_a_proj = nn.Linear(hidden_size, head_dim, bias=False) + self.f_b_proj = nn.Linear(head_dim, projection_size, bias=False) + self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32)) + self.b_proj = nn.Linear(hidden_size, num_heads, bias=False) + + if use_full_rank_gate: + self.g_proj = nn.Linear(hidden_size, projection_size, bias=False) + else: + self.g_a_proj = nn.Linear(hidden_size, head_dim, bias=False) + self.g_b_proj = nn.Linear(head_dim, projection_size, bias=False) + + self.o_norm = _MetaSafeFusedRMSNormGated(head_dim, eps=rms_norm_eps, activation="sigmoid") + self.o_proj = nn.Linear(projection_size, hidden_size, bias=False) + + # Installed together by the FP8 weight loader (fused [q | k | v | g] + # decode GEMM). Declared here so the decode path never sees a + # half-installed pair. + self.qkvg_proj: Optional[nn.Module] = None + self.qkvg_split_sizes: Optional[list[int]] = None + + if dtype is not None: + _meta_safe_cast_dtype(self, dtype) + + # The optimized decode/verify kernels are specialized for the Kimi + # K3 shape (K == V == 128). Reduced-dim test configurations must + # fall back to FLA instead of hard-failing inside the kernels. + kernel_shape_ok = self.head_k_dim == 128 and self.head_dim == 128 + self._dispatch = KDAKernelDispatch( + use_optimized_prefill=use_optimized_prefill, + use_optimized_decode=use_optimized_decode and kernel_shape_ok, + use_optimized_verify=kernel_shape_ok, + ) + + # ------------------------------------------------------------------ + # Introspection helpers used by the test smoke. + # ------------------------------------------------------------------ + + @property + def prefill_kernel_path(self) -> str: + return self._dispatch.prefill_kernel_path + + @property + def decode_kernel_path(self) -> str: + return self._dispatch.decode_kernel_path + + @property + def verify_kernel_path(self) -> str: + return self._dispatch.verify_kernel_path + + @property + def sm_100_optimized_supported(self) -> bool: + return is_kda_optimized_supported() + + def prefill_kernel_source(self) -> str: + return self._dispatch.get_prefill_source() + + def decode_kernel_source(self) -> str: + return self._dispatch.get_decode_source() + + def prefill_chunk_kda(self, **kwargs): + """Kernel-level chunked prefill via the dispatch. + + Used by the executor runtime (``KimiKDARuntime``), which owns the + projections, convs, and cache pools itself and only needs the + delta-rule inner loop. States are exchanged in the V-first + ``[N, H, V, K]`` pool layout on both dispatch paths — see + ``KDAKernelDispatch.prefill_chunk_kda``. + """ + return self._dispatch.prefill_chunk_kda(**kwargs) + + # ------------------------------------------------------------------ + # Prefill entry (Goal 2.1 pass path). + # ------------------------------------------------------------------ + + def forward_prefill( + self, + hidden_states: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Prefill forward matching HF ``KimiDeltaAttention.forward`` in chunk mode. + + Parameters + ---------- + hidden_states : ``(B, T, hidden_size)`` for equal-length prefill or + ``(1, sum(seq_lens), hidden_size)`` when ``cu_seqlens`` is given. + cu_seqlens : optional cumulative sequence lengths for varlen inputs. + + Returns + ------- + ``(B, T, hidden_size)`` output tensor (equal-length case) or + ``(1, sum(seq_lens), hidden_size)`` (varlen case). + """ + if cu_seqlens is not None: + cu_seqlens = cu_seqlens.to(device=hidden_states.device, dtype=torch.long) + + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + + q, _ = self.q_conv1d( + x=q_proj_states, + cache=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + k, _ = self.k_conv1d( + x=k_proj_states, + cache=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + v, _ = self.v_conv1d( + x=v_proj_states, + cache=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + + g = self.f_b_proj(self.f_a_proj(hidden_states)) + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) + beta = self.b_proj(hidden_states).float() + + q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) + k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + lower_bound = ( + self.gate_lower_bound_override + if self.gate_lower_bound_override is not None + else self.gate_lower_bound + ) + safe_gate = lower_bound is not None + scale = self.head_k_dim**-0.5 + + o, _final_state = self._dispatch.prefill_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + scale=scale, + initial_state=None, + safe_gate=safe_gate, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_size=64, + ) + + if self.use_full_rank_gate: + g_out = self.g_proj(hidden_states) + else: + g_out = self.g_b_proj(self.g_a_proj(hidden_states)) + g_out = rearrange(g_out, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, g_out) + + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + return o + + # ------------------------------------------------------------------ + # Decode entry (Goal 2.1 pass path). + # ------------------------------------------------------------------ + + def forward_decode( + self, + hidden_states: torch.Tensor, + cache: Optional[KimiKDACachedState] = None, + ssm_state_indices: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, KimiKDACachedState]: + """T=1 cached-decode forward. Returns ``(o, new_cache)``. + + ``hidden_states`` shape ``(B, 1, hidden_size)``. Cache is + ``KimiKDACachedState`` in HF layout; ``None`` fields become zero + tensors. + """ + b, q_len, _ = hidden_states.shape + assert q_len == 1, f"KimiKDALinearAttention.forward_decode expects T=1, got T={q_len}" + + if self._dispatch.decode_kernel_path == KimiKDAKernelPath.OPTIMIZED: + return self._decode_via_optimized(hidden_states, cache, b, ssm_state_indices) + if ssm_state_indices is not None: + raise ValueError("ssm_state_indices requires the optimized KDA decode kernel") + return self._decode_via_fla(hidden_states, cache, b) + + # ------------------------------------------------------------------ + # Internals — optimized decode dispatch. + # ------------------------------------------------------------------ + + def _decode_via_optimized( + self, + hidden_states: torch.Tensor, + cache: Optional[KimiKDACachedState], + b: int, + ssm_state_indices: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, KimiKDACachedState]: + dev = hidden_states.device + H = self.num_heads + HV = self.num_heads + K_dim = self.head_dim + V_dim = self.head_dim + W = self.conv_size + + projection_size = H * K_dim + projection_v_size = HV * V_dim + + # q/k/v and the full-rank output gate all read this same normed hidden. + # When their weights are read at FP8 block-scale (Blackwell decode), the + # loader fuses them into one ``qkvg_proj`` GEMM: one activation quant and + # one GEMM launch replace four, which is what the launch-bound + # generation step needs. The split is output-identical to the per + # projection GEMMs (same activation, same weight slices). The forget + # gate (f_a/f_b), beta and low-rank output gate stay BF16, so they keep + # their own calls. + fused_qkvg = self.qkvg_proj if self.qkvg_split_sizes is not None else None + if fused_qkvg is not None: + parts = fused_qkvg(hidden_states).split(self.qkvg_split_sizes, dim=-1) + q_proj_states, k_proj_states, v_proj_states = parts[0], parts[1], parts[2] + onorm_g_hidden = ( + parts[3] if self.use_full_rank_gate else self.g_b_proj(self.g_a_proj(hidden_states)) + ) + else: + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + onorm_g_hidden = ( + self.g_proj(hidden_states) + if self.use_full_rank_gate + else self.g_b_proj(self.g_a_proj(hidden_states)) + ) + + g_hidden = self.f_b_proj(self.f_a_proj(hidden_states)) + + beta_hidden = self.b_proj(hidden_states).float() + + def _kernel_input(proj: torch.Tensor, h: int, d: int) -> torch.Tensor: + x = rearrange(proj, "b t (h d) -> t b h d", h=h, d=d) + return x.to(dtype=torch.bfloat16).contiguous() + + x_q_full = _kernel_input(q_proj_states, H, K_dim) + x_k_full = _kernel_input(k_proj_states, H, K_dim) + x_v_full = _kernel_input(v_proj_states, HV, V_dim) + g_full = _kernel_input(g_hidden, H, K_dim) + onorm_g_full = _kernel_input(onorm_g_hidden, HV, V_dim) + beta_full = rearrange(beta_hidden, "b t h -> t b h").to(torch.bfloat16).contiguous() + + w_q_t_full = ( + self.q_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + w_k_t_full = ( + self.k_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + w_v_t_full = ( + self.v_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + + if cache is not None and cache.conv_state_q is not None: + hf_cs_q_pre = cache.conv_state_q.to(torch.bfloat16) + else: + hf_cs_q_pre = torch.zeros(b, projection_size, W, device=dev, dtype=torch.bfloat16) + if cache is not None and cache.conv_state_k is not None: + hf_cs_k_pre = cache.conv_state_k.to(torch.bfloat16) + else: + hf_cs_k_pre = torch.zeros(b, projection_size, W, device=dev, dtype=torch.bfloat16) + if cache is not None and cache.conv_state_v is not None: + hf_cs_v_pre = cache.conv_state_v.to(torch.bfloat16) + else: + hf_cs_v_pre = torch.zeros(b, projection_v_size, W, device=dev, dtype=torch.bfloat16) + + cs_q_full = _hf_conv_to_kernel_conv(hf_cs_q_pre, b, projection_size, W, dev, torch.bfloat16) + cs_k_full = _hf_conv_to_kernel_conv(hf_cs_k_pre, b, projection_size, W, dev, torch.bfloat16) + cs_v_full = _hf_conv_to_kernel_conv( + hf_cs_v_pre, b, projection_v_size, W, dev, torch.bfloat16 + ) + + x_q_col = q_proj_states.transpose(1, 2).to(torch.bfloat16) + x_k_col = k_proj_states.transpose(1, 2).to(torch.bfloat16) + x_v_col = v_proj_states.transpose(1, 2).to(torch.bfloat16) + new_hf_cs_q = _roll_hf_conv( + hf_cs_q_pre, x_q_col, b, projection_size, W, dev, torch.bfloat16 + ) + new_hf_cs_k = _roll_hf_conv( + hf_cs_k_pre, x_k_col, b, projection_size, W, dev, torch.bfloat16 + ) + new_hf_cs_v = _roll_hf_conv( + hf_cs_v_pre, x_v_col, b, projection_v_size, W, dev, torch.bfloat16 + ) + + if cache is not None and cache.recurrent_state is not None: + if ssm_state_indices is not None: + if self.wrong_state_layout: + raise ValueError("ssm_state_indices is incompatible with wrong_state_layout") + state_full = cache.recurrent_state + else: + state_full = cache.recurrent_state.to(dtype=torch.float32).contiguous() + else: + if ssm_state_indices is not None: + raise ValueError("ssm_state_indices requires a recurrent state pool") + state_full = torch.zeros(b, HV, V_dim, K_dim, device=dev, dtype=torch.float32) + + # The decode op requires fp32 A_log/dt_bias even in a bf16-cast module. + A_log_full = self.A_log.detach().float().contiguous() + dt_bias_full = self.dt_bias.detach().float().contiguous() + onorm_weight_full = self.o_norm.weight.detach().to(torch.float32).contiguous() + lower_bound = ( + self.gate_lower_bound_override + if self.gate_lower_bound_override is not None + else self.gate_lower_bound + ) + + kernel_state = ( + state_full.transpose(-1, -2).contiguous() if self.wrong_state_layout else state_full + ) + o_bfhvk = self._dispatch.decode_kda( + x_q=x_q_full, + x_k=x_k_full, + x_v=x_v_full, + w_q_t=w_q_t_full, + w_k_t=w_k_t_full, + w_v_t=w_v_t_full, + bias_q=None, + bias_k=None, + bias_v=None, + cs_q=cs_q_full, + cs_k=cs_k_full, + cs_v=cs_v_full, + A_log=A_log_full, + g=g_full, + dt_bias=dt_bias_full, + beta=beta_full, + state=kernel_state, + onorm_g=onorm_g_full, + onorm_weight=onorm_weight_full, + out=None, + ssm_state_indices=ssm_state_indices, + cu_seqlens=None, + scale=K_dim**-0.5, + onorm_eps=self.o_norm.eps, + lower_bound=lower_bound, + use_beta_sigmoid_in_kernel=True, + verbose=False, + update_conv_cache=False, + ) + state_full = ( + kernel_state.transpose(-1, -2).contiguous() if self.wrong_state_layout else kernel_state + ) + + o_flat = rearrange(o_bfhvk, "b t h d -> b t (h d)") + o = self.o_proj(o_flat) + + new_cache = KimiKDACachedState( + conv_state_q=new_hf_cs_q, + conv_state_k=new_hf_cs_k, + conv_state_v=new_hf_cs_v, + recurrent_state=state_full, + ) + return o, new_cache + + # ------------------------------------------------------------------ + # Internals — FLA fallback decode (non-sm_100 path). + # ------------------------------------------------------------------ + + def _decode_via_fla( + self, + hidden_states: torch.Tensor, + cache: Optional[KimiKDACachedState], + b: int, + ) -> Tuple[torch.Tensor, KimiKDACachedState]: + """FLA ``fused_recurrent_kda`` decode path — used when sm_100 is unavailable. + + Matches HF ``KimiDeltaAttention`` in ``fused_recurrent`` mode: uses + the ``ShortConvolution.step`` semantics and dispatches the delta + update to ``fla.ops.kda.fused_recurrent_kda``. + """ + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + + conv_q_in = cache.conv_state_q if cache is not None else None + conv_k_in = cache.conv_state_k if cache is not None else None + conv_v_in = cache.conv_state_v if cache is not None else None + recurrent_in = cache.recurrent_state if cache is not None else None + + q, new_conv_q = self.q_conv1d(x=q_proj_states, cache=conv_q_in, output_final_state=True) + k, new_conv_k = self.k_conv1d(x=k_proj_states, cache=conv_k_in, output_final_state=True) + v, new_conv_v = self.v_conv1d(x=v_proj_states, cache=conv_v_in, output_final_state=True) + + g_hidden = self.f_b_proj(self.f_a_proj(hidden_states)) + if self.use_full_rank_gate: + onorm_g_hidden = self.g_proj(hidden_states) + else: + onorm_g_hidden = self.g_b_proj(self.g_a_proj(hidden_states)) + beta = self.b_proj(hidden_states).float() + + g = rearrange(g_hidden, "... (h d) -> ... h d", d=self.head_dim) + q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) + k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + lower_bound = ( + self.gate_lower_bound_override + if self.gate_lower_bound_override is not None + else self.gate_lower_bound + ) + + o, new_recurrent = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_in, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=lower_bound, + state_v_first=True, + ) + + onorm_g = rearrange(onorm_g_hidden, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, onorm_g) + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + + new_cache = KimiKDACachedState( + conv_state_q=new_conv_q, + conv_state_k=new_conv_k, + conv_state_v=new_conv_v, + recurrent_state=new_recurrent, + ) + return o, new_cache + + # ------------------------------------------------------------------ + # Weight helper for random-weight parity tests. + # ------------------------------------------------------------------ + + def copy_weights_from(self, source: nn.Module) -> "dict[str, Tuple[Tuple[int, ...], str]]": + """Copy every named parameter/buffer from ``source`` into ``self``. + + Because ``KimiKDALinearAttention`` mirrors the HF reference's + parameter names 1:1, the mapping is identity: every source name is + assigned to the identically named target. Shape mismatches raise + loudly. Returns a ``{name: (shape, dtype)}`` provenance dict. + """ + src: dict[str, torch.Tensor] = {} + for name, p in source.named_parameters(recurse=True): + src[name] = p.data + for name, buf in source.named_buffers(recurse=True): + src[name] = buf + + dst: dict[str, torch.Tensor] = {} + for name, p in self.named_parameters(recurse=True): + dst[name] = p.data + for name, buf in self.named_buffers(recurse=True): + dst[name] = buf + + missing_on_dst = sorted(set(src) - set(dst)) + missing_on_src = sorted(set(dst) - set(src)) + if missing_on_dst: + raise KeyError( + f"copy_weights_from: source params missing on target: {missing_on_dst[:5]}" + ) + if missing_on_src: + raise KeyError( + f"copy_weights_from: target params missing on source: {missing_on_src[:5]}" + ) + + provenance: "dict[str, Tuple[Tuple[int, ...], str]]" = {} + for name, srct in src.items(): + dstt = dst[name] + if srct.shape != dstt.shape: + raise ValueError( + f"shape mismatch for {name}: source {tuple(srct.shape)} " + f"vs target {tuple(dstt.shape)}" + ) + dstt.copy_(srct.to(dtype=dstt.dtype, device=dstt.device)) + provenance[name] = (tuple(srct.shape), str(srct.dtype)) + return provenance diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 5b97e32a6ee3..26b57e1e486a 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -261,6 +261,14 @@ def __init__(self, max_batch_size: int, chunk_size: int): self.state_indices = torch.zeros(max_batch_size, dtype=torch.int32, device="cuda") + # int64 mirror of state_indices, refreshed once per prepare() so + # per-layer consumers that need long indices (index_select / + # index_copy_) do not each launch an int32->int64 cast kernel + # inside the decode CUDA graph (69 KDA layers x ~1.7us for Kimi K3). + self._state_indices_long = torch.zeros(max_batch_size, + dtype=torch.long, + device="cuda") + self.state_indices_long = self._state_indices_long[:0] # Stable data_ptr() of the CUDA tensor we alias (if any) — used to # detect cache-manager buffer reallocation that would silently break # CUDA graph replays. @@ -398,6 +406,12 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices[:batch_size].copy_( self.state_indices_cpu[:batch_size], non_blocking=True) + # Refresh the int64 mirror once per step (outside the decode graph) + # so layers can index pools without a per-layer cast kernel. + self._state_indices_long[:batch_size].copy_( + self.state_indices[:batch_size]) + self.state_indices_long = self._state_indices_long[:batch_size] + self._prepare_replay_work_items(kv_cache_manager, batch_size, num_contexts) diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index d454f833726b..5b36641f3d0e 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -433,6 +433,8 @@ def __init__( reduce_output: bool = True, num_groups: int = 1, o_lora_rank: int = 1024, + fuse_qkv_a_proj: bool = True, + rms_norm_eps: Optional[float] = None, ): """ Initialize the MLA module. @@ -459,6 +461,12 @@ def __init__( config (ModelConfig): The model configuration. num_groups (int): The number of groups. o_lora_rank (int): The dimension of the compressed output. + fuse_qkv_a_proj (bool): Whether q_a and kv_a share one fused + projection. Set to ``False`` for checkpoints that store a + separate ``q_a_proj``. + rms_norm_eps (Optional[float]): Override the RMSNorm epsilon from + the pretrained config. If neither source provides a value + (e.g. config.pretrained_config is None), falls back to 1e-6. """ super().__init__() self.layer_idx = layer_idx @@ -482,6 +490,7 @@ def __init__( self.dense_bias = dense_bias self.num_groups = num_groups self.o_lora_rank = o_lora_rank + self.fuse_qkv_a_proj = fuse_qkv_a_proj if dense_bias is None: self.dense_bias = bias @@ -523,6 +532,15 @@ def __init__( sparse_algorithm = getattr(sparse_params, "algorithm", None) self.is_dsa = sparse_algorithm == "dsa" self.is_deepseek_v4 = sparse_algorithm == "deepseek_v4" + if (self.is_dsa or self.is_deepseek_v4) and not fuse_qkv_a_proj: + # forward_dsa_proj assumes the fused [q_a | kv_a | k_pe] + # projection layout; the separate q_a_proj layout (Kimi K3 MLA) + # is not wired into the DSA / DeepSeek-V4 sparse paths. + raise NotImplementedError( + "DSA and DeepSeek-V4 sparse attention require " + "fuse_qkv_a_proj=True; the separate q_a_proj layout is not " + "supported with them." + ) self._disable_dsv4_epilogue_fusion = self.is_deepseek_v4 and _is_env_truthy( "TRTLLM_DSV4_DISABLE_FMHA_EPILOGUE_FUSION" ) @@ -584,7 +602,13 @@ def __init__( ) self.n_local_groups = self.num_groups // tp_size - rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps", 1e-6) + if rms_norm_eps is None: + rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps", None) + if rms_norm_eps is None: + # No explicit value and pretrained_config is None or lacks + # rms_norm_eps (e.g. unit tests constructing MLA directly): + # keep the historical default. + rms_norm_eps = 1e-6 quant_config = config.get_quant_config() self.quant_config = quant_config @@ -594,9 +618,12 @@ def __init__( self.use_cute_dsl_bf16_gemm = config.use_cute_dsl_bf16_gemm if not self.is_lite: + kv_a_out_features = self.kv_lora_rank + self.qk_rope_head_dim + if self.fuse_qkv_a_proj: + kv_a_out_features += self.q_lora_rank self.kv_a_proj_with_mqa = Linear( hidden_size, - self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim, + kv_a_out_features, bias=bias, dtype=dtype, quant_config=quant_config, @@ -607,6 +634,20 @@ def __init__( use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, ) + if not self.fuse_qkv_a_proj: + self.q_a_proj = Linear( + hidden_size, + self.q_lora_rank, + bias=bias, + dtype=dtype, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + use_custom_cublas_mm=True, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, + use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, + ) + self.q_a_layernorm = RMSNorm( hidden_size=self.q_lora_rank, eps=rms_norm_eps, dtype=dtype ) @@ -1427,9 +1468,15 @@ def forward_impl( compressed_kv = self.kv_a_layernorm(compressed_kv) q = hidden_states else: - q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 - ) + if self.fuse_qkv_a_proj: + q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1 + ) + else: + q = self.q_a_proj(hidden_states) + compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( + [self.kv_lora_rank, self.qk_rope_head_dim], -1 + ) q, compressed_kv = maybe_execute_in_parallel( lambda: self._q_a_layernorm_maybe_fused(q), diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index be555fa55c42..9da5db99c10d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -48,8 +48,8 @@ get_spec_decoder, should_use_separate_draft_kv_cache) from ..utils import is_gdn_replay_enabled from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params, - is_gemma4_hybrid, is_hybrid_linear, is_mla, - is_nemotron_hybrid, is_qwen3_hybrid) + is_gemma4_hybrid, is_hybrid_linear, is_kimi_linear, + is_mla, is_nemotron_hybrid, is_qwen3_hybrid) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -155,6 +155,31 @@ def get_kv_cache_manager_cls( raise ValueError("Mamba additional snapshot offsets require " "use_kv_cache_manager_v2=True; V1 supports only " "periodic_snapshot_interval.") + + # Kimi K3 (KDA + MLA hybrid): block reuse uses the unified C++ pool + # (CppMambaHybridCacheManager) like the other hybrid linear models — + # per-block KDA state snapshots every mamba_state_cache_interval + # tokens with FORCE_CHUNK context chunking. Without block reuse the + # Mixed manager (separate KV / recurrent-state pools) stays the + # default. SA speculative decoding is validated on the Mixed + # manager's SpeculativeState scratch path only; reuse + SA is + # unvalidated. + if is_kimi_linear(config) and not use_v2: + if is_disagg: + # Fail fast instead of bypassing the disagg transceiver + # validation below with an unvalidated route. + raise NotImplementedError( + "Disaggregated serving is not supported for Kimi K3 yet " + "(TRTLLM-14815).") + if kv_cache_config.enable_block_reuse: + logger.info( + "Using CppMambaHybridCacheManager for Kimi K3 hybrid " + "model (block reuse enabled)") + return CppMambaHybridCacheManager + logger.info( + "Using MixedMambaHybridCacheManager for Kimi K3 hybrid model") + return MixedMambaHybridCacheManager + # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. if is_disagg: @@ -2147,7 +2172,88 @@ def _create_kv_cache_manager( if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg - if is_mla(config): + if is_kimi_linear(config): + # Kimi K3 hybrid: KDA (Kimi Delta Attention) recurrent/conv states on + # the mamba side of the hybrid manager, absorbed-MQA MLA latent cache + # (num_kv_heads=1, head_dim = kv_lora_rank + qk_rope_head_dim, + # SELFKONLY) on the paged-KV side. Must come before the is_mla(...) + # route: the kimi_linear config carries MLA fields, but only 24 of + # its 93 layers are MLA. + if max_beam_width > 1: + raise ValueError( + "MambaHybridCacheManager + beam search is not supported yet.") + if not estimating_kv_cache and kv_connector_manager is not None: + raise NotImplementedError( + "Connector manager is not supported for MambaHybridCacheManager." + ) + mamba_params = extract_mamba_kv_cache_params( + config, + spec_config=spec_config, + quant_config=quant_config, + ) + mamba_layer_mask, full_attention_layer_mask = ( + _get_mamba_cache_layer_masks( + mamba_params, + mapping, + spec_config, + is_draft, + )) + num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 + else mamba_params.num_mamba_layers) + # Kimi K3 KDA state sharding follows the attention-family TP + # semantics (Qwen3-Next pattern): replicated under attention-DP, + # head-sharded across tp_size otherwise. That is exactly the cache + # manager's own internal gate (`tp_size = 1 if enable_attention_dp + # else tp_size`, then num_heads / n_groups / conv_dim divide by + # it), so the params pass through unscaled. + # KDA fused multi-token verify (trtllm::kda_mtp_decode): when the + # kernel can run here, allocate the per-slot replay caches instead + # of the legacy per-step intermediate verification buffers. The + # kernel replays accepted drafts from these caches and commits + # states in place, replacing the intermediate-buffer + promotion + # flow for KDA layers. + kimi_extra_kwargs = {} + if spec_config is not None and issubclass(kv_cache_manager_cls, + MixedMambaHybridCacheManager): + from ..modules.kimi_kda._kda_kernels import \ + is_kda_mtp_verify_available + if is_kda_mtp_verify_available(): + kimi_extra_kwargs["kda_replay_num_spec"] = ( + spec_config.tokens_per_gen_step - 1) + kv_cache_manager = kv_cache_manager_cls( + # mamba (KDA) cache parameters + mamba_params.state_size, + mamba_params.conv_kernel, + mamba_params.num_heads, + mamba_params.n_groups, + mamba_params.head_dim, + num_mamba_layers, + mamba_layer_mask, + mamba_params.dtype, + mamba_params.mamba_ssm_cache_dtype, + # kv cache parameters (MLA latent cache) + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, + num_kv_heads=1, + head_dim=config.kv_lora_rank + config.qk_rope_head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + is_draft=is_draft, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + spec_config=spec_config, + is_estimating_kv_cache=estimating_kv_cache, + execution_stream=execution_stream, + # Reuse the qwen3_next [Q | K | V] conv-state section layout; + # all three KDA sections have identical width. + model_type="qwen3_next", + **kimi_extra_kwargs, + **manager_extra_kwargs, + ) + elif is_mla(config): kv_cache_manager = kv_cache_manager_cls( kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 16be530f7c44..383bc3834e5e 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -20,7 +20,68 @@ def is_gemma4_hybrid(config): def is_hybrid_linear(config): - return is_nemotron_hybrid(config) or is_qwen3_hybrid(config) + return is_nemotron_hybrid(config) or is_qwen3_hybrid(config) or \ + is_kimi_linear(config) + + +def is_kimi_linear(config): + """True for Kimi K3 ("kimi_linear") hybrid KDA + MLA text models. + + Handles both the flattened text config (model_type "kimi_linear") and the + composite VLM config (model_type "kimi_k3" with a nested text_config). + """ + model_type = getattr(config, "model_type", None) + if model_type == "kimi_linear": + return getattr(config, "linear_attn_config", None) is not None + if model_type == "kimi_k3": + text_config = getattr(config, "text_config", None) + return text_config is not None and is_kimi_linear(text_config) + return False + + +def unwrap_kimi_text_config(config): + """Return the flattened Kimi text config. + + ``is_kimi_linear`` accepts both the flattened text config and the + composite "kimi_k3" config with a nested ``text_config``; consumers read + text-level fields (``linear_attn_config``, ``num_hidden_layers``, ...), + so they must unwrap the composite form first. + """ + if getattr(config, "model_type", None) == "kimi_k3": + text_config = getattr(config, "text_config", None) + if text_config is not None: + return text_config + return config + + +def get_kimi_linear_layer_masks(config): + """Return (full_attention_layer_mask, kda_layer_mask) for Kimi K3. + + The config's ``linear_attn_config`` carries 1-indexed ``kda_layers`` and + ``full_attn_layers`` lists; every decoder layer must be exactly one of + the two. + """ + config = unwrap_kimi_text_config(config) + lin = config.linear_attn_config + kda_layers = set(lin["kda_layers"]) + full_attn_layers = set(lin["full_attn_layers"]) + full_mask, kda_mask = [], [] + for layer_idx in range(config.num_hidden_layers): + is_kda = (layer_idx + 1) in kda_layers + is_full = (layer_idx + 1) in full_attn_layers + if is_kda == is_full: + raise ValueError( + f"Kimi K3 layer {layer_idx} (1-indexed {layer_idx + 1}) must " + f"be exactly one of KDA / full attention; got is_kda={is_kda} " + f"is_full={is_full}") + kda_mask.append(is_kda) + full_mask.append(is_full) + return full_mask, kda_mask + + +def get_kimi_linear_num_attention_layers(config): + full_mask, _ = get_kimi_linear_layer_masks(config) + return sum(full_mask) def _coerce_torch_dtype(dtype): @@ -260,7 +321,8 @@ def extract_mamba_kv_cache_params( ) -> MambaKVCacheParams: """Build the mamba-related inputs for kv_cache_manager_cls. - Supports Nemotron-hybrid and Qwen3-hybrid (Qwen3-Next + Qwen3.5). + Supports Nemotron-hybrid, Qwen3-hybrid (Qwen3-Next + Qwen3.5) and + Kimi K3 (kimi_linear). Args: config: HuggingFace model config of a hybrid Mamba model. @@ -288,6 +350,23 @@ def extract_mamba_kv_cache_params( n_groups = config.linear_num_key_heads head_dim = config.linear_value_head_dim target_full_attn_mask, mamba_mask = get_qwen3_hybrid_layer_masks(config) + elif is_kimi_linear(config): + # Kimi K3 KDA (Kimi Delta Attention) state, mapped onto the Mamba + # cache-manager parametrization (see PythonMambaCacheManager): + # conv_dim = head_dim*num_heads + 2*n_groups*state_size + # = 3 * num_heads * head_dim -> [q | k | v] short-conv + # ssm state shape = [num_heads, head_dim, state_size] + # = [H, V, K] fp32 delta-rule recurrent state. + # conv_kernel is set to short_conv_kernel_size + 1 so the pool's + # (conv_kernel - 1) columns hold the FULL FLA ShortConvolution cache + # window of `short_conv_kernel_size` columns. + lin = unwrap_kimi_text_config(config).linear_attn_config + state_size = lin["head_dim"] + conv_kernel = lin["short_conv_kernel_size"] + 1 + num_heads = lin["num_heads"] + n_groups = lin["num_heads"] + head_dim = lin["head_dim"] + target_full_attn_mask, mamba_mask = get_kimi_linear_layer_masks(config) else: raise ValueError( f"{type(config).__name__} is not a supported hybrid Mamba config") @@ -307,6 +386,15 @@ def extract_mamba_kv_cache_params( mamba_ssm_cache_dtype = (resolve_ssm_cache_dtype(config) or resolve_hf_torch_dtype(config) or torch.bfloat16) + if is_kimi_linear(config) and mamba_ssm_cache_dtype != torch.float32: + # The KDA delta-rule recurrent state must be kept in fp32 for + # numerical parity with the HF reference (fla chunk/fused_recurrent + # KDA kernels carry the state in fp32). + logger.info( + f"Kimi K3: overriding mamba_ssm_cache_dtype " + f"{mamba_ssm_cache_dtype} -> torch.float32 (KDA recurrent state " + "must be fp32)") + mamba_ssm_cache_dtype = torch.float32 return MambaKVCacheParams( state_size=state_size, @@ -509,6 +597,15 @@ def load_pretrained_config(model_name_or_path: str, model_name_or_path, **kwargs) _normalize_qwen35_vl_config(model_config, inner_arch="Qwen3_5ForCausalLM") + elif model_type in ("kimi_k3", "kimi_linear"): + # Kimi K3: the checkpoint ships a composite VLM config + # (model_type "kimi_k3" with text/vision sub-configs). TRT-LLM runs + # the text model only, so flatten to the in-tree KimiLinearConfig + # (this also avoids trust_remote_code for the config). + from tensorrt_llm._torch.configs import KimiLinearConfig + text_dict = dict(config_dict.get("text_config") or config_dict) + model_config = KimiLinearConfig.from_dict(text_dict) + model_config.architectures = ["KimiLinearForCausalLM"] elif model_type in _CONFIG_REGISTRY: config_class = _CONFIG_REGISTRY[model_type] model_config = config_class.from_pretrained(model_name_or_path, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 7000b2bdaf0f..026522eaf972 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -343,19 +343,44 @@ def at_layer_idx(self, layer: int): class SpeculativeState(State): """Speculative state with intermediate states for draft tokens. - Supports two SSM update paths (only one set of tensors is allocated): + Supports three SSM update paths (only one set of tensors is + allocated): - Legacy: caches full intermediate SSM states (intermediate_ssm) - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) + - KDA replay: per-slot draft-token caches consumed by the fused + ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted + drafts and commits states in place (kda_conv_*, kda_*_cache) """ _SHARED_FIELDS = frozenset({ "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" }) - intermediate_conv_window: torch.Tensor # always allocated + # Allocated for the legacy and Mamba2-replay paths; None for the + # KDA replay path (the kernel commits conv windows in place). + intermediate_conv_window: torch.Tensor | None = None # Legacy path: full intermediate SSM states at each step intermediate_ssm: torch.Tensor | None = None + # KDA replay path (fused multi-token verify, kimi_linear). + # Pool invariant under this path: `temporal` holds the state after + # the LAST GOLDEN token; the accepted drafts recorded in + # prev_num_accepted_tokens are pending in these caches and are + # replayed by the kernel at the start of the next verify round. + # Extended conv caches [layers, slots, dim, (W-1) + num_spec] fp32, + # dim-contiguous (stride(dim) == 1): columns [0, W-1) are the + # committed raw-input window, tail columns the pending drafts' raw + # inputs. + kda_conv_q: torch.Tensor | None = None + kda_conv_k: torch.Tensor | None = None + kda_conv_v: torch.Tensor | None = None + # Post-processed per-draft quantities for replay: + # [layers, slots, num_spec, 3, H*K] (q/k/gate), [.., num_spec, H*V], + # [.., num_spec, H] — all fp32. + kda_qkg_cache: torch.Tensor | None = None + kda_v_cache: torch.Tensor | None = None + kda_beta_cache: torch.Tensor | None = None + # Replay path: compact double-buffered cache # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). # 0 means temporal saved state is actually the last state, not two back. @@ -389,12 +414,27 @@ def __init__( model_type: str = "nemotron_hybrid", use_replay_state_update: bool = False, mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, ) -> None: self.mamba_ssm_cache_dtype = ssm_cache_dtype self.speculative_num_draft_tokens = speculative_num_draft_tokens self.spec_state_size = spec_state_size self._use_replay_state_update = use_replay_state_update + # KDA replay path (kimi_linear fused multi-token verify). Mutually + # exclusive with use_replay_state_update; requires speculative mode. + self._kda_replay_num_spec = kda_replay_num_spec + self._use_kda_replay_update = kda_replay_num_spec is not None + if self._use_kda_replay_update: + assert not use_replay_state_update, ( + "kda_replay_num_spec and use_replay_state_update are " + "mutually exclusive") + assert speculative_num_draft_tokens is not None, ( + "KDA replay caches require speculative decoding") + assert kda_replay_num_spec == speculative_num_draft_tokens, ( + f"KDA replay cache width ({kda_replay_num_spec}) must match " + f"the draft length ({speculative_num_draft_tokens}): the " + "fused verify kernel is compiled with a static NUM_SPEC") self.replay_history_size: Optional[int] = None self.replay_step_width: Optional[int] = None # When True, allocate the per-slot Philox seed buffer even outside @@ -483,13 +523,17 @@ def __init__( T = speculative_num_draft_tokens + 1 self.replay_step_width = T - # Conv intermediate cache — same for both paths - intermediate_conv_window_cache = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - conv_state_shape, - dtype=dtype, - device=device, - ) + # Conv intermediate cache — legacy and Mamba2-replay paths only. + # The KDA replay kernel commits conv windows in place, so the + # per-step window scratch is not needed. + intermediate_conv_window_cache = None + if not self._use_kda_replay_update: + intermediate_conv_window_cache = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + + conv_state_shape, + dtype=dtype, + device=device, + ) # SSM speculative cache — path-specific tensors spec_kwargs = {} @@ -497,7 +541,63 @@ def __init__( # so the MTP path can still read it via layer_cache. if self._mamba_ssm_rand_seed is not None: spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed - if self._use_replay_state_update: + if self._use_kda_replay_update: + # KDA replay caches for the fused multi-token verify kernel + # (trtllm::kda_mtp_decode). Per cache slot (persistent + # across rounds — the next round replays from them), unlike + # the batch-row-indexed intermediate buffers. + M = self._kda_replay_num_spec + # conv_dim covers the [q | k | v] short-conv sections; the + # kernel consumes them as three dim-contiguous caches. + assert conv_dim % 3 == 0, ( + "KDA replay caches expect the [q | k | v] conv-state " + "sectioning (3 equal sections)") + section_dim = conv_dim // 3 + # d_conv is short_conv_kernel_size + 1 for kimi_linear (the + # pool trick that stores the full FLA window); the kernel's + # conv width is the FLA window size. + w_kernel = d_conv - 1 + extended_s = w_kernel - 1 + M + + def _dim_contiguous_conv_cache(): + return torch.zeros(num_local_layers, + max_batch_size, + extended_s, + section_dim, + dtype=torch.float32, + device=device).transpose(-1, -2) + + spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device) + spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_qkg_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + 3, + section_dim, + dtype=torch.float32, + device=device) + spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.float32, + device=device) + spec_kwargs['kda_beta_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + nheads, + dtype=torch.float32, + device=device) + ssm_spec_cache = [ + spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], + spec_kwargs['kda_conv_v'], spec_kwargs['kda_qkg_cache'], + spec_kwargs['kda_v_cache'], spec_kwargs['kda_beta_cache'] + ] + spec_path_label = "kda-replay" + elif self._use_replay_state_update: assert n_groups % tp_size == 0, \ "replay state update requires n_groups divisible by tp_size" n_groups_per_rank = n_groups // tp_size @@ -659,6 +759,10 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + # Fresh request: no drafts pending in the replay caches. + self.mamba_cache.prev_num_accepted_tokens[block] = 0 if self._mamba_ssm_rand_seed is not None: # Deterministic per-slot rotation on fresh assignment. # `block` is pulled from mamba_cache_free_blocks, which @@ -670,6 +774,60 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): self._seed_request_counter, block, self._seed_rank_offset)) + @torch.inference_mode() + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + """Seed the fused-verify KDA replay conv caches from the conv pool. + + On a disaggregated generation server the ctx->gen transfer populates + only the base ``conv`` / ``temporal`` pools (see + ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` + replay caches consumed by ``trtllm::kda_mtp_decode`` are normally + seeded by a *local* prefill/decode via + ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a + previous occupant's window) for a transferred request — corrupting + the recurrent state from the first verify step onward. Call this + after the state transfer completes, before the first generation + forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row + stores the full FLA window (width W); its last ``W - 1`` columns are + the committed window of the replay caches. Pending-draft scratch is + cleared (no drafts are pending for a freshly transferred request). + """ + if not (self._use_kda_replay_update + and isinstance(self.mamba_cache, self.SpeculativeState)): + return + blocks = [ + self.mamba_cache_index[rid] for rid in request_ids + if rid in self.mamba_cache_index + ] + if not blocks: + return + conv = self.mamba_cache.conv # [L, slots, 3D, W] + idx = torch.tensor(sorted(set(blocks)), + dtype=torch.long, + device=conv.device) + d = conv.shape[2] // 3 + committed = conv.shape[3] - 1 # W - 1 + cs = conv.index_select(1, idx) + for cache, section in ( + (self.mamba_cache.kda_conv_q, cs[:, :, :d]), + (self.mamba_cache.kda_conv_k, cs[:, :, d:2 * d]), + (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d:]), + ): + # cache: [L, slots, D, committed + num_spec]; zero the draft + # tail columns and seed the committed window in one copy. + seeded = torch.zeros( + (cache.shape[0], idx.numel()) + cache.shape[2:], + dtype=cache.dtype, + device=cache.device) + seeded[:, :, :, :committed] = section[:, :, :, 1:].to(cache.dtype) + cache.index_copy_(1, idx, seeded) + for buf in (self.mamba_cache.kda_qkg_cache, + self.mamba_cache.kda_v_cache, + self.mamba_cache.kda_beta_cache): + buf.index_fill_(1, idx, 0) + self.mamba_cache.prev_num_accepted_tokens[idx] = 0 + def prepare_resources(self, scheduled_batch: ScheduledRequests): requests = (scheduled_batch.context_requests + scheduled_batch.generation_requests) @@ -708,6 +866,9 @@ def add_dummy_requests(self, request_ids: List[int], **kwargs): and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 continue if self._is_padding_sentinel(r): block = self._padding_slot @@ -723,6 +884,9 @@ def add_dummy_requests(self, request_ids: List[int], **kwargs): and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 def free_resources(self, request: LlmRequest): request_id = request.py_request_id @@ -869,7 +1033,8 @@ def _drop(tensor): self.mamba_cache = self.SpeculativeState( conv=empty, temporal=empty, - intermediate_conv_window=empty, + intermediate_conv_window=_drop( + self.mamba_cache.intermediate_conv_window), intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), prev_num_accepted_tokens=_drop( self.mamba_cache.prev_num_accepted_tokens), @@ -879,6 +1044,12 @@ def _drop(tensor): old_B=_drop(self.mamba_cache.old_B), old_dt=_drop(self.mamba_cache.old_dt), old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), + kda_conv_q=_drop(self.mamba_cache.kda_conv_q), + kda_conv_k=_drop(self.mamba_cache.kda_conv_k), + kda_conv_v=_drop(self.mamba_cache.kda_conv_v), + kda_qkg_cache=_drop(self.mamba_cache.kda_qkg_cache), + kda_v_cache=_drop(self.mamba_cache.kda_v_cache), + kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), ) else: self.mamba_cache = self.State(conv=empty, temporal=empty) @@ -903,6 +1074,21 @@ def update_mamba_states( state_indices_d = state_indices[num_contexts:num_contexts + num_gens] src_state_indices = self.intermediate_state_indices[:num_gens] + if self._use_kda_replay_update: + # KDA replay: the fused verify kernel already committed the SSM + # state and conv windows (in place, after the golden token) and + # cached this round's drafts. All that remains is recording how + # many of those drafts the sampler accepted, so the next round's + # kernel launch replays exactly that prefix. + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + prev = self.mamba_cache.prev_num_accepted_tokens + current = prev[state_indices_d] + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + prev[state_indices_d] = torch.where(is_dummy_request, current, + accepted) + return + if self._use_replay_state_update: is_dummy_request = self._dummy_request_mask[ num_contexts:num_contexts + num_gens] @@ -956,6 +1142,7 @@ def __init__( model_type: str = "nemotron_hybrid", use_replay_state_update: bool = False, mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, ) -> None: max_num_sequences = max_batch_size * mapping.pp_size @@ -976,6 +1163,7 @@ def __init__( model_type=model_type, use_replay_state_update=use_replay_state_update, mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, ) def get_max_resource_count(self) -> int: @@ -1014,6 +1202,10 @@ def mamba_cache_index(self) -> Dict[int, int]: def get_conv_states(self, layer_idx: int) -> torch.Tensor: return self._impl.get_conv_states(layer_idx) + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: return self._impl.get_ssm_states(layer_idx) @@ -1394,6 +1586,7 @@ def __init__( is_draft: bool = False, use_replay_state_update: bool = False, mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, # Per-pool configurations forwarded to the C++ KVCacheManager ctor. # Lets a single manager host pools with mixed shapes (e.g. Gemma4 # hybrid attention). See KVCacheManager.__init__. @@ -1405,6 +1598,15 @@ def __init__( "mamba hybrid cache requires block reuse to be disabled in KV cache config" ) + # Host-drafter spec modes (NGram) have no spec worker to call + # update_mamba_states; this manager promotes accepted states itself + # in update_resources. One-model modes (MTP/Eagle/DFlash) and the + # suffix-automaton worker promote from their spec workers and must + # NOT be promoted twice. + self._promote_states_in_update_resources = ( + spec_config is not None + and getattr(spec_config, "decoding_type", None) == "NGram") + pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) MambaCacheManager.__init__( @@ -1427,6 +1629,7 @@ def __init__( model_type=model_type, use_replay_state_update=use_replay_state_update, mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, ) # initialize kv cache manager @@ -1472,6 +1675,63 @@ def update_resources(self, kv_cache_dtype_byte_size: float = None): KVCacheManager.update_resources(self, scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) + self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) + + def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): + """Commit accepted verification states for host-drafter spec modes. + + One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton + worker call update_mamba_states themselves right after on-device + acceptance. Host-drafter modes (NGram) have no spec worker: + acceptance lands on the requests as + ``py_num_accepted_draft_tokens`` during sampler update, and this + hook — running right after, alongside the KV rewind — promotes the + accepted step's intermediate state into the live pools (or, on the + KDA replay path, records the accepted-draft count for the next + round's replay). Without it, the pools would keep the + pre-verification state and the next forward would resume from a + stale prefix. + """ + if not self._promote_states_in_update_resources: + return + if not self.is_speculative() or attn_metadata is None: + return + gen_requests = scheduled_batch.generation_requests + if not gen_requests: + return + drafted = [ + r for r in gen_requests + if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 + ] + if not drafted: + # Drafter skipped this step (e.g. speculation gated off): the + # forward ran the plain in-place decode path; nothing to promote. + return + assert len(drafted) == len(gen_requests), ( + "mixed drafted/undrafted generation batch is not supported for " + "hybrid state promotion (drafts are padded to the static max)") + device = self._impl.mamba_cache.temporal.device + num_contexts = len(scheduled_batch.context_requests) + num_accepted = torch.tensor( + [0] * num_contexts + + [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], + dtype=torch.int32, + device=device) + # Batch-ordered slots (contexts then gens), matching the ordering + # the forward used for the intermediate scratch buffers. Requests + # that finished this step were already freed by response handling + # (which runs before update_resources) and are gone from the index; + # their rows must stay in place for alignment, so redirect them to + # the reserved padding slot (a harmless scratch write). + slot_index = self.mamba_cache_index + padding_slot = self._impl._padding_slot + state_indices = torch.tensor([ + slot_index.get(r.py_request_id, padding_slot) + for r in scheduled_batch.context_requests + gen_requests + ], + dtype=torch.int32, + device=device) + self.update_mamba_states(attn_metadata, num_accepted, state_indices) @triton.jit diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index eb3d4193bd8c..153fe93e82ae 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -581,7 +581,14 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], ) for pc in self.pool_configurations ] - if kv_cache_type != CacheTypeCpp.SELF: + # Hybrid linear-attention managers intentionally carry two windows + # (the RECURRENT_STATES sentinel window plus the full attention + # window), including with the SELFKONLY cache type (e.g. Kimi K3: + # MLA latent cache + KDA recurrent state). The C++ WindowBlockManager + # treats kSELFKONLY as a self cache (kv_factor=1) with regular + # per-window pools, so the single-window normalization below (meant + # for cross-KV caches) must not fire for them. + if kv_cache_type != CacheTypeCpp.SELF and not self.is_linear_attention: assert len( blocks_per_window ) == 1, "Only one window size is supported for non-self KV cache" diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 641b948a47b0..2a754d020c7b 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -69,6 +69,10 @@ class ActType_TrtllmGen(IntEnum): Relu2 = 1 # act = x0 * sigmoid(x0) Silu = 2 + # SiTu gated activation (Kimi K3), gate on x1: + # act = (beta * tanh(x0 / beta)) * (alpha * tanh(x1 / alpha) * sigmoid(x1)) + # alpha/beta come from the per-expert gemm1_alpha/gemm1_beta runtime buffers. + SiTu = 3 # IMPORTANT: when adding a new activation type, please update this function. diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index b5cabcd506e0..cbb97e45ed18 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -101,6 +101,17 @@ def __init__( # ``num_experts_per_worker``. This is what unlocks dwdp_size values # that do not divide ``num_experts`` evenly (e.g. dwdp_size = 3, 5) # and the IPC-era redundancy mode where adjacent peer ranges overlap. + # Record whether the caller explicitly requested a MoE TP/EP split + # (the -1 sentinels mean "auto"). Model code (e.g. Kimi K3) uses this + # to distinguish an explicit ``moe_tensor_parallel_size`` / + # ``moe_expert_parallel_size`` request from the auto default, which + # resolves to the same ``(moe_tp=tp_size, moe_ep=1)`` values. Only + # meaningful on mappings built from raw user args: ``to_dict()`` emits + # resolved values, so a ``from_dict`` round-trip marks this True. + self.moe_tp_ep_user_specified = (dwdp_size <= 1 + and (moe_tp_size != -1 + or moe_ep_size != -1)) + if dwdp_size > 1: moe_tp_size = 1 moe_ep_size = 1 diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index 09d50a1727e3..95d7d9955f37 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -38,8 +38,46 @@ def update_quant_config_from_compressed_tensors( ) group_config = next(iter(config_groups.values())) weights_quant_config = group_config["weights"] - inputs_quant_config = group_config["input_activations"] weights_quant_strategy = weights_quant_config["strategy"] + + # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. Handled + # before the weight-algo branches so recipes that early-return (MXFP4 + # pack-quantized) still pick up the KV-cache quantization. + kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") + if kv_cache_scheme is not None: + if kv_cache_scheme.get("num_bits") == 8 and kv_cache_scheme.get("type") == "float": + if quant_config.kv_cache_quant_algo in (None, QuantAlgo.FP8): + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + else: + raise ValueError( + f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, " + "conflicting with FP8 KV cache from HF quant config." + ) + else: + raise ValueError(f"Unsupported kv_cache_scheme: {kv_cache_scheme}.") + + # MXFP4 pack-quantized (weight-only): FP4 E2M1 weights packed two per + # uint8 with per-32-group uint8 E8M0 scales and no activation + # quantization (e.g. Kimi K3 routed experts). Handled before reading the + # input-activation strategy, which is null for weight-only recipes. + if hf_quant_config.get("format") == "mxfp4-pack-quantized" or ( + weights_quant_config["num_bits"] == 4 + and weights_quant_config.get("type") == "float" + and weights_quant_strategy == "group" + and group_config.get("input_activations") is None + ): + group_size = weights_quant_config["group_size"] + if group_size != 32: + raise ValueError(f"Unsupported group_size: {group_size}. Supported: 32 for MXFP4.") + quant_config.quant_algo = QuantAlgo.W4A16_MXFP4 + quant_config.group_size = group_size + hf_exclude_modules = hf_quant_config.get("modules_to_not_convert", None) + quant_config.exclude_modules = list( + set((hf_exclude_modules or []) + hf_quant_config.get("ignore", [])) + ) + return + + inputs_quant_config = group_config["input_activations"] inputs_quant_strategy = inputs_quant_config["strategy"] if weights_quant_config["num_bits"] == 8: @@ -85,20 +123,6 @@ def update_quant_config_from_compressed_tensors( "Supported: 8 (FP8) or 4 (NVFP4)." ) - # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. - kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") - if kv_cache_scheme is not None: - if kv_cache_scheme.get("num_bits") == 8 and kv_cache_scheme.get("type") == "float": - if quant_config.kv_cache_quant_algo in (None, QuantAlgo.FP8): - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - else: - raise ValueError( - f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, " - "conflicting with FP8 KV cache from HF quant config." - ) - else: - raise ValueError(f"Unsupported kv_cache_scheme: {kv_cache_scheme}.") - hf_exclude_modules = hf_quant_config.get("modules_to_not_convert", None) if hf_exclude_modules is not None: quant_config.exclude_modules = list( diff --git a/tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py new file mode 100644 index 000000000000..abd128c3edc3 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py @@ -0,0 +1,540 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""Parity: in-tree ``trtllm::kda_mtp_decode`` CuTe kernel vs two references. + +References: + +1. A pure-torch fp32 CPU golden (vendored from the kernel drop's + ``cpu_reference`` self-check). +2. An FLA sequential reference built from the exact op sequence + ``KimiKDARuntime._forward_verify`` uses in-tree: per-step fp32 causal + conv + SiLU followed by ``fla.ops.kda.fused_recurrent_kda`` with + ``use_qk_l2norm/use_gate/use_beta_sigmoid`` in kernel, ``lower_bound``, + ``state_v_first=True``. + +Agreement of all three establishes both that the kernel is internally +correct and that it computes the same function the model's sequential +verify path computes — i.e. it is a drop-in replacement. + +Round-2 tests exercise the kernel's replay mode (``num_accepted_tokens`` +mixed per request), chained from round-1 CPU-golden outputs. H=8 (K3 +state-TP4) and H=6 (TP16) are the production per-rank head counts outside +the drop's benchmark-tuned set {2, 12, 32}; the vendored v_row hoist fix +makes them compile, and this test validates them numerically. + +Requires: 1 GPU (sm100 for the CuTe kernel), fla-core, nvidia-cutlass-dsl, +cuda-bindings. Skips cleanly otherwise. +""" + +import pytest +import torch +import torch.nn.functional as F + +_HAVE_DEPS = True +_DEP_ERR = None +try: + import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 + from fla.ops.kda import fused_recurrent_kda # noqa: F401 +except ImportError as e: + _HAVE_DEPS = False + _DEP_ERR = str(e) + + +def _is_blackwell(): + if not torch.cuda.is_available(): + return False + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor in (100, 103) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), + pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103"), + pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), +] + +# (N, H): the drop's benchmark-tuned shapes plus K3's production per-rank +# head counts (96 KDA heads: 12 at TP8, 8 at state-TP 12, 6 at TP16). +SHAPES = [ + (128, 2), + (32, 12), + (32, 32), + (32, 8), + (32, 6), +] +M = 2 # NUM_SPEC — spec-token count per verify round + + +def _silu(x): + return x * torch.sigmoid(x) + + +def make_conv_data(B, H, K=128, V=128, M=2, W=4, lower_bound=-5.0, scale=None, seed=2025): + """Random inputs + caches in the op's layout contract (from the drop).""" + torch.manual_seed(seed) + scale = K**-0.5 if scale is None else scale + device = "cuda" + dim = H * K + T = 2 * M + 1 + T_total = B * T + S = W - 1 + M + + x_q = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device=device) + x_k = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device=device) + x_v = torch.randn(1, T_total, H, V, dtype=torch.bfloat16, device=device) + g = torch.randn(1, T_total, H, K, dtype=torch.bfloat16, device=device) + beta = torch.randn(1, T_total, H, dtype=torch.bfloat16, device=device) + w_q = torch.randn(dim, W, dtype=torch.float32, device=device) + w_k = torch.randn(dim, W, dtype=torch.float32, device=device) + w_v = torch.randn(dim, W, dtype=torch.float32, device=device) + A_log = torch.randn(H, dtype=torch.float32, device=device) + dt_bias = torch.randn(dim, dtype=torch.float32, device=device) + # dim-contiguous extended conv caches: allocate [B, S, dim], transpose. + cs_q = torch.randn(B, S, dim, dtype=torch.float32, device=device).transpose(1, 2) + cs_k = torch.randn(B, S, dim, dtype=torch.float32, device=device).transpose(1, 2) + cs_v = torch.randn(B, S, dim, dtype=torch.float32, device=device).transpose(1, 2) + initial_state_kfirst = torch.randn(B, H, K, V, dtype=torch.float32, device=device) + initial_state_cute = initial_state_kfirst.permute(0, 1, 3, 2).contiguous() + cu_seqlens = torch.arange(0, B * T + 1, T, dtype=torch.int32, device=device) + ssm_state_indices = torch.arange(B, dtype=torch.int32, device=device) + num_accepted_tokens = torch.zeros(B, dtype=torch.int32, device=device) + qkg_cache = torch.zeros(B, M, 3, dim, dtype=torch.float32, device=device) + v_cache = torch.zeros(B, M, H * V, dtype=torch.float32, device=device) + beta_cache = torch.zeros(B, M, H, dtype=torch.float32, device=device) + return { + "x_q": x_q, + "x_k": x_k, + "x_v": x_v, + "w_q": w_q, + "w_k": w_k, + "w_v": w_v, + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "cs_q": cs_q, + "cs_k": cs_k, + "cs_v": cs_v, + "initial_state_kfirst": initial_state_kfirst, + "initial_state_cute": initial_state_cute, + "qkg_cache": qkg_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + "cu_seqlens": cu_seqlens, + "ssm_state_indices": ssm_state_indices, + "num_accepted_tokens": num_accepted_tokens, + "B": B, + "H": H, + "K": K, + "V": V, + "M": M, + "W": W, + "T": T, + "T_total": T_total, + "lower_bound": lower_bound, + "scale": scale, + } + + +def cpu_reference(data): + """fp32 pure-torch golden with replay semantics (from the drop).""" + B, H, K, V, M, W = (data["B"], data["H"], data["K"], data["V"], data["M"], data["W"]) + T = data["T"] + lower_bound = data["lower_bound"] + scale = data["scale"] + w_q = data["w_q"].float().cpu() + w_k = data["w_k"].float().cpu() + w_v = data["w_v"].float().cpu() + A_log = data["A_log"].float().cpu() + dt_bias = data["dt_bias"].float().cpu() + x_q = data["x_q"].float().cpu() + x_k = data["x_k"].float().cpu() + x_v = data["x_v"].float().cpu() + g = data["g"].float().cpu() + beta = data["beta"].float().cpu() + cs_q = data["cs_q"].float().cpu().clone() + cs_k = data["cs_k"].float().cpu().clone() + cs_v = data["cs_v"].float().cpu().clone() + qkg_cache = data["qkg_cache"].float().cpu().clone() + v_cache = data["v_cache"].float().cpu().clone() + beta_cache = data["beta_cache"].float().cpu().clone() + ht = data["initial_state_kfirst"].float().cpu().clone() + ht_commit = ht.clone() + out = torch.zeros(1, B * T, H, V, dtype=torch.float32) + + for n in range(B): + bos = n * T + slot = n + commit_len = int(data["num_accepted_tokens"][n].item()) + T_loop = commit_len + 1 + M + for h in range(H): + hk = h * K + hv = h * V + hist_q = cs_q[slot, hk : hk + K, : W - 1].clone() + hist_k = cs_k[slot, hk : hk + K, : W - 1].clone() + hist_v = cs_v[slot, hv : hv + V, : W - 1].clone() + h_state = ht[slot, h].clone() + for i_t in range(T_loop): + if i_t < commit_len: + q_t = qkg_cache[slot, i_t, 0, hk : hk + K] + k_t = qkg_cache[slot, i_t, 1, hk : hk + K] + gk_t = qkg_cache[slot, i_t, 2, hk : hk + K] + v_t = v_cache[slot, i_t, hv : hv + V] + beta_t = beta_cache[slot, i_t, h] + xq_raw = cs_q[slot, hk : hk + K, W - 1 + i_t] + xk_raw = cs_k[slot, hk : hk + K, W - 1 + i_t] + xv_raw = cs_v[slot, hv : hv + V, W - 1 + i_t] + hist_q = torch.cat([hist_q[:, 1:], xq_raw.unsqueeze(-1)], dim=1) + hist_k = torch.cat([hist_k[:, 1:], xk_raw.unsqueeze(-1)], dim=1) + hist_v = torch.cat([hist_v[:, 1:], xv_raw.unsqueeze(-1)], dim=1) + else: + token = bos + i_t + xq_raw = x_q[0, token, h] + xk_raw = x_k[0, token, h] + xv_raw = x_v[0, token, h] + cq = (torch.cat([hist_q, xq_raw.unsqueeze(-1)], dim=-1) * w_q[hk : hk + K]).sum( + dim=-1 + ) + ck = (torch.cat([hist_k, xk_raw.unsqueeze(-1)], dim=-1) * w_k[hk : hk + K]).sum( + dim=-1 + ) + cv = (torch.cat([hist_v, xv_raw.unsqueeze(-1)], dim=-1) * w_v[hv : hv + V]).sum( + dim=-1 + ) + q_t = F.normalize(cq / (1.0 + torch.exp(-cq)), p=2, dim=-1) * scale + k_t = F.normalize(ck / (1.0 + torch.exp(-ck)), p=2, dim=-1) + v_t = cv / (1.0 + torch.exp(-cv)) + gr = g[0, token, h] + dt_bias[hk : hk + K] + gk_t = lower_bound * torch.sigmoid(gr * torch.exp(A_log[h])) + beta_t = torch.sigmoid(beta[0, token, h]) + hist_q = torch.cat([hist_q[:, 1:], xq_raw.unsqueeze(-1)], dim=1) + hist_k = torch.cat([hist_k[:, 1:], xk_raw.unsqueeze(-1)], dim=1) + hist_v = torch.cat([hist_v[:, 1:], xv_raw.unsqueeze(-1)], dim=1) + + decay = torch.exp(gk_t) + h_state = h_state * decay.unsqueeze(1) + sum_hk = (h_state * k_t.unsqueeze(1)).sum(dim=0) + v_new = (v_t - sum_hk) * beta_t + h_state = h_state + k_t.unsqueeze(1) * v_new.unsqueeze(0) + o_t = (h_state * q_t.unsqueeze(1)).sum(dim=0) + if i_t >= commit_len: + out[0, bos + i_t, h] = o_t + if i_t == commit_len: + ht_commit[slot, h] = h_state + cs_q[slot, hk : hk + K, : W - 1] = hist_q + cs_k[slot, hk : hk + K, : W - 1] = hist_k + cs_v[slot, hv : hv + V, : W - 1] = hist_v + if i_t > commit_len: + cache_pos = i_t - commit_len - 1 + qkg_cache[slot, cache_pos, 0, hk : hk + K] = q_t + qkg_cache[slot, cache_pos, 1, hk : hk + K] = k_t + qkg_cache[slot, cache_pos, 2, hk : hk + K] = gk_t + v_cache[slot, cache_pos, hv : hv + V] = v_t + beta_cache[slot, cache_pos, h] = beta_t + cs_q[slot, hk : hk + K, W - 1 + cache_pos] = xq_raw + cs_k[slot, hk : hk + K, W - 1 + cache_pos] = xk_raw + cs_v[slot, hv : hv + V, W - 1 + cache_pos] = xv_raw + + return { + "out": out, + "recurrent_state": ht_commit, + "qkg_cache": qkg_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + "cs_q": cs_q, + "cs_k": cs_k, + "cs_v": cs_v, + } + + +def cute_run(data, zero_accepted_hint=False): + """Run the in-tree op on cloned caches; return the drop-format dict.""" + import tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops # noqa: F401 + + state = data["initial_state_cute"].clone() + cs = {} + for name in ("cs_q", "cs_k", "cs_v"): + src = data[name] + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + cs[name] = dst + qkg_cache = data["qkg_cache"].clone() + v_cache = data["v_cache"].clone() + beta_cache = data["beta_cache"].clone() + out = torch.ops.trtllm.kda_mtp_decode( + x_q=data["x_q"], + x_k=data["x_k"], + x_v=data["x_v"], + w_q=data["w_q"], + w_k=data["w_k"], + w_v=data["w_v"], + cs_q=cs["cs_q"], + cs_k=cs["cs_k"], + cs_v=cs["cs_v"], + g=data["g"], + beta=data["beta"], + A_log=data["A_log"], + dt_bias=data["dt_bias"], + recurrent_state=state, + qkg_cache=qkg_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=data["ssm_state_indices"], + cu_seqlens=data["cu_seqlens"], + num_spec=data["M"], + num_accepted_tokens=data["num_accepted_tokens"], + lower_bound=data["lower_bound"], + scale=data["scale"], + zero_accepted_hint=zero_accepted_hint, + ) + return { + "out": out, + # committed state back in K-first layout for CPU-golden comparison + "recurrent_state": state.permute(0, 1, 3, 2).contiguous(), + "state_v_first": state, + "qkg_cache": qkg_cache, + "v_cache": v_cache, + "beta_cache": beta_cache, + "cs_q": cs["cs_q"], + "cs_k": cs["cs_k"], + "cs_v": cs["cs_v"], + } + + +def _fla_sequential_reference(data, num_accepted): + """Per-request sequential conv+SiLU (fp32 torch) + fused_recurrent_kda. + + Mirrors ``KimiKDARuntime._forward_verify``'s op sequence, extended with + the replay prefix: for request ``n`` with ``a = num_accepted[n]``, the + processed token sequence is ``a`` cached raw tokens (re-convolved from + the extended conv-cache slots) followed by the ``1 + M`` new tokens. + Returns out rows (new tokens only) and the committed state (after the + first new token, FLA/cute ``[B, H, V, K]`` layout). + """ + from fla.ops.kda import fused_recurrent_kda + + B, H, K, V, W = data["B"], data["H"], data["K"], data["V"], data["W"] + T = data["T"] + dim = H * K + dev = data["x_q"].device + out = torch.zeros(1, B * T, H, V, dtype=torch.float32, device=dev) + committed = torch.zeros(B, H, V, K, dtype=torch.float32, device=dev) + + w_q, w_k, w_v = (data[k].float() for k in ("w_q", "w_k", "w_v")) + x_q, x_k, x_v = (data[k].float() for k in ("x_q", "x_k", "x_v")) + g_all, beta_all = data["g"], data["beta"] + + for n in range(B): + a = int(num_accepted[n]) + bos = n * T + hist = { + "q": data["cs_q"][n, :, : W - 1].float().clone(), + "k": data["cs_k"][n, :, : W - 1].float().clone(), + "v": data["cs_v"][n, :, : W - 1].float().clone(), + } + state = data["initial_state_cute"][n : n + 1].float().clone() + + for i_t in range(a + 1 + M): + if i_t < a: # replay a cached token (raw x from cache slots) + xq = data["cs_q"][n, :, W - 1 + i_t].float() + xk = data["cs_k"][n, :, W - 1 + i_t].float() + xv = data["cs_v"][n, :, W - 1 + i_t].float() + tok = None + g_t = data["qkg_cache"][n, i_t, 2].float() + beta_t = data["beta_cache"][n, i_t].float() + replay = True + else: + tok = bos + i_t + xq = x_q[0, tok].reshape(dim) + xk = x_k[0, tok].reshape(dim) + xv = x_v[0, tok].reshape(H * V) + replay = False + + def conv_step(hist_s, x_raw, w): + window = torch.cat([hist_s, x_raw.unsqueeze(-1)], dim=-1) + y = (window * w).sum(dim=-1) + return y, window[:, 1:] + + cq, hist["q"] = conv_step(hist["q"], xq, w_q) + ck, hist["k"] = conv_step(hist["k"], xk, w_k) + cv, hist["v"] = conv_step(hist["v"], xv, w_v) + + if replay: + # Replayed tokens use the cached post-processed k/g/v/beta + # exactly as the kernel does (delta rule applied directly). + k_t = data["qkg_cache"][n, i_t, 1].float().view(H, K) + gk_t = g_t.view(H, K) + v_t = data["v_cache"][n, i_t].float().view(H, V) + st = state[0] + decay = torch.exp(gk_t) + st = st * decay.unsqueeze(1) + sum_hk = torch.einsum("hvk,hk->hv", st, k_t) + v_new = (v_t - sum_hk) * beta_t.unsqueeze(-1) + st = st + torch.einsum("hk,hv->hvk", k_t, v_new) + state = st.unsqueeze(0) + else: + # fp32 hand-off into FLA: the comparison target is the + # mathematical function, not the model's bf16 dataflow. + q_in = _silu(cq).view(1, 1, H, K) + k_in = _silu(ck).view(1, 1, H, K) + v_in = _silu(cv).view(1, 1, H, V) + o_t, state = fused_recurrent_kda( + q=q_in, + k=k_in, + v=v_in, + g=g_all[0, tok].view(1, 1, H, K), + beta=beta_all[0, tok].view(1, 1, H).float(), + A_log=data["A_log"], + dt_bias=data["dt_bias"], + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=data["lower_bound"], + state_v_first=True, + ) + out[0, tok] = o_t[0, 0].float() + if i_t == a: # first new (golden) token -> committed state + committed[n] = state[0].float() + return out, committed + + +def _cute_layout_state(cpu_out): + # CPU golden reports K-first [B, H, K, V]; cute/FLA layout is [B,H,V,K] + return cpu_out["recurrent_state"].permute(0, 1, 3, 2).contiguous() + + +def _assert_close(name, a, b, atol, rtol=0.0): + diff = (a.float() - b.float()).abs() + denom = b.float().abs().clamp_min(1.0) + ok = (diff <= atol + rtol * denom).all() + assert ok, ( + f"{name}: max_abs={diff.max().item():.3e} " + f"(atol={atol}, worst rel={((diff / denom).max()):.3e})" + ) + + +@pytest.mark.parametrize("B,H", SHAPES, ids=lambda v: str(v)) +def test_round1_zero_accepted(B, H): + """Fresh verify round (no replay): kernel vs CPU golden vs FLA seq.""" + data = make_conv_data(B, H, M=M, seed=2025) + T = data["T"] + + cpu = {k: v.cuda() for k, v in cpu_reference(data).items()} + cute_out = cute_run(data) + fla_out, fla_committed = _fla_sequential_reference(data, data["num_accepted_tokens"].cpu()) + + new_rows = torch.cat([torch.arange(n * T, n * T + 1 + M) for n in range(B)]).cuda() + + # Kernel vs the fp32 golden (tight: fp32 accumulation). + _assert_close( + "out(cute vs cpu)", cute_out["out"][0, new_rows], cpu["out"][0, new_rows], atol=2e-2 + ) + _assert_close( + "state(cute vs cpu)", cute_out["recurrent_state"], cpu["recurrent_state"], atol=1e-4 + ) + for name in ("qkg_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + _assert_close(f"{name}(cute vs cpu)", cute_out[name], cpu[name], atol=2e-2) + + # Kernel vs the FLA sequential path (the in-tree _forward_verify math). + _assert_close( + "out(cute vs fla)", + cute_out["out"][0, new_rows].float(), + fla_out[0, new_rows], + atol=5e-2, + rtol=5e-2, + ) + _assert_close( + "state(cute vs fla)", cute_out["state_v_first"], fla_committed, atol=5e-3, rtol=5e-3 + ) + + +@pytest.mark.parametrize("B,H", [(128, 2), (32, 12), (32, 6)], ids=lambda v: str(v)) +def test_round2_replay(B, H): + """Replay round: mixed num_accepted per request, chained from a CPU- + golden round 1. Validates the kernel's cache-replay state math (the + path the drop's own self-check never exercised).""" + data = make_conv_data(B, H, M=M, seed=7) + T = data["T"] + + # Round 1 (all zero accepted) on the CPU golden to produce the caches + # and committed state that seed round 2. + r1 = cpu_reference(data) + + # Round 2 inputs: fresh tokens, caches/state/conv from round 1. + data2 = make_conv_data(B, H, M=M, seed=8) + for name in ("qkg_cache", "v_cache", "beta_cache"): + data2[name] = r1[name].cuda().contiguous() + for name in ("cs_q", "cs_k", "cs_v"): + # Preserve the contract's dim-contiguous (transposed) layout. + src = r1[name].cuda() + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=src.dtype, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data2[name] = dst + data2["initial_state_kfirst"] = r1["recurrent_state"].cuda().contiguous() + data2["initial_state_cute"] = r1["recurrent_state"].cuda().permute(0, 1, 3, 2).contiguous() + # Mixed acceptance: 0, 1, 2 cycling across requests. + accept = torch.arange(B, dtype=torch.int32) % (M + 1) + data2["num_accepted_tokens"] = accept.cuda() + + cpu2 = {k: v.cuda() for k, v in cpu_reference(data2).items()} + cute2 = cute_run(data2) + fla_out2, fla_committed2 = _fla_sequential_reference(data2, accept) + + rows = torch.cat( + [torch.arange(n * T + int(accept[n]), n * T + int(accept[n]) + 1 + M) for n in range(B)] + ).cuda() + + _assert_close("out2(cute vs cpu)", cute2["out"][0, rows], cpu2["out"][0, rows], atol=2e-2) + _assert_close( + "state2(cute vs cpu)", cute2["recurrent_state"], cpu2["recurrent_state"], atol=1e-4 + ) + _assert_close( + "out2(cute vs fla)", cute2["out"][0, rows].float(), fla_out2[0, rows], atol=5e-2, rtol=5e-2 + ) + _assert_close( + "state2(cute vs fla)", cute2["state_v_first"], fla_committed2, atol=5e-3, rtol=5e-3 + ) + + +@pytest.mark.parametrize("B,H", [(32, 12)], ids=lambda v: str(v)) +def test_zero_accepted_hint_variant(B, H): + """The zero_accepted_hint fast variant matches the general variant.""" + data = make_conv_data(B, H, M=M, seed=11) + general = cute_run(data, zero_accepted_hint=False) + fast = cute_run(data, zero_accepted_hint=True) + for name in ( + "out", + "recurrent_state", + "qkg_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + _assert_close(f"{name}(fast vs general)", fast[name], general[name], atol=1e-5) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-x"])) diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py new file mode 100644 index 000000000000..9006ed8fb7ee --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py @@ -0,0 +1,235 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""Runtime-level parity: KimiKDARuntime fused verify vs sequential verify. + +Simulates two chained speculative-verification rounds through +``KimiKDARuntime._forward_verify`` in both worlds: + +* Sequential world: the legacy intermediate-buffer path + (``_forward_verify_sequential``) plus the manager's legacy promotion + (copy the accepted step's conv window / SSM state into the live pools). +* Fused world: the ``trtllm::kda_mtp_decode`` replay path + (``_forward_verify_fused``) with per-slot replay caches, in-place state + commit after the golden token, and only the accepted-draft count + recorded between rounds. + +Identical hidden states are fed to both worlds; with mixed per-request +acceptance between rounds, matching round-2 outputs proves the fused +path's replay bookkeeping (shifted ``cu_seqlens`` layout, conv-window +seeding, pending-count plumbing) reproduces the promoted-state semantics. + +Requires 1 Blackwell GPU, fla-core, nvidia-cutlass-dsl. Skips otherwise. +""" + +from types import SimpleNamespace + +import pytest +import torch + +_HAVE_DEPS = True +_DEP_ERR = None +try: + import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 + from fla.ops.kda import fused_recurrent_kda # noqa: F401 +except ImportError as e: + _HAVE_DEPS = False + _DEP_ERR = str(e) + + +def _is_blackwell(): + if not torch.cuda.is_available(): + return False + prop = torch.cuda.get_device_properties(0) + return prop.major * 10 + prop.minor in (100, 103) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU"), + pytest.mark.skipif(not _is_blackwell(), reason="needs sm100/sm103"), + pytest.mark.skipif(not _HAVE_DEPS, reason=f"deps: {_DEP_ERR}"), +] + +HIDDEN = 512 +H = 8 # per-rank head count outside the drop's tuned set — general variant +K = 128 +W = 4 +M = 2 # draft tokens per round +LB = -5.0 + + +def _make_runtime(seed): + from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime + + cfg = SimpleNamespace( + hidden_size=HIDDEN, + rms_norm_eps=1e-5, + linear_attn_config=dict( + num_heads=H, + head_dim=K, + short_conv_kernel_size=W, + use_full_rank_gate=True, + gate_lower_bound=LB, + ), + ) + rt = KimiKDARuntime(cfg, layer_idx=0).to("cuda") + gen = torch.Generator(device="cuda").manual_seed(seed) + with torch.no_grad(): + for name, p in rt.named_parameters(): + if name.endswith("A_log"): + p.copy_( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5 + ) + elif name.endswith("dt_bias"): + p.copy_( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1 + ) + else: + p.copy_( + ( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) + * 0.03 + ).to(p.dtype) + ) + return rt + + +def _make_pools(B, seed): + gen = torch.Generator(device="cuda").manual_seed(seed) + d = H * K + conv_pool = ( + torch.randn(B, 3 * d, W, generator=gen, device="cuda", dtype=torch.float32) * 0.5 + ).to(torch.bfloat16) + ssm_pool = torch.randn(B, H, K, K, generator=gen, device="cuda", dtype=torch.float32) + ssm_pool *= torch.linspace(0.5, 1.5, K, device="cuda").view(1, 1, K, 1) + return conv_pool, ssm_pool + + +def _make_fused_layer_cache(B, conv_pool): + """Replay caches shaped like PythonMambaCacheManager's KDA allocation, + with the committed conv window seeded from the base pool (the prefill + seeding contract: FLA window columns [1, W) -> committed columns).""" + d = H * K + S = W - 1 + M + + def _conv_cache(section): + cache = torch.zeros(B, S, d, device="cuda", dtype=torch.float32).transpose(-1, -2) + cache[:, :, : W - 1] = conv_pool[:, section * d : (section + 1) * d, 1:].float() + return cache + + return SimpleNamespace( + kda_conv_q=_conv_cache(0), + kda_conv_k=_conv_cache(1), + kda_conv_v=_conv_cache(2), + kda_qkg_cache=torch.zeros(B, M, 3, d, device="cuda", dtype=torch.float32), + kda_v_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.float32), + kda_beta_cache=torch.zeros(B, M, H, device="cuda", dtype=torch.float32), + prev_num_accepted_tokens=torch.zeros(B, dtype=torch.int32, device="cuda"), + intermediate_conv_window=None, + intermediate_ssm=None, + ) + + +def _make_seq_layer_cache(B): + d = H * K + return SimpleNamespace( + kda_qkg_cache=None, + intermediate_conv_window=torch.zeros( + B, M + 1, 3 * d, W, device="cuda", dtype=torch.bfloat16 + ), + intermediate_ssm=torch.zeros(B, M + 1, H, K, K, device="cuda", dtype=torch.float32), + ) + + +def _promote_sequential(layer_cache, conv_pool, ssm_pool, accept): + """The manager's legacy promotion: accepted step's states -> pools.""" + B = conv_pool.shape[0] + rows = torch.arange(B, device="cuda") + conv_pool.copy_(layer_cache.intermediate_conv_window[rows, accept]) + ssm_pool.copy_(layer_cache.intermediate_ssm[rows, accept]) + + +def _rep(name, a, b): + a, b = a.float(), b.float() + cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() + rel = ((a - b).norm() / (b.norm() + 1e-12)).item() + print(f" {name}: cos={cos:.6f} rel_l2={rel:.3e}") + return cos > 0.999 and rel < 3e-2 + + +def test_fused_vs_sequential_two_rounds(): + torch.manual_seed(0) + B = 4 + T = M + 1 + rt = _make_runtime(seed=1) + slot_indices = torch.arange(B, dtype=torch.long, device="cuda") + + conv_pool_seq, ssm_pool_seq = _make_pools(B, seed=2) + conv_pool_fused = conv_pool_seq.clone() + ssm_pool_fused = ssm_pool_seq.clone() + cache_seq = _make_seq_layer_cache(B) + cache_fused = _make_fused_layer_cache(B, conv_pool_fused) + + gen = torch.Generator(device="cuda").manual_seed(3) + + def tokens(scale=0.5): + return ( + torch.randn(B * T, HIDDEN, generator=gen, device="cuda", dtype=torch.float32) * scale + ).to(torch.bfloat16) + + ok = True + with torch.no_grad(): + # ---- Round 1 (no pending drafts) ---- + x1 = tokens() + out1_seq = rt._forward_verify_sequential( + x1, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices + ) + out1_fused = rt._forward_verify( + x1, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices + ) + print("round 1:") + ok &= _rep("out", out1_fused, out1_seq) + + # ---- Acceptance: 0, 1, 2, 0 drafts across the 4 requests ---- + accept = torch.tensor([0, 1, 2, 0], dtype=torch.long, device="cuda") + _promote_sequential(cache_seq, conv_pool_seq, ssm_pool_seq, accept) + cache_fused.prev_num_accepted_tokens.copy_(accept.to(torch.int32)) + + # ---- Round 2 (fused path replays the accepted drafts) ---- + x2 = tokens() + out2_seq = rt._forward_verify_sequential( + x2, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices + ) + out2_fused = rt._forward_verify( + x2, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices + ) + print("round 2 (mixed replay):") + ok &= _rep("out", out2_fused, out2_seq) + + # Committed pool state cross-check: fused pool holds the state after + # round-2's golden token; reproduce it in the sequential world by + # promoting with accept=0 (golden only). + _promote_sequential( + cache_seq, conv_pool_seq, ssm_pool_seq, torch.zeros(B, dtype=torch.long, device="cuda") + ) + ok &= _rep("committed ssm", ssm_pool_fused, ssm_pool_seq) + + assert ok + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-x", "-s"])) diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py new file mode 100644 index 000000000000..0b20cfc34bd2 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KDA speculative-verification parity: _forward_verify vs sequential decode. + +The verify path must produce, for every step t, exactly the state and output +that t sequential single-token _forward_decode calls would produce — the two +paths call the same FLA kernels with the same [B, 1] shapes, so agreement is +expected to near-bitwise tolerance. A real mismatch here means the verify +implementation (state threading, conv stepping, intermediate writes) is +wrong; e2e text divergence on truncated models alone does NOT indicate a bug +(batched MoE reduction order flips argmax on noise logits). + +Needs 1 GPU + fla-core; runs with random weights (no checkpoint). +""" + +import pytest +import torch + +pytest.importorskip("fla") + +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime + + +class _Cfg: + hidden_size = 256 + rms_norm_eps = 1e-6 + linear_attn_config = { + "num_heads": 4, + "head_dim": 64, + "short_conv_kernel_size": 4, + "use_full_rank_gate": True, + "gate_lower_bound": None, + } + + +class _LayerCache: + def __init__(self, slots, dim3, w, h, v, k, t_max, device): + self.conv = torch.zeros(slots, dim3, w, dtype=torch.bfloat16, device=device) + self.temporal = torch.zeros(slots, h, v, k, dtype=torch.float32, device=device) + self.intermediate_conv_window = torch.zeros( + slots, t_max, dim3, w, dtype=torch.bfloat16, device=device + ) + self.intermediate_ssm = torch.zeros( + slots, t_max, h, v, k, dtype=torch.float32, device=device + ) + + +@pytest.mark.parametrize("batch", [1, 3]) +@pytest.mark.parametrize("t_steps", [2, 3]) +def test_kda_verify_matches_sequential_decode(batch, t_steps): + if not torch.cuda.is_available(): + pytest.skip("needs a GPU") + torch.manual_seed(0) + device = "cuda" + cfg = _Cfg() + lin = cfg.linear_attn_config + h = lin["num_heads"] + dim = h * lin["head_dim"] + w = lin["short_conv_kernel_size"] + + runtime = KimiKDARuntime(cfg, layer_idx=0).to(device) + slots = batch + 2 # non-trivial slot mapping + cache = _LayerCache(slots, 3 * dim, w, h, lin["head_dim"], lin["head_dim"], t_steps, device) + slot_indices = torch.arange(2, 2 + batch, device=device, dtype=torch.long) + + # Random-but-fixed starting state and inputs. + torch.nn.init.normal_(cache.conv[2 : 2 + batch], std=0.02) + torch.nn.init.normal_(cache.temporal[2 : 2 + batch], std=0.02) + x = torch.randn(batch, t_steps, cfg.hidden_size, dtype=torch.bfloat16, device=device) * 0.1 + + # --- Reference: t sequential in-place decodes on a cloned pool. --- + ref_conv = cache.conv.clone() + ref_ssm = cache.temporal.clone() + ref_outs, ref_conv_steps, ref_ssm_steps = [], [], [] + for t in range(t_steps): + out = runtime._forward_decode(x[:, t], ref_conv, ref_ssm, slot_indices) + ref_outs.append(out) + ref_conv_steps.append(ref_conv.index_select(0, slot_indices).clone()) + ref_ssm_steps.append(ref_ssm.index_select(0, slot_indices).clone()) + + # --- Verify path: one call, intermediates into the scratch buffers. --- + pristine_conv = cache.conv.clone() + pristine_ssm = cache.temporal.clone() + out_verify = runtime._forward_verify( + x.reshape(batch * t_steps, cfg.hidden_size), + t_steps, + cache, + cache.conv, + cache.temporal, + slot_indices, + ) + + # Live pools must be untouched by verification. + torch.testing.assert_close(cache.conv, pristine_conv, rtol=0, atol=0) + torch.testing.assert_close(cache.temporal, pristine_ssm, rtol=0, atol=0) + + out_verify = out_verify.reshape(batch, t_steps, cfg.hidden_size) + for t in range(t_steps): + torch.testing.assert_close(out_verify[:, t], ref_outs[t], rtol=2e-2, atol=2e-2) + torch.testing.assert_close( + cache.intermediate_conv_window[:batch, t], ref_conv_steps[t], rtol=2e-2, atol=2e-2 + ) + torch.testing.assert_close( + cache.intermediate_ssm[:batch, t], ref_ssm_steps[t], rtol=2e-2, atol=2e-2 + ) diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py new file mode 100644 index 000000000000..cea663614621 --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parity tests for the fused Kimi K3 attention-residual op.""" + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3RMSNorm, _apply_attn_res_fused +from tensorrt_llm._torch.modules.kimi_k3_attn_res import apply_attn_res_reference + +HIDDEN_SIZE = 7168 +RMS_EPS = 1e-6 + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="Kimi K3 is supported only on Blackwell (SM100/SM103)", +) + + +def _similarity(actual: torch.Tensor, expected: torch.Tensor) -> tuple[float, float]: + actual_float = actual.float() + expected_float = expected.float() + cosine = torch.nn.functional.cosine_similarity( + actual_float.flatten(), expected_float.flatten(), dim=0 + ).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + return cosine, relative_l2 + + +@pytest.mark.parametrize( + ("num_tokens", "num_snapshots"), + [ + (1, 0), + (1, 1), + (1, 3), + (1, 7), + (1, 11), + (64, 0), + (128, 3), + (1024, 11), + (300, 5), + (16384, 11), + ], +) +@torch.no_grad() +def test_fused_attn_res_matches_torch_reference(num_tokens: int, num_snapshots: int) -> None: + torch.manual_seed(0) + projection = nn.Linear(HIDDEN_SIZE, 1, bias=False, dtype=torch.bfloat16, device="cuda") + norm = KimiK3RMSNorm(HIDDEN_SIZE, eps=RMS_EPS).to(device="cuda", dtype=torch.bfloat16) + projection.weight.mul_(0.02) + + prefix_sum = torch.randn(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + # Kernel-native [num_snapshots, num_tokens, H] layout — the layout the + # model's `_apply_attn_res_fused` wrapper takes (snapshots are stacked on + # dim 0 at runtime). The reference implementation uses the HF layout with + # the snapshot axis in the middle, hence the transpose below. + block_residual = ( + torch.randn( + num_snapshots, + num_tokens, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + + expected = apply_attn_res_reference( + prefix_sum, + block_residual.transpose(0, 1), + projection.weight, + norm.weight, + RMS_EPS, + ) + actual = _apply_attn_res_fused(prefix_sum, block_residual, projection, norm) + + assert actual is not None + cosine, relative_l2 = _similarity(actual, expected) + assert cosine > 0.999 + assert relative_l2 < 3e-2 diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py new file mode 100644 index 000000000000..fdf5b7a8b961 --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py @@ -0,0 +1,382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime-integration soundness tests for the optimized KDA prefill op. + +Covers the two classes of bug that op-math parity tests cannot see: + +1. STREAM soundness. The executor runs the model on a dedicated non-blocking + ``torch.cuda.Stream`` (``py_executor.execution_stream``); the CuTe DSL + kernels must launch on torch's current stream, not the DSL default stream. + A default-stream launch races with the projections producing q/k/v/g/beta + and with consumers of the output — silent, intermittent corruption e2e + (GSM8K 68.99 vs 97.01 on the FLA control) while single-stream unit tests + all pass. ``test_nondefault_stream_parity`` reproduces this + deterministically by delaying the execution stream with ``_sleep`` so any + kernel launched on the default stream reads not-yet-written inputs. + +2. id()-keyed cache soundness under the runtime's fresh-tensor-per-call + pattern. CPython recycles an object's address (= its id) as soon as it is + freed; the ``_prune_on_gc`` weakref guard must pop the cache entry at + dealloc, before a new tensor can alias the key. The recycled-id tests + engineer exactly that aliasing: cache a verdict under id(t1), free t1, + re-allocate until a new tensor lands on the same id, refill it with data + for which the stale verdict would be WRONG, and parity-check against FLA. + ``test_single_seq_same_object_repeated`` additionally covers the + Phase 2.1 poisoning case (same cu_seqlens object, two calls). +""" + +import gc + +import pytest +import torch + +pytest.importorskip("fla") + +from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import KDAKernelDispatch # noqa: E402 + +NUM_HEADS = 96 +HEAD_K_DIM = 128 +LOWER_BOUND = -5.0 + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="Kimi K3 is supported only on Blackwell (SM100/SM103)", +) + + +def _op_module(): + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops + + return cute_dsl_kimi_k3_custom_ops + + +@pytest.fixture(scope="module") +def dispatch_pair(): + optimized = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False) + assert optimized.prefill_kernel_path == "optimized" + reference = KDAKernelDispatch(use_optimized_prefill=False, use_optimized_decode=False) + assert reference.prefill_kernel_path == "fla" + return optimized, reference + + +@pytest.fixture(scope="module") +def gate_params(): + torch.manual_seed(0) + a_log = torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.5 + dt_bias = torch.randn(NUM_HEADS * HEAD_K_DIM, dtype=torch.float32, device="cuda") * 0.1 + return a_log, dt_bias + + +def _make_inputs(total_t: int, seed: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + h, k = NUM_HEADS, HEAD_K_DIM + + def rnd(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) + + q = rnd(1, total_t, h, k) + key = rnd(1, total_t, h, k) + v = rnd(1, total_t, h, k) + g = rnd(1, total_t, h, k) + beta = rnd(1, total_t, h, dtype=torch.float32) + return q, key, v, g, beta + + +def _run(dispatch, gate_params, q, k, v, g, beta, cu): + """No .clone() on cu — these tests exercise object identity on purpose.""" + a_log, dt_bias = gate_params + return dispatch.prefill_chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + A_log=a_log, + dt_bias=dt_bias, + scale=HEAD_K_DIM**-0.5, + initial_state=None, + safe_gate=True, + lower_bound=LOWER_BOUND, + cu_seqlens=cu, + ) + + +def _assert_close(name, actual, expected): + actual, expected = actual.float(), expected.float() + cos = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0).item() + rel = ((actual - expected).norm() / (expected.norm() + 1e-12)).item() + assert cos > 0.999 and rel < 3e-2, f"{name}: cos={cos:.6f} rel_l2={rel:.3e}" + + +def _make_cu(lens): + return torch.tensor( + [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" + ) + + +def _flush_tensor_cache_pins(): + """Evict external pins on cu_seqlens tensors. + + Both the dispatcher (fla.ops.utils.index) and the op + (tensorrt_llm._torch.modules.fla.index) run cu_seqlens through + ``@tensor_cache`` helpers that keep the 4 most recent (args, result) + tuples alive. Until those entries are evicted, a cu_seqlens object + stays pinned after the call — which is SOUND (the id cannot be + recycled while our id-keyed entry exists) but makes the recycled-id + scenario unreachable. Churn the caches with fresh dummy tensors so the + pins drop and the finalizer-prune path can be exercised. + """ + from fla.ops.utils.index import prepare_chunk_indices as fla_pci + from fla.ops.utils.index import prepare_chunk_offsets as fla_pco + + from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices as intree_pci + + for _ in range(5): + dummy = torch.tensor([0, 64], dtype=torch.long, device="cuda") + fla_pci(dummy, 64) + fla_pco(dummy, 64) + intree_pci(dummy, 64) + + +def _release_cu(mod, cu_holder): + """Free the cu_seqlens tensor held in the one-element list ``cu_holder`` + and return (old_id, pruned). ``pruned`` False means something still pins + the tensor — sound, but the recycled-id scenario is unreachable; the + caller must then prove the id cannot be recycled and skip. + """ + old_id = id(cu_holder[0]) + cu_holder.clear() + _flush_tensor_cache_pins() + gc.collect() + return old_id, old_id not in mod._varlen_pure_cache + + +def _alloc_with_recycled_id(target_id, make, attempts=512): + """Allocate via ``make()`` until an object lands on ``target_id``. + + Wrong-address candidates are kept alive so the allocator cannot hand the + same wrong slot back; pymalloc free-lists normally return the freed + address on the first attempt. Returns None if unattainable. + """ + hold = [] + for _ in range(attempts): + cand = make() + if id(cand) == target_id: + return cand + hold.append(cand) + return None + + +# --------------------------------------------------------------------------- +# 1. Stream soundness — deterministic repro of the e2e accuracy regression. +# --------------------------------------------------------------------------- + + +def _make_eqlen_inputs(batch, t, seed): + gen = torch.Generator(device="cuda").manual_seed(seed) + h, k = NUM_HEADS, HEAD_K_DIM + + def rnd(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) + + return ( + rnd(batch, t, h, k), + rnd(batch, t, h, k), + rnd(batch, t, h, k), + rnd(batch, t, h, k), + rnd(batch, t, h, dtype=torch.float32), + ) + + +@pytest.mark.parametrize("regime", ["eqlen", "varlen"]) +@torch.no_grad() +def test_nondefault_stream_parity(dispatch_pair, gate_params, regime): + """Run the optimized prefill on a fresh non-default stream (like the + executor's ``execution_stream``) whose queue is held back by a GPU sleep, + with the inputs produced BEHIND that sleep, and read the outputs on that + same stream right after the call. A kernel the op launches on the default + stream instead of the current stream reads inputs before they exist + (eqlen regime — the op has no host sync there, so the repro is + deterministic) and/or its output is consumed before it is written + (varlen regime — the ``out * 1.0`` capture below races a stray + default-stream K4).""" + optimized, reference = dispatch_pair + if regime == "eqlen": + lens = None + make = lambda seed: _make_eqlen_inputs(2, 256, seed) # noqa: E731 + else: + lens = [100, 257, 300] # multi-seq masked path, like eval traffic + make = lambda seed: _make_inputs(sum(lens), seed) # noqa: E731 + + q, k, v, g, beta = make(31337) + cu = _make_cu(lens) if lens else None + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, cu) + # Warmup compile of the optimized variants on the default stream so the + # streamed run below launches immediately (a ~100 s JIT inside the + # streamed region would let the sleep expire and hide the race). + _run(optimized, gate_params, q, k, v, g, beta, _make_cu(lens) if lens else None) + torch.cuda.synchronize() + + exec_stream = torch.cuda.Stream() # non-blocking, like the executor's + with torch.cuda.stream(exec_stream): + torch.cuda._sleep(1 << 28) # hold the stream back ~100 ms + # Produced on exec_stream, pending behind the sleep: a default-stream + # launch would read these buffers before they are written. + q2, k2, v2 = q * 1.0, k * 1.0, v * 1.0 + g2, beta2 = g * 1.0, beta * 1.0 + out_opt, state_opt = _run( + optimized, gate_params, q2, k2, v2, g2, beta2, _make_cu(lens) if lens else None + ) + # Consume on the execution stream immediately, exactly like the + # runtime's output-gate matmul and ssm_pool.index_copy_ do. + out_opt = out_opt * 1.0 + state_opt = state_opt * 1.0 + torch.cuda.synchronize() + + _assert_close(f"stream_{regime}/out", out_opt, out_ref) + _assert_close(f"stream_{regime}/state", state_opt, state_ref) + + +# --------------------------------------------------------------------------- +# 2. id()-keyed cache soundness under recycled ids / reused objects. +# --------------------------------------------------------------------------- + + +@torch.no_grad() +def test_recycled_id_varlen_pure_cache(dispatch_pair, gate_params): + """Aligned lens cache VARLEN_PURE=True under id(cu). Free cu, land a new + cu_seqlens on the recycled id with NON-aligned lens: a stale hit would + run the mask-free compile variant on partial chunks.""" + optimized, reference = dispatch_pair + mod = _op_module() + + aligned = [128, 256, 192] + cu_holder = [_make_cu(aligned)] + q, k, v, g, beta = _make_inputs(sum(aligned), seed=11) + _run(optimized, gate_params, q, k, v, g, beta, cu_holder[0]) + torch.cuda.synchronize() + assert mod._varlen_pure_cache.get(id(cu_holder[0])) is True, ( + "test setup: expected VARLEN_PURE=True cached under id(cu1)" + ) + + old_id, pruned = _release_cu(mod, cu_holder) + nonaligned = [100, 257, 219] # same n_seqs so the entry shape matches + if not pruned: + # Entry outlived our release: something still pins the tensor. That + # is sound ONLY if no new tensor can land on its id. + clash = _alloc_with_recycled_id(old_id, lambda: _make_cu(nonaligned), attempts=64) + assert clash is None, ( + "UNSOUND: id recycled while a stale _varlen_pure_cache entry for it is still present" + ) + pytest.skip("cu_seqlens still pinned externally; recycled-id scenario unreachable this run") + + cu2 = _alloc_with_recycled_id(old_id, lambda: _make_cu(nonaligned)) + if cu2 is None: + pytest.skip("could not obtain a recycled id for cu_seqlens") + q, k, v, g, beta = _make_inputs(sum(nonaligned), seed=12) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, cu2) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, _make_cu(nonaligned)) + _assert_close("recycled_pure/out", out_opt, out_ref) + _assert_close("recycled_pure/state", state_opt, state_ref) + assert mod._varlen_pure_cache.get(id(cu2)) is False + + +@torch.no_grad() +def test_recycled_id_single_seq_caches(dispatch_pair, gate_params): + """Phase 2.1 caches: a single ALIGNED seq caches (pure=True, seqlen=256) + under id(cu). Recycle the id with a single NON-aligned seq: a stale hit + would skip the g sentinel pad / use the wrong single-seq padding.""" + optimized, reference = dispatch_pair + mod = _op_module() + + cu_holder = [_make_cu([256])] + q, k, v, g, beta = _make_inputs(256, seed=21) + _run(optimized, gate_params, q, k, v, g, beta, cu_holder[0]) + torch.cuda.synchronize() + assert mod._varlen_pure_cache.get(id(cu_holder[0])) is True + + old_id, pruned = _release_cu(mod, cu_holder) + if not pruned: + clash = _alloc_with_recycled_id(old_id, lambda: _make_cu([300]), attempts=64) + assert clash is None, ( + "UNSOUND: id recycled while stale Phase 2.1 cache entries for it are still present" + ) + pytest.skip("cu_seqlens still pinned externally; recycled-id scenario unreachable this run") + assert old_id not in mod._varlen_single_seqlen_cache + + cu2 = _alloc_with_recycled_id(old_id, lambda: _make_cu([300])) + if cu2 is None: + pytest.skip("could not obtain a recycled id for cu_seqlens") + q, k, v, g, beta = _make_inputs(300, seed=22) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, cu2) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, _make_cu([300])) + _assert_close("recycled_single/out", out_opt, out_ref) + _assert_close("recycled_single/state", state_opt, state_ref) + + +@torch.no_grad() +def test_single_seq_same_object_repeated(dispatch_pair, gate_params): + """Two calls with the SAME cu_seqlens object and a non-aligned single + seq. The Phase 2.1 path must sentinel-pad g on BOTH calls; poisoning + _varlen_pure_cache with the per-call override made the second call skip + the pad while still compiling the mask-free variant.""" + optimized, reference = dispatch_pair + cu = _make_cu([300]) + for i in range(2): + q, k, v, g, beta = _make_inputs(300, seed=40 + i) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, cu) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, _make_cu([300])) + _assert_close(f"sameobj_call{i}/out", out_opt, out_ref) + _assert_close(f"sameobj_call{i}/state", state_opt, state_ref) + + +@torch.no_grad() +def test_recycled_id_input_wrap_cache(): + """_ct_cached wrappers are keyed by (id(tensor), etype); the wrapper pins + the tensor's storage, so entries for live tensors are immortal by + design. Soundness requires the entry to be popped when the keyed tensor + dies before its id can be recycled — otherwise a new tensor on the + recycled id would silently reuse a wrapper over freed storage.""" + mod = _op_module() + import cutlass + + t1 = torch.zeros(64, 64, dtype=torch.bfloat16, device="cuda") + w1 = mod._ct_cached(t1, cutlass.BFloat16) + assert mod._ct_cached(t1, cutlass.BFloat16) is w1 # cache hit path + old_id = id(t1) + key = (old_id, cutlass.BFloat16) + assert key in mod._input_wrap_cache + del t1 + gc.collect() + + if key in mod._input_wrap_cache: + # The wrapper (still referenced by the cache) pins the tensor object, + # so the entry survives — sound only if the id can NOT be recycled + # while the entry is present. Prove no new tensor lands on old_id. + t2 = _alloc_with_recycled_id( + old_id, + lambda: torch.zeros(64, 64, dtype=torch.bfloat16, device="cuda"), + attempts=64, + ) + assert t2 is None, ( + "UNSOUND: id recycled while a stale _input_wrap_cache entry for " + "it is still present — a new tensor would reuse a wrapper over " + "another tensor's storage" + ) + # Dropping the last wrapper ref must let the finalizer prune. + del w1 + mod._input_wrap_cache.pop(key, None) + gc.collect() + else: + # Finalizer fired at t1's dealloc — entry pruned before any recycle. + del w1 + + # Fresh tensor (recycled id or not) must get a fresh wrapper. + t3 = torch.zeros(64, 64, dtype=torch.bfloat16, device="cuda") + w3 = mod._ct_cached(t3, cutlass.BFloat16) + assert mod._input_wrap_cache[(id(t3), cutlass.BFloat16)] is w3 diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py new file mode 100644 index 000000000000..bc092f195b58 --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parity tests for the optimized Kimi K3 KDA prefill op.""" + +import pytest +import torch + +pytest.importorskip("fla") + +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import KimiKDALinearAttention # noqa: E402 + +NUM_HEADS = 96 +HEAD_DIM = 128 +CONV_KERNEL_SIZE = 4 +HIDDEN_SIZE = 7168 + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="Kimi K3 is supported only on Blackwell (SM100/SM103)", +) + + +def _make_attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttention]: + common = { + "hidden_size": HIDDEN_SIZE, + "num_heads": NUM_HEADS, + "head_dim": HEAD_DIM, + "conv_kernel_size": CONV_KERNEL_SIZE, + "use_full_rank_gate": True, + "gate_lower_bound": -5.0, + "rms_norm_eps": 1e-5, + "dtype": torch.bfloat16, + } + optimized = KimiKDALinearAttention(**common).to("cuda") + reference = KimiKDALinearAttention(**common, use_optimized_prefill=False).to("cuda") + reference.load_state_dict(optimized.state_dict()) + + assert optimized.prefill_kernel_path == "optimized" + assert reference.prefill_kernel_path == "fla" + assert reference.decode_kernel_path == optimized.decode_kernel_path + return optimized, reference + + +def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + actual_float = actual.float() + expected_float = expected.float() + cosine = torch.nn.functional.cosine_similarity( + actual_float.flatten(), expected_float.flatten(), dim=0 + ).item() + relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item() + assert cosine > 0.999 + assert relative_l2 < 3e-2 + + +@torch.no_grad() +def test_optimized_prefill_matches_fla_reference() -> None: + torch.manual_seed(0) + optimized, reference = _make_attention_pair() + + # Keep B=2 across a T transition: eqlen mBeta/mAqk/mAkk batch strides + # depend on T and therefore require distinct compiled kernel variants. + for batch_size, sequence_length in [(2, 256), (2, 512), (1, 1024)]: + hidden_states = ( + torch.randn( + batch_size, + sequence_length, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = optimized.forward_prefill(hidden_states) + expected = reference.forward_prefill(hidden_states) + _assert_close(actual, expected) + + hidden_states = torch.randn(1, 300, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + actual = optimized.forward_prefill(hidden_states) + expected = reference.forward_prefill(hidden_states) + _assert_close(actual, expected) + + sequence_lengths = [128, 256, 192] + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + assert optimized.prefill_kernel_source() + + +@torch.no_grad() +def test_kda_prefill_op_empty_token_batch(): + """T=0 call: no output rows, recurrent state passes through unchanged. + + The runtime can emit a context batch with an empty token payload + (observed under the overlap scheduler + logprobs flows). The op used + to raise ``RuntimeError: step must be nonzero`` from its + ``arange(step=T)`` buffer setup; the FLA fallback tolerates the call. + """ + num_heads, head_k, head_v = 4, 128, 128 + q = torch.empty(1, 0, num_heads, head_k, dtype=torch.bfloat16, device="cuda") + k = torch.empty_like(q) + g = torch.empty(1, 0, num_heads, head_k, dtype=torch.float32, device="cuda") + v = torch.empty(1, 0, num_heads, head_v, dtype=torch.bfloat16, device="cuda") + beta = torch.empty(1, 0, num_heads, dtype=torch.float32, device="cuda") + initial_state = torch.randn(1, num_heads, head_k, head_v, dtype=torch.float32, device="cuda") + + output, final_state = torch.ops.trtllm.kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=head_k**-0.5, + initial_state=initial_state, + output_final_state=True, + ) + assert output.shape == (1, 0, num_heads, head_v) + torch.testing.assert_close(final_state, initial_state) + # State must not alias the input (the caller copies it back into the + # pool the initial state may be a view of). + assert final_state.data_ptr() != initial_state.data_ptr() + + _, final_state_zero = torch.ops.trtllm.kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=head_k**-0.5, + initial_state=None, + output_final_state=True, + ) + torch.testing.assert_close(final_state_zero, torch.zeros_like(initial_state)) + + +@torch.no_grad() +def test_kda_prefill_op_empty_token_batch_variants(): + """Regression coverage for the T=0 guard's other entry shapes. + + - varlen (cu_seqlens present): n_seqs derives from cu_seqlens, not B + - output_final_state=False: the op's empty-tensor final-state fallback + - use_fused_k1234=True: the guard must precede the fused-path branch + """ + num_heads, head_k, head_v = 4, 128, 128 + q = torch.empty(1, 0, num_heads, head_k, dtype=torch.bfloat16, device="cuda") + k = torch.empty_like(q) + g = torch.empty(1, 0, num_heads, head_k, dtype=torch.float32, device="cuda") + v = torch.empty(1, 0, num_heads, head_v, dtype=torch.bfloat16, device="cuda") + beta = torch.empty(1, 0, num_heads, dtype=torch.float32, device="cuda") + common = dict(q=q, k=k, v=v, g=g, beta=beta, scale=head_k**-0.5) + + # Varlen: two zero-length sequences -> final_state per sequence. + cu_seqlens = torch.tensor([0, 0, 0], dtype=torch.long, device="cuda") + _, final_state = torch.ops.trtllm.kda_prefill( + **common, initial_state=None, output_final_state=True, cu_seqlens=cu_seqlens + ) + assert final_state.shape == (2, num_heads, head_k, head_v) + assert (final_state == 0).all() + + # output_final_state=False: op returns the empty placeholder tensor. + output, final_state = torch.ops.trtllm.kda_prefill( + **common, initial_state=None, output_final_state=False + ) + assert output.shape == (1, 0, num_heads, head_v) + assert final_state.numel() == 0 + + # Fused path: the guard must fire before _launch_fused_k1234. + output, final_state = torch.ops.trtllm.kda_prefill( + **common, initial_state=None, output_final_state=True, use_fused_k1234=True + ) + assert output.shape == (1, 0, num_heads, head_v) + assert final_state.shape == (1, num_heads, head_k, head_v) + + +@torch.no_grad() +def test_kda_mixer_empty_prefill(): + """Runtime-shaped regression: the mixer dispatch with an empty token + payload (a crashing call shape observed at runtime) must run + end-to-end on the optimized path.""" + optimized, _ = _make_attention_pair() + hidden_states = torch.empty(1, 0, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + out = optimized.forward_prefill(hidden_states) + assert out.shape == (1, 0, HIDDEN_SIZE) + + +@torch.no_grad() +def test_kda_prefill_op_partial_final_chunk_large_batch(): + """Regression: varlen batches whose FINAL chunk is partial. + + The chunk-tile kernels access the full 64-row tile of every chunk and + neutralize invalid rows only after the access, so the batch's final + partial chunk touches up to 63 rows past the logical packed length — + OOB reads on the beta input (now bounds-checked in fused_k123) and on + the A_kk/A_qk scratch (now allocated with one chunk of slack). The + runtime's autotuner-warmup shape [max_seq_len - 1, 1] = [8191, 1] hit + this as CUDA_ERROR_ILLEGAL_ADDRESS whenever the following page was + unmapped. + + - [8191, 1]: the exact autotuner-warmup composition (one max_seq_len-1 + context plus a 1-token remainder). + - [8000, 150, 42]: interior partial chunks (cross-sequence rows) plus + a partial final chunk, at eval-like scale. + """ + optimized, reference = _make_attention_pair() + for sequence_lengths in ([8191, 1], [8000, 150, 42]): + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +def test_kda_prefill_op_small_varlen_batch(): + """Small varlen batches (short-prompt contexts) through the dispatch. + + The persistent K123 scheduler needs >= 4 total chunks, so the dispatch + routes NT < 4 batches to the FLA path ([6,12], [1,2,3], [30] here); + the [6,12,20,25] case carries exactly NT=4 with total T < 64, running + the optimized masked path where building the eqlen chunk-offset + scratch used to raise ``step must be nonzero`` (arange step + ``T // 64 == 0``) — its output is parity-checked against FLA. + """ + optimized, reference = _make_attention_pair() + for sequence_lengths in ([6, 12], [1, 2, 3], [30], [6, 12, 20, 25]): + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +def test_kda_prefill_op_shape_growth_and_cu_dtype_transitions(): + """Cross-call transitions through one process's compile caches. + + Regression for the cu/ci-dtype cache-key bug: the K123/akk_inv compile + caches were keyed shape-independently but NOT on the cu_seqlens / + chunk_indices dtype, while the compiled kernels bake the element type + (int64 reads use stride 8, int32 stride 4). Reusing an int64-compiled + kernel on int32 cu/ci misaddressed every cu/ci element — garbage seq + ids / chunk starts -> cudaErrorIllegalAddress on the first call after + the flip (memcheck: 4-byte read one element past the 2-entry int32 cu); + the reverse direction (int32-compiled, int64 passed) corrupted + silently. Shape growth alone (same dtype) was already sound. + + The sequence below covers, in one process: buffer-cache growth + (T 1171 -> 8191), int64 -> int32 flip on the grown shape, shrink with + a flip back, and a multi-seq int32 batch. Every call is parity-checked + against FLA (catches the silent-corruption direction too). + """ + torch.manual_seed(0) + optimized, reference = _make_attention_pair() + cases = [ + ([517, 654], torch.long), # small batch, int64 cu (dump-replay-like) + ([8191], torch.int32), # buffer growth + dtype flip (crashed pre-fix) + ([1171], torch.long), # shrink + flip back (silent corruption pre-fix) + ([150, 900, 333, 640], torch.int32), # multi-seq int32 + ] + for sequence_lengths, cu_dtype in cases: + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=cu_dtype, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py new file mode 100644 index 000000000000..b712bbbe760a --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py @@ -0,0 +1,397 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""State parity tests: optimized KDA prefill dispatch vs FLA, kernel level. + +Complements test_kda_prefill_op.py (module outputs, no cache) with the cases +the executor runtime actually exercises through KDAKernelDispatch: + + * final_state parity in the pool's V-first [N, H, V, K] layout — K == V == + 128 for Kimi K3, so a layout mix-up is invisible to shape checks and only + caught numerically; + * a LARGE carried initial_state — the K4 sign error fixed by dev-tech + commit e45ae259 (NV = U - W@S, was U + W@S) is masked by tiny/zero + initial states and random inputs, and only blows up when a real-magnitude + state is carried in; + * non-64-aligned varlen (single-seq pad path and multi-seq masked path); + * fresh input tensors on every call — the runtime never reuses tensor + objects, so this exercises per-call wrapper rebuild instead of the + benchmark-style stable-object path, and would catch both stale-cache + corruption and per-call pinning leaks. +""" + +import pytest +import torch + +pytest.importorskip("fla") + +from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import KDAKernelDispatch # noqa: E402 + +NUM_HEADS = 96 +HEAD_K_DIM = 128 +LOWER_BOUND = -5.0 + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} + + +pytestmark = pytest.mark.skipif( + not _has_supported_gpu(), + reason="Kimi K3 is supported only on Blackwell (SM100/SM103)", +) + + +@pytest.fixture(scope="module") +def dispatch_pair(): + optimized = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False) + assert optimized.prefill_kernel_path == "optimized" + reference = KDAKernelDispatch(use_optimized_prefill=False, use_optimized_decode=False) + assert reference.prefill_kernel_path == "fla" + return optimized, reference + + +@pytest.fixture(scope="module") +def gate_params(): + torch.manual_seed(0) + a_log = torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.5 + dt_bias = torch.randn(NUM_HEADS * HEAD_K_DIM, dtype=torch.float32, device="cuda") * 0.1 + return a_log, dt_bias + + +def _make_inputs(batch: int, total_t: int, h0_scale, seed: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + h, k = NUM_HEADS, HEAD_K_DIM + + def rnd(*shape, dtype=torch.bfloat16, scale=1.0): + return ( + torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) * scale + ) + + q = rnd(batch, total_t, h, k) + key = rnd(batch, total_t, h, k) + v = rnd(batch, total_t, h, k) + g = rnd(batch, total_t, h, k) + beta = rnd(batch, total_t, h, dtype=torch.float32) + h0 = None + if h0_scale is not None: + # Pool V-first [N, H, V, K] layout, deliberately non-symmetric in the + # last two dims (row-index modulation) so a [K,V]<->[V,K] transpose + # mix-up cannot cancel out. + h0 = rnd(batch, h, k, k, dtype=torch.float32, scale=h0_scale) + h0 = h0 * torch.linspace(0.5, 1.5, k, device="cuda").view(1, 1, k, 1) + return q, key, v, g, beta, h0 + + +def _run(dispatch, gate_params, q, k, v, g, beta, h0, cu): + a_log, dt_bias = gate_params + return dispatch.prefill_chunk_kda( + q=q.clone(), + k=k.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + A_log=a_log, + dt_bias=dt_bias, + scale=HEAD_K_DIM**-0.5, + initial_state=h0.clone() if h0 is not None else None, + safe_gate=True, + lower_bound=LOWER_BOUND, + cu_seqlens=cu.clone() if cu is not None else None, + ) + + +def _assert_close(name, actual, expected): + actual, expected = actual.float(), expected.float() + cos = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0).item() + rel = ((actual - expected).norm() / (expected.norm() + 1e-12)).item() + assert cos > 0.999 and rel < 3e-2, f"{name}: cos={cos:.6f} rel_l2={rel:.3e}" + + +CASES = [ + # (label, batch, eqlen T or None, varlen seq lens or None, h0_scale) + ("eqlen_b2_t256_no_state", 2, 256, None, None), + ("eqlen_b2_t256_large_state", 2, 256, None, 1.0), + ("eqlen_b1_t300_pad_large_state", 1, 300, None, 1.0), + ("varlen_aligned_large_state", 1, None, [128, 256, 192], 1.0), + ("varlen_nonaligned_large_state", 1, None, [100, 257, 64], 1.0), + ("varlen_single_nonaligned_large_state", 1, None, [300], 1.0), +] + + +@pytest.mark.parametrize("label,batch,eqlen_t,lens,h0_scale", CASES, ids=[c[0] for c in CASES]) +@torch.no_grad() +def test_state_parity(dispatch_pair, gate_params, label, batch, eqlen_t, lens, h0_scale): + optimized, reference = dispatch_pair + if lens is not None: + total_t = sum(lens) + cu = torch.tensor( + [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" + ) + n_seqs = len(lens) + else: + total_t, cu, n_seqs = eqlen_t, None, batch + case_idx = [c[0] for c in CASES].index(label) + q, k, v, g, beta, h0 = _make_inputs( + 1 if lens is not None else batch, total_t, h0_scale, seed=100 + case_idx + ) + if h0 is not None and lens is not None: + h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( + 0.5, 1.5, n_seqs, device="cuda" + ).view(n_seqs, 1, 1, 1) + + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu) + _assert_close(f"{label}/out", out_opt, out_ref) + _assert_close(f"{label}/state", state_opt, state_ref) + + +@torch.no_grad() +def test_fresh_tensors_every_call(dispatch_pair, gate_params): + """Same shape, new tensor objects + contents per call. Catches stale + id-keyed cache hits (wrong data); a persistently climbing allocation + here would indicate per-call activation pinning.""" + optimized, reference = dispatch_pair + for i in range(8): + q, k, v, g, beta, h0 = _make_inputs(2, 256, 1.0, seed=1000 + i) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, None) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, None) + _assert_close(f"iter{i}/out", out_opt, out_ref) + _assert_close(f"iter{i}/state", state_opt, state_ref) + + +# Eval-scale coverage (2026-07-23): GSM8K on the optimized prefill scored +# 68.99 vs 97.01 on the FLA path while every existing unit test passed — +# the gap is the eval regime: varlen batches packing dozens of +# mixed-length sequences, and chunked prefill chaining a sequence's state +# across calls. These tests parity-check exactly those two regimes. + +_EVAL_SCALE_LENS = [(97 * (i + 3)) % 911 + 45 for i in range(24)] + + +@pytest.mark.parametrize("h0_scale", [None, 1.0], ids=["no_state", "state"]) +@torch.no_grad() +def test_eval_scale_packed_varlen_parity(dispatch_pair, gate_params, h0_scale): + """24 packed mixed-length sequences (45..955 tokens), like one + max_num_tokens=8192 eval context batch.""" + optimized, reference = dispatch_pair + lens = _EVAL_SCALE_LENS + n_seqs, total_t = len(lens), sum(lens) + cu = torch.tensor( + [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" + ) + q, k, v, g, beta, h0 = _make_inputs(1, total_t, h0_scale, seed=4242) + if h0 is not None: + h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( + 0.5, 1.5, n_seqs, device="cuda" + ).view(n_seqs, 1, 1, 1) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu) + _assert_close("evalscale/out", out_opt, out_ref) + _assert_close("evalscale/state", state_opt, state_ref) + + +# GSM8K-shaped coverage (2026-07-23, second regression pass): the stream +# fix lifted GSM8K 68.99 -> 77.79, still below the 97.01 FLA control, and +# the partial-score trajectory declined with completion length — the +# residual defect skews toward LONG sequences (8-shot GSM8K prompts are +# ~1-1.2k tokens = 16-20 chunks/seq; the eval-scale test above caps at +# 955). These cases pin per-seq NT at 16-20. The weak-gate variants +# (g_scale=0.05) reduce the per-chunk decay so state accumulates across +# many chunks — accumulation-order/precision defects that strong random +# gates (effectively local memory) cannot surface. + +_LONG_MIXED_LENS = [1150, 1200, 980, 1100, 1279, 1024, 1216, 1090] + + +@pytest.mark.parametrize("g_scale", [1.0, 0.05], ids=["g_normal", "g_weak"]) +@pytest.mark.parametrize( + "lens,h0_scale", + [([1150], None), ([1150], 1.0), (_LONG_MIXED_LENS, None), (_LONG_MIXED_LENS, 1.0)], + ids=["single_long", "single_long_state", "mixed_long", "mixed_long_state"], +) +@torch.no_grad() +def test_long_sequence_parity(dispatch_pair, gate_params, lens, h0_scale, g_scale): + optimized, reference = dispatch_pair + n_seqs, total_t = len(lens), sum(lens) + cu = torch.tensor( + [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" + ) + q, k, v, g, beta, h0 = _make_inputs(1, total_t, h0_scale, seed=9000 + len(lens)) + g = (g.float() * g_scale).to(g.dtype) + if h0 is not None: + h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( + 0.5, 1.5, n_seqs, device="cuda" + ).view(n_seqs, 1, 1, 1) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu) + _assert_close("long/out", out_opt, out_ref) + _assert_close("long/state", state_opt, state_ref) + + +# Poisoned-scratch FINAL-STATE parity (2026-07-23, review follow-up on the +# fused_k123 beta-guard lines ~1520): the store guard (t < ci_eos) leaves +# k_scaled/kg/q_scaled scratch rows past each seq's end UNWRITTEN for +# partial final chunks. K4 is supposed to neutralize those rows via its +# per-seq bounded TMA (token extent = cu[s+1], hardware zero-fill); these +# cases PROVE that with garbage actually present in the stale rows: a +# preceding call with the SAME buffer-cache key (same B/T/NT/n_seqs) and +# huge-magnitude inputs fills the scratch, then the victim call's final +# state is asserted against FLA. Every eval prompt ends in a partial +# chunk and the state feeds all decode steps, so state corruption here is +# invisible to output-only checks. +# +# Victim/poison pairs share (T_total, NT_total, n_seqs) so they hit the +# same _buf_cache entry; the poison split is chosen so its guarded stores +# leave different rows written than the victim's layout expects. + +_POISON_STATE_CASES = [ + # (victim lens, poison lens) — equal sum and equal total chunk count + ([8191, 1], [8127, 65]), + ([8000, 150, 42], [7999, 129, 64]), + ([1186, 30], [1150, 66]), +] + + +def _nt(lens): + return sum((n + 63) // 64 for n in lens) + + +@pytest.mark.parametrize( + "victim,poison", _POISON_STATE_CASES, ids=["v8191_1", "v8000_150_42", "v1186_30"] +) +@torch.no_grad() +def test_final_state_parity_poisoned_scratch(dispatch_pair, gate_params, victim, poison): + optimized, reference = dispatch_pair + assert ( + sum(victim) == sum(poison) and _nt(victim) == _nt(poison) and len(victim) == len(poison) + ), "pairs must share the buffer-cache key" + total_t = sum(victim) + + # Poison pass: same buffer-cache key, huge magnitudes. + q, k, v, g, beta, _ = _make_inputs(1, total_t, None, seed=31) + cu_p = torch.tensor( + [0] + torch.cumsum(torch.tensor(poison), 0).tolist(), dtype=torch.long, device="cuda" + ) + _run(optimized, gate_params, q * 100, k * 100, v * 1000, g, beta * 10, None, cu_p) + torch.cuda.synchronize() + + # Victim pass: fresh inputs, partial final chunks per seq. + q, k, v, g, beta, h0 = _make_inputs(1, total_t, 1.0, seed=32) + n_seqs = len(victim) + h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( + 0.5, 1.5, n_seqs, device="cuda" + ).view(n_seqs, 1, 1, 1) + cu_v = torch.tensor( + [0] + torch.cumsum(torch.tensor(victim), 0).tolist(), dtype=torch.long, device="cuda" + ) + out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu_v) + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu_v) + _assert_close("poisoned/out", out_opt, out_ref) + _assert_close("poisoned/state", state_opt, state_ref) + + +def _run_headcount_case(dispatch_pair, heads): + optimized, reference = dispatch_pair + h, k = heads, HEAD_K_DIM + gen = torch.Generator(device="cuda").manual_seed(4711 + heads) + a_log = torch.randn(h, generator=gen, dtype=torch.float32, device="cuda") * 0.5 + dt_bias = torch.randn(h * k, generator=gen, dtype=torch.float32, device="cuda") * 0.1 + lens = [1150, 731, 1024, 987] + total_t = sum(lens) + cu = torch.tensor( + [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" + ) + + def rnd(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) + + q = rnd(1, total_t, h, k) + key = rnd(1, total_t, h, k) + v = rnd(1, total_t, h, k) + g = rnd(1, total_t, h, k) + beta = rnd(1, total_t, h, dtype=torch.float32) + h0 = rnd(len(lens), h, k, k, dtype=torch.float32) * torch.linspace( + 0.5, 1.5, k, device="cuda" + ).view(1, 1, k, 1) + + def run(dispatch): + return dispatch.prefill_chunk_kda( + q=q.clone(), + k=key.clone(), + v=v.clone(), + g=g.clone(), + beta=beta.clone(), + A_log=a_log, + dt_bias=dt_bias, + scale=k**-0.5, + initial_state=h0.clone(), + safe_gate=True, + lower_bound=LOWER_BOUND, + cu_seqlens=cu.clone(), + ) + + out_opt, state_opt = run(optimized) + out_ref, state_ref = run(reference) + _assert_close(f"h{heads}/out", out_opt, out_ref) + _assert_close(f"h{heads}/state", state_opt, state_ref) + + +@pytest.mark.parametrize("heads", [6, 12], ids=["tp16_heads6", "tp8_heads12"]) +@torch.no_grad() +def test_long_sequence_parity_tp_headcount(dispatch_pair, heads): + """Per-rank head count under tensor parallelism: the e2e eval runs + tp_size=16 -> H = 96/16 = 6 heads per rank, but every other test here + compiles the DSL kernels at H=96 (H is a compile-time specializer). + Parity-check the H=6 / H=12 compiles on a GSM8K-shaped batch.""" + _run_headcount_case(dispatch_pair, heads) + + +@torch.no_grad() +def test_headcount_recompile_parity(dispatch_pair): + """Cross-head-count compile-cache isolation, order-explicit: run two + different per-rank head counts back-to-back in the same process and + parity-check the SECOND one. Regression test for the K4 persistent + kernel cache key losing H (0e44bf64a6 follow-up): s_ct bakes the + [H, K, V] state shape/strides at compile time, so reusing the + first-compiled head count's kernel for a different H misaddresses + every (seq, head) state tile. The tp_headcount params above only catch + this via pytest execution order; this test pins the order even when + run in isolation. Heads 8 then 4 avoid cache hits from other tests.""" + _run_headcount_case(dispatch_pair, 8) + _run_headcount_case(dispatch_pair, 4) + + +@torch.no_grad() +def test_chunked_continuation_parity(dispatch_pair, gate_params): + """Chunked prefill: run [3000] as [2048] then [952] with the first + call's final state carried into the second (what the runtime's + chunked-prefill continuation does through the pool), and compare the + second chunk's output and final state against a single full-sequence + reference run.""" + optimized, reference = dispatch_pair + total_t, split = 3000, 2048 + q, k, v, g, beta, _ = _make_inputs(1, total_t, None, seed=777) + cu_full = torch.tensor([0, total_t], dtype=torch.long, device="cuda") + out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, None, cu_full) + + cu1 = torch.tensor([0, split], dtype=torch.long, device="cuda") + cu2 = torch.tensor([0, total_t - split], dtype=torch.long, device="cuda") + sl1 = slice(0, split) + sl2 = slice(split, total_t) + out1, state1 = _run( + optimized, gate_params, q[:, sl1], k[:, sl1], v[:, sl1], g[:, sl1], beta[:, sl1], None, cu1 + ) + out2, state2 = _run( + optimized, + gate_params, + q[:, sl2], + k[:, sl2], + v[:, sl2], + g[:, sl2], + beta[:, sl2], + state1, + cu2, + ) + _assert_close("chunked/out_chunk1", out1, out_ref[:, sl1]) + _assert_close("chunked/out_chunk2", out2, out_ref[:, sl2]) + _assert_close("chunked/state", state2, state_ref) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py new file mode 100644 index 000000000000..9b412e1acf2e --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py @@ -0,0 +1,689 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 native TRTLLM-Gen SiTU MoE tests. + +Covers the acceptance criteria of the SiTU cubin integration plan +(`tensorrt_llm/_torch/modules/kimi_k3_moe/SITU_CUBIN_INTEGRATION_PLAN.md`): + +* runner-local ActType numeric stability (SwiGlu/Relu2/Silu unchanged, + SiTu appended); +* the native SiTU runner returns valid tactics and actually launches a + kernel whose name contains ``siTuGlu``; +* fused output matches the Python ``SituAndMul`` reference within + MXFP8/MXFP4 quantization tolerance, for the default AND asymmetric + non-default ``activation_situ_beta`` / ``activation_situ_linear_beta``; +* swapping the FC1 halves (gate-first packing) breaks accuracy + (mutation test — proves the w3-first convention is load-bearing); +* module-level semantics (latent projections, RMSNorm, shared experts) + survive the fused path; +* the fused path fails loudly without loaded weights (no silent + random-weight fallback). +""" + +import dataclasses +import os +from types import SimpleNamespace + +import pytest +import torch +from utils.util import check_accuracy + +from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory +from tensorrt_llm._torch.modules.kimi_k3_moe import KimiK3SparseMoeBlock +from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import ( + is_native_situ_supported, + make_situ_alpha_beta, + padded_fused_shapes, +) +from tensorrt_llm._torch.modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate +from tensorrt_llm._torch.utils import ActType_TrtllmGen + +situ_supported = pytest.mark.skipif( + not is_native_situ_supported(), + reason="native SiTU cubins require SM100/SM103 (Blackwell)", +) + + +@dataclasses.dataclass +class _K3Config: + """Minimal config carrying the fields KimiK3SparseMoeBlock reads.""" + + hidden_size: int = 512 + num_experts: int = 8 + num_experts_per_token: int = 2 + moe_intermediate_size: int = 256 + moe_renormalize: bool = True + moe_router_activation_func: str = "sigmoid" + routed_scaling_factor: float = 1.0 + num_expert_group: int = 1 + topk_group: int = 1 + num_shared_experts: int = None + routed_expert_hidden_size: int = None + latent_moe_use_norm: bool = False + rms_norm_eps: float = 1e-5 + activation_situ_beta: float = 4.0 + activation_situ_linear_beta: float = 25.0 + + +def _init_block_weights(block: KimiK3SparseMoeBlock, seed: int = 1234): + """Fill gate/latent/shared weights and the MXFP4 expert bank.""" + gen = torch.Generator(device="cpu").manual_seed(seed) + + def randn_like_param(p, scale): + return torch.randn(p.shape, generator=gen, dtype=torch.float32) * scale + + with torch.no_grad(): + block.gate.weight.copy_(randn_like_param(block.gate.weight, 0.05)) + block.gate.e_score_correction_bias.copy_( + randn_like_param(block.gate.e_score_correction_bias, 0.1) + ) + if block.use_latent_moe: + block.routed_expert_down_proj.weight.copy_( + randn_like_param(block.routed_expert_down_proj.weight, 0.05).to( + block.routed_expert_down_proj.weight.dtype + ) + ) + block.routed_expert_up_proj.weight.copy_( + randn_like_param(block.routed_expert_up_proj.weight, 0.05).to( + block.routed_expert_up_proj.weight.dtype + ) + ) + if block.routed_expert_norm is not None: + block.routed_expert_norm.weight.fill_(1.0) + if block.shared_experts is not None: + block.shared_experts.gate_up_proj.weight.copy_( + randn_like_param(block.shared_experts.gate_up_proj.weight, 0.05).to( + block.shared_experts.gate_up_proj.weight.dtype + ) + ) + block.shared_experts.down_proj.weight.copy_( + randn_like_param(block.shared_experts.down_proj.weight, 0.05).to( + block.shared_experts.down_proj.weight.dtype + ) + ) + + isize, hsize = block.expert_bank.intermediate_size, block.expert_bank.hidden_size + for e in range(block.num_experts): + w1 = torch.randn(isize, hsize, generator=gen, dtype=torch.float32) * 0.1 + w2 = torch.randn(hsize, isize, generator=gen, dtype=torch.float32) * 0.1 + w3 = torch.randn(isize, hsize, generator=gen, dtype=torch.float32) * 0.1 + block.expert_bank.store_expert(e, w1, w2, w3) + + +def _make_block_pair(config, device): + """Fused block + reference block sharing identical weights.""" + fused = KimiK3SparseMoeBlock( + config, use_fused_cubin=True, dtype=torch.bfloat16, device=device + ).to(device) + ref = KimiK3SparseMoeBlock( + config, use_fused_cubin=False, dtype=torch.bfloat16, device=device + ).to(device) + _init_block_weights(fused) + ref.load_state_dict(fused.state_dict()) + fused.build_fused_weights() + return fused, ref + + +def test_act_type_enum_values_stable(): + assert int(ActType_TrtllmGen.SwiGlu) == 0 + assert int(ActType_TrtllmGen.Relu2) == 1 + assert int(ActType_TrtllmGen.Silu) == 2 + assert int(ActType_TrtllmGen.SiTu) == 3 + + +def test_padded_fused_shapes(): + assert padded_fused_shapes(512, 256) == (512, 512, 256) + assert padded_fused_shapes(128, 256) == (512, 128, 256) + assert padded_fused_shapes(2880, 96) == (3072, 2944, 128) + + +def test_kimi_gate_reuses_deepseek_v3_routing(): + config = _K3Config(num_experts=16, num_experts_per_token=4) + gate = KimiK3MoEGate(config) + torch.manual_seed(23) + with torch.no_grad(): + gate.weight.normal_(std=0.1) + gate.e_score_correction_bias.normal_(std=0.05) + hidden_states = torch.randn(2, 7, config.hidden_size) + + expected_ids, expected_weights = gate(hidden_states) + routing_method = gate.routing_method + # Exercise the portable PyTorch short path; the production path keeps + # is_fused=True and uses the same routing contract. + routing_method.routing_impl.is_fused = False + actual_ids, actual_weights = routing_method.apply(gate.compute_logits(hidden_states)) + + expected_order = expected_ids.argsort(dim=-1) + actual_order = actual_ids.argsort(dim=-1) + assert actual_ids.dtype == torch.int32 + torch.testing.assert_close( + expected_ids.gather(1, expected_order).to(actual_ids.dtype), + actual_ids.gather(1, actual_order), + ) + torch.testing.assert_close( + expected_weights.gather(1, expected_order), + actual_weights.gather(1, actual_order), + ) + + +def test_communication_factory_accepts_model_selected_method(monkeypatch): + mapping = SimpleNamespace( + enable_attention_dp=True, + dp_size=16, + moe_tp_size=1, + moe_ep_size=16, + ) + model_config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(hidden_size=3584), + torch_dtype=torch.bfloat16, + quant_config=None, + max_num_tokens=4096, + moe_max_num_tokens=65536, + use_cuda_graph=False, + use_low_precision_moe_combine=False, + ) + selected = object() + method = None + + def create_forced_method(force_method, *args, **kwargs): + nonlocal method + method = force_method + return selected + + monkeypatch.delenv("TRTLLM_FORCE_COMM_METHOD", raising=False) + monkeypatch.setattr( + CommunicationFactory, + "_create_forced_method", + staticmethod(create_forced_method), + ) + actual = CommunicationFactory.create_strategy( + model_config=model_config, + num_experts=896, + num_slots=896, + top_k=16, + expert_size_per_partition=56, + hidden_size=3584, + communication_method="ALLGATHER", + ) + + assert method == "ALLGATHER" + assert actual is selected + + +@situ_supported +def test_make_situ_alpha_beta_contract(): + alpha, beta = make_situ_alpha_beta( + local_num_experts=8, + situ_beta=4.0, + situ_linear_beta=25.0, + device=torch.device("cuda"), + ) + for buf, val in ((alpha, 4.0), (beta, 25.0)): + assert buf.is_cuda and buf.dtype == torch.float32 and buf.is_contiguous() + assert buf.shape == (8,) + assert torch.all(buf == val) + with pytest.raises(RuntimeError, match="must be > 0"): + make_situ_alpha_beta( + local_num_experts=8, + situ_beta=-1.0, + situ_linear_beta=25.0, + device=torch.device("cuda"), + ) + with pytest.raises(RuntimeError, match="must be > 0"): + make_situ_alpha_beta( + local_num_experts=8, + situ_beta=4.0, + situ_linear_beta=0.0, + device=torch.device("cuda"), + ) + + +@situ_supported +def test_situ_runner_returns_valid_tactics(): + runner = torch.classes.trtllm.MxE4m3MxE2m1BlockScaleMoERunner(int(ActType_TrtllmGen.SiTu), True) + # Representative low-latency and throughput shapes (topK, hidden, + # intermediate, localExperts, numTokens, validHidden, validIntermediate). + for num_tokens in (1, 8, 512): + tactics = runner.get_valid_configs(2, 512, 256, 8, num_tokens, 512, 256) + assert len(tactics) > 0, f"no valid SiTU tactic for num_tokens={num_tokens}" + + +_LAUNCH_EVIDENCE_SCRIPT = r""" +import torch +from test_kimi_k3_situ_moe import _K3Config, _make_block_pair + +device = torch.device("cuda") +config = _K3Config() +fused, _ = _make_block_pair(config, device) +x = torch.randn(1, 16, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 +fused(x) +torch.cuda.synchronize() +assert fused._cubin_call_count == 1 +""" + + +@situ_supported +def test_fused_forward_launches_situ_kernel(): + """Launch evidence: the FC1 kernel actually selected must be a siTuGlu cubin. + + Runs in a subprocess because the C++ logger level is fixed at process + start (TLLM_LOG_LEVEL) and TLLM_BATCHED_GEMM_PRINT_NAME logs at INFO. + """ + import subprocess + import sys + + env = dict(os.environ) + env["TLLM_BATCHED_GEMM_PRINT_NAME"] = "1" + env["TLLM_LOG_LEVEL"] = "INFO" + this_dir = os.path.dirname(os.path.abspath(__file__)) + unittest_root = os.path.abspath(os.path.join(this_dir, "..", "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([this_dir, unittest_root, env.get("PYTHONPATH", "")]) + result = subprocess.run( + [sys.executable, "-c", _LAUNCH_EVIDENCE_SCRIPT], + capture_output=True, + text=True, + env=env, + cwd=os.path.dirname(os.path.abspath(__file__)), + timeout=600, + ) + log = result.stdout + result.stderr + assert result.returncode == 0, f"fused forward failed:\n{log[-4000:]}" + assert "siTuGlu" in log, ( + "expected the FC1 launch log to name a siTuGlu kernel; got:\n" + log[-4000:] + ) + + +@situ_supported +@pytest.mark.parametrize( + "situ_beta,situ_linear_beta", + [ + (4.0, 25.0), # Kimi K3 defaults + (2.5, 7.0), # asymmetric non-defaults — catches alpha/beta swaps + ], + ids=["default_alpha_beta", "asymmetric_alpha_beta"], +) +@pytest.mark.parametrize("num_tokens", [1, 15, 256], ids=lambda n: f"tokens{n}") +def test_fused_matches_reference(num_tokens, situ_beta, situ_linear_beta): + device = torch.device("cuda") + config = _K3Config( + activation_situ_beta=situ_beta, + activation_situ_linear_beta=situ_linear_beta, + ) + fused, ref = _make_block_pair(config, device) + + torch.manual_seed(7) + x = torch.randn(1, num_tokens, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 + out_fused = fused(x) + out_ref = ref(x) + + # Error budget: MXFP8 activation quantization (FC1 input and FC1->FC2 + # intermediate) on top of shared canonical MXFP4 weights. + check_accuracy(out_fused, out_ref, atol=0.1, rtol=0.15, percent=0.95) + + +@situ_supported +def test_fused_matches_reference_with_latent_and_shared_experts(): + device = torch.device("cuda") + config = _K3Config( + hidden_size=512, + routed_expert_hidden_size=256, + latent_moe_use_norm=True, + num_shared_experts=2, + ) + fused, ref = _make_block_pair(config, device) + + torch.manual_seed(11) + x = torch.randn(2, 33, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 + out_fused = fused(x) + out_ref = ref(x) + assert out_fused.shape == x.shape and out_fused.dtype == x.dtype + check_accuracy(out_fused, out_ref, atol=0.1, rtol=0.15, percent=0.95) + + +@situ_supported +def test_fc1_swap_mutation_breaks_accuracy(): + """Packing gate-first (HF order) must break the numerics. + + Guards the w3-first packing convention: if the kernel accepted either + order, this mutation would silently pass and the convention assert in + the docs would be untestable. + """ + device = torch.device("cuda") + config = _K3Config() + fused, ref = _make_block_pair(config, device) + + # Rebuild the fused buffers with w1/w3 swapped. + from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import pack_routed_expert_weights + + swapped = pack_routed_expert_weights( + w1_packed=fused.expert_bank.w3_packed, + w1_scales=fused.expert_bank.w3_scales, + w3_packed=fused.expert_bank.w1_packed, + w3_scales=fused.expert_bank.w1_scales, + w2_packed=fused.expert_bank.w2_packed, + w2_scales=fused.expert_bank.w2_scales, + device=device, + ) + fused.gemm1_weights = swapped["gemm1_weights"] + fused.gemm1_weights_scale = swapped["gemm1_weights_scale"] + + torch.manual_seed(13) + x = torch.randn(1, 64, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 + out_fused = fused(x) + out_ref = ref(x) + with pytest.raises(Exception, match="Mismatch percentage"): + check_accuracy(out_fused, out_ref, atol=0.1, rtol=0.15, percent=0.95) + + +@situ_supported +def test_swiglu_act_mutation_breaks_accuracy(): + """Running the same weights through SwiGlu kernels must not match SiTU.""" + device = torch.device("cuda") + config = _K3Config() + fused, ref = _make_block_pair(config, device) + + import tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels as mk + + torch.manual_seed(17) + x = torch.randn(1, 64, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 + + orig = mk.invoke_native_situ_moe + + def swiglu_invoke(**kwargs): + kwargs["act_type"] = int(ActType_TrtllmGen.SwiGlu) + return orig(**kwargs) + + from tensorrt_llm._torch.modules import kimi_k3_moe + + kimi_k3_moe.kimi_k3_moe_block.invoke_native_situ_moe = swiglu_invoke + try: + out_fused = fused(x) + finally: + kimi_k3_moe.kimi_k3_moe_block.invoke_native_situ_moe = orig + + out_ref = ref(x) + with pytest.raises(Exception, match="Mismatch percentage"): + check_accuracy(out_fused, out_ref, atol=0.1, rtol=0.15, percent=0.95) + + +@situ_supported +def test_fused_forward_without_weights_raises(): + device = torch.device("cuda") + config = _K3Config() + block = KimiK3SparseMoeBlock( + config, use_fused_cubin=True, dtype=torch.bfloat16, device=device + ).to(device) + _init_block_weights(block) # bank filled, but fused buffers NOT built + x = torch.randn(1, 4, config.hidden_size, dtype=torch.bfloat16, device=device) + with pytest.raises(RuntimeError, match="fused weights were never built"): + block(x) + + +# --------------------------------------------------------------------------- +# Routed MoE TP/EP split selection (CPU-only). +# --------------------------------------------------------------------------- + + +def test_mapping_records_moe_tp_ep_user_specified(): + from tensorrt_llm.mapping import Mapping + + # Auto default: -1 sentinels resolve to (moe_tp=tp, moe_ep=1) but must + # NOT be flagged as a user request. + auto = Mapping(world_size=8, tp_size=8) + assert auto.moe_tp_size == 8 and auto.moe_ep_size == 1 + assert not auto.moe_tp_ep_user_specified + + tp = Mapping(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1) + assert tp.moe_tp_ep_user_specified + + # Setting only one side still counts as explicit. + ep = Mapping(world_size=8, tp_size=8, moe_ep_size=8) + assert ep.moe_tp_ep_user_specified + assert ep.moe_tp_size == 1 and ep.moe_ep_size == 8 + + +def test_kimi_k3_moe_split_selection(monkeypatch): + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _K3_MOE_EP_ENV, + _K3_MOE_TP_ENV, + KimiK3MoERuntime, + ) + from tensorrt_llm.mapping import Mapping + + monkeypatch.delenv(_K3_MOE_TP_ENV, raising=False) + monkeypatch.delenv(_K3_MOE_EP_ENV, raising=False) + + # Auto mapping default stays EP-only (the historical K3 layout), even + # though the resolved mapping says moe_tp=8. + auto = Mapping(world_size=8, tp_size=8) + assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (1, 8) + + # Explicit pure-TP and hybrid requests are honored. + tp = Mapping(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1) + assert KimiK3MoERuntime._select_moe_tp_ep(tp) == (8, 1) + tep = Mapping(world_size=8, tp_size=8, moe_tp_size=4, moe_ep_size=2) + assert KimiK3MoERuntime._select_moe_tp_ep(tep) == (4, 2) + + # Env override wins; a single side derives the other from tp_size. + monkeypatch.setenv(_K3_MOE_TP_ENV, "4") + assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) + monkeypatch.delenv(_K3_MOE_TP_ENV) + monkeypatch.setenv(_K3_MOE_EP_ENV, "2") + assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) + + +# --------------------------------------------------------------------------- +# MoE tensor-parallel shard parity (ConfigurableMoE / TRTLLM-Gen, GPU). +# +# Production K3 TP8 geometry per rank: ALL experts, intermediate 3072/8=384, +# latent hidden 3584, group-32 packed MXFP4 weights column-sharded (w1/w3) +# and row-sharded (w2) by the stock TRTLLM-Gen quant-method loaders. These +# tests run the identical shard shapes on ONE GPU by loading each simulated +# rank through the real `load_packed_mxfp4_expert` path with a proxy module +# exposing (tp_size, tp_rank), then summing the per-rank partial outputs. +# --------------------------------------------------------------------------- + +_TP_HIDDEN = 3584 +_TP_INTERMEDIATE = 3072 +_TP_EXPERTS = 8 +_TP_TOPK = 2 + + +def _make_packed_expert_bank(num_experts, intermediate, hidden, seed=101): + """Random group-32 packed MXFP4 tensors in checkpoint layout (uint8).""" + gen = torch.Generator().manual_seed(seed) + + def nibbles(*shape): + return torch.randint(0, 256, shape, generator=gen, dtype=torch.uint8) + + def scales(*shape): + # UE8M0 exponents 2^-9..2^-4 keep bf16 outputs well-conditioned. + return torch.randint(118, 124, shape, generator=gen, dtype=torch.uint8) + + bank = [] + for _ in range(num_experts): + bank.append( + { + "w1": nibbles(intermediate, hidden // 2), + "w1_sf": scales(intermediate, hidden // 32), + "w3": nibbles(intermediate, hidden // 2), + "w3_sf": scales(intermediate, hidden // 32), + "w2": nibbles(hidden, intermediate // 2), + "w2_sf": scales(hidden, intermediate // 32), + } + ) + return bank + + +def _make_test_gate(num_experts=_TP_EXPERTS, seed=71): + """One deterministically-initialized gate SHARED by all modules under + comparison: the fused routing kernel applies the gate's + e_score_correction_bias per module, so a per-module `torch.empty` + (garbage) bias would silently route the shard and whole-expert modules + to different experts.""" + cfg = _K3Config( + hidden_size=_TP_HIDDEN, + num_experts=num_experts, + num_experts_per_token=min(_TP_TOPK, num_experts), + ) + gate = KimiK3MoEGate(cfg) + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + gate.weight.copy_(torch.randn(gate.weight.shape, generator=gen, dtype=torch.float32) * 0.05) + gate.e_score_correction_bias.copy_( + torch.randn(gate.e_score_correction_bias.shape, generator=gen, dtype=torch.float32) + * 0.1 + ) + return gate.cuda() + + +def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS): + """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping.""" + from transformers.configuration_utils import PretrainedConfig + + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + pretrained_config = PretrainedConfig() + pretrained_config.num_experts = num_experts + pretrained_config.hidden_size = _TP_HIDDEN + pretrained_config.intermediate_size = intermediate_size + pretrained_config.torch_dtype = torch.bfloat16 + model_config = ModelConfig( + pretrained_config=pretrained_config, + mapping=Mapping(), + moe_backend="TRTLLM", + ) + moe = create_moe( + routing_method=gate.routing_method, + num_experts=num_experts, + hidden_size=_TP_HIDDEN, + intermediate_size=intermediate_size, + dtype=torch.bfloat16, + reduce_results=True, + model_config=model_config, + override_quant_config=QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8), + layer_idx=0, + trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, + trtllm_gen_activation_alpha=4.0, + trtllm_gen_activation_beta=25.0, + communication_method=None, + ).cuda() + assert isinstance(moe, ConfigurableMoE) + return moe + + +def _load_bank(moe, bank, tp_size=1, tp_rank=0): + """Load packed experts through the production per-expert adapter. + + ``tp_size > 1`` simulates one MoE-TP rank: the proxy exposes the shard + coordinates so the stock `load_weight_shard` slicing runs exactly as it + would on a real multi-rank mapping, while the backing module holds the + shard-sized parameters. + """ + backend = moe.backend + proxy = SimpleNamespace( + expert_size_per_partition=backend.expert_size_per_partition, + initial_local_expert_ids=backend.initial_local_expert_ids, + scaling_vector_size=backend.scaling_vector_size, + tp_size=tp_size, + tp_rank=tp_rank, + w3_w1_weight=backend.w3_w1_weight, + w2_weight=backend.w2_weight, + w3_w1_weight_scale=backend.w3_w1_weight_scale, + w2_weight_scale=backend.w2_weight_scale, + ) + for expert_id, tensors in enumerate(bank): + backend.quant_method.load_packed_mxfp4_expert( + proxy, + global_expert_id=expert_id, + local_slot_id=expert_id, + w1_weight=tensors["w1"], + w1_weight_scale=tensors["w1_sf"], + w2_weight=tensors["w2"], + w2_weight_scale=tensors["w2_sf"], + w3_weight=tensors["w3"], + w3_weight_scale=tensors["w3_sf"], + ) + backend._weights_transformed = False + moe.post_load_weights() + return moe + + +@situ_supported +@pytest.mark.parametrize("tp_size", [2, 8], ids=lambda n: f"tp{n}") +def test_tp_shard_loader_matches_manual_slice(tp_size): + """The stock shard loaders must equal a manual contiguous slice. + + Guards the group-32 packed-byte / scale slicing assumptions: w1/w3 + column-shard along intermediate rows, w2 row-shard along the packed + intermediate bytes and per-32-group scales. + """ + ipp = _TP_INTERMEDIATE // tp_size + num_experts = 2 + bank = _make_packed_expert_bank(num_experts, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate(num_experts=num_experts) + + for tp_rank in (0, tp_size - 1): + via_shard = _make_routed_moe(ipp, gate, num_experts=num_experts) + _load_bank(via_shard, bank, tp_size=tp_size, tp_rank=tp_rank) + + rows = slice(tp_rank * ipp, (tp_rank + 1) * ipp) + cols_packed = slice(tp_rank * (ipp // 2), (tp_rank + 1) * (ipp // 2)) + cols_sf = slice(tp_rank * (ipp // 32), (tp_rank + 1) * (ipp // 32)) + manual_bank = [ + { + "w1": e["w1"][rows].contiguous(), + "w1_sf": e["w1_sf"][rows].contiguous(), + "w3": e["w3"][rows].contiguous(), + "w3_sf": e["w3_sf"][rows].contiguous(), + "w2": e["w2"][:, cols_packed].contiguous(), + "w2_sf": e["w2_sf"][:, cols_sf].contiguous(), + } + for e in bank + ] + via_manual = _make_routed_moe(ipp, gate, num_experts=num_experts) + _load_bank(via_manual, manual_bank) + + for name in ("w3_w1_weight", "w2_weight", "w3_w1_weight_scale", "w2_weight_scale"): + a = getattr(via_shard.backend, name).data + b = getattr(via_manual.backend, name).data + assert torch.equal(a, b), f"{name} mismatch for tp_size={tp_size} tp_rank={tp_rank}" + + +@situ_supported +@pytest.mark.parametrize("num_tokens", [1, 16], ids=lambda n: f"tokens{n}") +def test_tp8_sharded_forward_matches_whole_expert(num_tokens): + """Sum of 8 TP-shard partial outputs == whole-expert reference. + + Per-element MXFP4/MXFP8 numerics are identical between the two layouts + (group-32 boundaries align: 384 % 32 == 0), so the only expected error + is bf16 rounding of the per-shard FC2 partial sums. + """ + tp_size = 8 + ipp = _TP_INTERMEDIATE // tp_size # 384 — the production TP8 shard size + bank = _make_packed_expert_bank(_TP_EXPERTS, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate() + + whole = _make_routed_moe(_TP_INTERMEDIATE, gate) + _load_bank(whole, bank) + + torch.manual_seed(3) + x = torch.randn(num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5 + router_logits = gate.compute_logits(x) + + out_whole = whole.forward(x, router_logits, all_rank_num_tokens=None) + + partial_sum = torch.zeros(num_tokens, _TP_HIDDEN, dtype=torch.float32, device="cuda") + for tp_rank in range(tp_size): + shard = _make_routed_moe(ipp, gate) + _load_bank(shard, bank, tp_size=tp_size, tp_rank=tp_rank) + out_shard = shard.forward(x, router_logits, all_rank_num_tokens=None) + partial_sum += out_shard.float() + del shard + torch.cuda.empty_cache() + + check_accuracy(partial_sum.to(torch.bfloat16), out_whole, atol=0.08, rtol=0.08, percent=0.98) diff --git a/tests/unittest/models/test_quant_config_utils.py b/tests/unittest/models/test_quant_config_utils.py index b2382cd70ab8..d10c6da77709 100644 --- a/tests/unittest/models/test_quant_config_utils.py +++ b/tests/unittest/models/test_quant_config_utils.py @@ -284,3 +284,30 @@ def test_update_quant_config_from_compressed_tensors_rejects_kv_cache_conflict() } ), ) + + +def test_update_quant_config_from_compressed_tensors_mxfp4_with_fp8_kv_cache(): + quant_config = QuantConfig() + update_quant_config_from_compressed_tensors( + quant_config, + _compressed_tensors_config( + weights={ + "num_bits": 4, + "type": "float", + "strategy": "group", + "group_size": 32, + }, + format="mxfp4-pack-quantized", + kv_cache_scheme={ + "num_bits": 8, + "type": "float", + }, + ignore=["lm_head"], + ), + ) + + assert quant_config.quant_algo == QuantAlgo.W4A16_MXFP4 + assert quant_config.group_size == 32 + # The MXFP4 branch returns early; kv_cache_scheme must still be honored. + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert set(quant_config.exclude_modules) == {"lm_head"}