From 03840efa9690acabd709d4646f776548b360b36b Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:39:33 -0700 Subject: [PATCH 1/9] [None][feat] Support multi-modal part of K3 Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- tensorrt_llm/_torch/configs/__init__.py | 15 +- tensorrt_llm/_torch/configs/kimi_k3.py | 140 +++++ tensorrt_llm/_torch/models/__init__.py | 1 + tensorrt_llm/_torch/models/_arch_index.py | 4 +- .../_torch/models/modeling_kimi_k25.py | 11 + .../_torch/models/modeling_kimi_k3_vl.py | 509 ++++++++++++++++++ .../_torch/models/modeling_kimi_linear.py | 8 +- .../_torch/pyexecutor/config_utils.py | 36 +- tensorrt_llm/evaluate/lm_eval.py | 19 +- tensorrt_llm/evaluate/post_processing.py | 59 ++ tests/unittest/others/test_lm_eval.py | 111 ++++ 11 files changed, 895 insertions(+), 18 deletions(-) create mode 100644 tensorrt_llm/_torch/configs/kimi_k3.py create mode 100644 tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index 4152f14ccb07..9110469e5fbb 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_k3 import KimiK3Config, KimiK3VisionConfig 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 @@ -55,11 +56,13 @@ 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 K3 composite multimodal config ("kimi_k3") and its text config + # ("kimi_linear"). pyexecutor.config_utils.load_pretrained_config keeps + # the composite KimiK3Config when the checkpoint ships text+vision + # sub-configs and multimodal is not disabled, and otherwise flattens to + # the text config. Registering both here lets AutoConfig / AutoTokenizer + # resolve them without trust_remote_code. + "kimi_k3": KimiK3Config, "kimi_linear": KimiLinearConfig, "laguna": LagunaConfig, # minicpmv4_6 is only registered in transformers>=5.7.0; register our @@ -93,6 +96,8 @@ def _register_custom_configs_with_transformers() -> None: "Gemma4UnifiedConfig", "Gemma4UnifiedTextConfig", "Gemma4UnifiedVisionConfig", + "KimiK3Config", + "KimiK3VisionConfig", "KimiLinearConfig", "LagunaConfig", "MiniCPMV4_6Config", diff --git a/tensorrt_llm/_torch/configs/kimi_k3.py b/tensorrt_llm/_torch/configs/kimi_k3.py new file mode 100644 index 000000000000..d865e7d3ab54 --- /dev/null +++ b/tensorrt_llm/_torch/configs/kimi_k3.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""In-tree composite config for Kimi K3 ("kimi_k3") multimodal checkpoints. + +Mirrors the checkpoint-shipped ``configuration_kimi_k3.KimiK3Config`` / +``KimiK3VisionConfig`` so TRT-LLM can parse the released Kimi K3 VLM checkpoint +without ``trust_remote_code`` for the config. The composite ``kimi_k3`` config +nests the in-tree text config (``KimiLinearConfig``) as ``text_config`` and a +flat ``KimiK3VisionConfig`` as ``vision_config``. + +``pyexecutor.config_utils.load_pretrained_config`` keeps this composite config +(architectures ``KimiK3ForConditionalGeneration``) when the checkpoint ships both +``text_config`` and ``vision_config`` and multimodal is not disabled; otherwise +it flattens to the text config (``KimiLinearForCausalLM``) as before. +""" + +from typing import Optional, Union + +from transformers.configuration_utils import PretrainedConfig + +from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig + + +class KimiK3VisionConfig(PretrainedConfig): + """Vision-tower + projector sub-config for Kimi K3. + + Flat fields mirroring the checkpoint-shipped ``KimiK3VisionConfig``. The + vision tower is a MoonViT-3D encoder (``vt_*`` fields) and the projector is a + ``patchmergerv2`` MLP (``mm_*`` / ``projector_*`` fields). + """ + + model_type = "kimi_k3_vision" + + def __init__( + self, + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + vt_num_attention_heads: int = 12, + vt_num_hidden_layers: int = 27, + vt_hidden_size: int = 1024, + vt_intermediate_size: int = 4096, + merge_kernel_size: tuple = (2, 2), + merge_type: str = "sd2_tpool", + _attn_implementation: str = "flash_attention_2", + # MM Projector parameters + mm_projector_type: str = "patchmergerv2", + mm_hidden_size: Optional[int] = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + # vision tower parameters + qkv_hidden_size: int = 1536, + norm_type: str = "rmsnorm", + attn_bias: bool = False, + patch_embed_proj_bias: bool = False, + mlp_type: str = "mlp2", + linear_bias: bool = False, + activation_func: str = "gelu_pytorch_tanh", + pos_emb_interpolation_mode: str = "bilinear", + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + text_hidden_size: int = 7168, + **kwargs, + ): + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.vt_num_attention_heads = vt_num_attention_heads + self.vt_num_hidden_layers = vt_num_hidden_layers + self.vt_hidden_size = vt_hidden_size + self.vt_intermediate_size = vt_intermediate_size + self.merge_kernel_size = tuple(merge_kernel_size) + self.merge_type = merge_type + self._attn_implementation = _attn_implementation + + # MM Projector config + self.mm_projector_type = mm_projector_type + self.mm_hidden_size = (mm_hidden_size + if mm_hidden_size is not None else vt_hidden_size) + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + # vision tower parameters + self.qkv_hidden_size = qkv_hidden_size + self.norm_type = norm_type + self.attn_bias = attn_bias + self.patch_embed_proj_bias = patch_embed_proj_bias + self.mlp_type = mlp_type + self.linear_bias = linear_bias + self.activation_func = activation_func + self.pos_emb_interpolation_mode = pos_emb_interpolation_mode + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + +class KimiK3Config(PretrainedConfig): + """Top-level composite config for the Kimi K3 multimodal model. + + ``text_config`` is the in-tree :class:`KimiLinearConfig` (the already + brought-up text core); ``vision_config`` is :class:`KimiK3VisionConfig`. + Sub-configs arrive as nested dicts from ``AutoConfig.from_pretrained`` and are + rebuilt with the in-tree classes here so no ``trust_remote_code`` is needed. + """ + + model_type = "kimi_k3" + + def __init__( + self, + text_config: Optional[Union[dict, KimiLinearConfig]] = None, + vision_config: Optional[Union[dict, KimiK3VisionConfig]] = None, + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + **kwargs, + ): + if isinstance(text_config, dict): + text_config = KimiLinearConfig(**text_config) + if isinstance(vision_config, dict): + vision_config = KimiK3VisionConfig(**vision_config) + self.text_config = text_config + self.vision_config = vision_config + + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + # The routed-expert MXFP4 quantization lives on the text config; surface + # it at the top level so TRT-LLM's quant-config extraction finds it. + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index c8bbc86672c4..4abb9565b7f8 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -49,6 +49,7 @@ "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", "KimiK25ForConditionalGeneration", + "KimiK3ForConditionalGeneration", "KimiLinearForCausalLM", "LagunaForCausalLM", "LlamaForCausalLM", diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index ed392ab3e51d..8915af4c6475 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -63,7 +63,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", "KimiK25ForConditionalGeneration": "modeling_kimi_k25", - "KimiK3ForConditionalGeneration": "modeling_kimi_linear", + "KimiK3ForConditionalGeneration": "modeling_kimi_k3_vl", "KimiLinearForCausalLM": "modeling_kimi_linear", "LagunaForCausalLM": "modeling_laguna", "Llama4ForConditionalGeneration": "modeling_llama", @@ -141,6 +141,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", "KimiK25ForConditionalGeneration": "modeling_kimi_k25", + "KimiK3ForConditionalGeneration": "modeling_kimi_k3_vl", "KimiLinearForCausalLM": "modeling_kimi_linear", "LagunaForCausalLM": "modeling_laguna", "LlamaForCausalLM": "modeling_llama", @@ -203,6 +204,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "gemma4_unified": "modeling_gemma4_unified", "hyperclovax_vlm": "modeling_hyperclovax", "kimi_k25": "modeling_kimi_k25", + "kimi_k3": "modeling_kimi_k3_vl", "llama4": "modeling_llama", "llava_llama": "modeling_vila", "llava_next": "modeling_llava_next", diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 723f242cb613..b0f31bf55893 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -883,6 +883,17 @@ def load_weights(self, weights: Dict[str, torch.Tensor]) -> None: converted: Dict[str, torch.Tensor] = {} for name, weight in mapped.items(): + # The runtime HF loader streams weights as lazy safetensors + # ``PySafeSlice`` objects (``safe_open(...).get_slice(name)``), which + # support slicing but none of the torch.Tensor attributes/methods the + # downstream consume path touches (``.chunk`` here, and ``.device`` / + # ``.dtype`` / ``.shape`` inside the child ``Linear.load_weights`` -> + # ``load_weight_shard``). Materialize EVERY value once, up front, so + # no raw PySafeSlice can leak into any child module. The vision tower + # + projector is small (168 tensors, non-quantized, tp=1 replicated), + # so full materialization is cheap; ``[:]`` on an already-real tensor + # is a harmless view, so the module-parity path is unaffected. + weight = weight[:] if ".wqkv." in name: prefix, suffix = name.split(".wqkv.", 1) q_weight, k_weight, v_weight = weight.chunk(3, dim=0) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py new file mode 100644 index 000000000000..bca6a545e760 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -0,0 +1,509 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Kimi K3 multimodal model for TensorRT-LLM PyTorch backend. + +Wires the Kimi K3 vision-language model on top of the already-brought-up +text-only core (:class:`KimiLinearForCausalLM`): + +- MoonViT3d vision encoder (native, no ``trust_remote_code`` for the model) +- PatchMergerMLPV2 vision->text projector +- ``KimiK3ForConditionalGeneration`` that fuses vision embeddings into the + KimiLinear text backbone. + +The structure closely mirrors the in-tree K2.5 model +(``modeling_kimi_k25.py``); this file only carries the K3-specific deltas and +reuses everything numerically identical from the K2.5 implementation: + + delta | K2.5 | K3 + ---------------------+----------------------------+--------------------------- + vision norms | LayerNorm | RMSNorm (torch.nn.RMSNorm) + attention head_dim | vt_hidden_size // heads | qkv_hidden_size // heads + patch-embed conv bias| True | False + vision qkv/o/MLP bias| True | False + projector | PatchMergerMLP (pre_norm) | PatchMergerMLPV2 (post_norm) + text backbone | DeepseekV3ForCausalLM | KimiLinearForCausalLM +""" + +import copy +from typing import Optional, Tuple + +import torch +import torch.nn as nn +from transformers import PretrainedConfig, PreTrainedModel + +from ...inputs import ( + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, + register_input_processor, +) +from ..attention_backend import AttentionMetadata +from ..attention_backend.utils import get_attention_backend +from ..model_config import ModelConfig +from ..modules.linear import Linear, TensorParallelMode +from ..modules.mlp import MLP +from .checkpoints.base_weight_loader import ConsumableWeightsDict +from .modeling_kimi_k25 import ( + _MEDIA_PLACEHOLDER_TOKEN_ID, + DISAGG, + KimiK25ForConditionalGeneration, + KimiK25InputProcessor, + KimiK25VisionAttention, + KimiK25VisionModel, + Learnable2DPosEmb, + MoonViT3dEncoder, + Rope2D, + _gelu_tanh, + _get_vision_tp_mapping, + _has_meta_tensors, +) +from .modeling_kimi_linear import KimiLinearForCausalLM +from .modeling_utils import ( + MetaInitException, + QuantConfig, + filter_weights, + register_auto_model, + register_vision_encoder, +) + +from ...logger import logger # noqa: E402 + + +# --------------------------------------------------------------------------- +# Native MoonViT3d Vision Encoder Components (K3 deltas) +# --------------------------------------------------------------------------- + + +class K3PatchEmbed3d(nn.Module): + """Conv2d patch embedding (bias-free) + learnable 2D position embedding. + + Mirrors ``MoonVision3dPatchEmbed`` with ``patch_embed_proj_bias=False``. + ``Learnable2DPosEmb`` is numerically identical to the reference + ``Learnable2DInterpPosEmbDivided_fixed`` for image inputs (bicubic interp). + """ + + def __init__( + self, + hidden_dim: int, + patch_size: int, + pos_emb_height: int, + pos_emb_width: int, + pos_emb_time: int, + ) -> None: + super().__init__() + self.proj = nn.Conv2d(3, hidden_dim, kernel_size=patch_size, stride=patch_size, bias=False) + self.pos_emb = Learnable2DPosEmb(pos_emb_height, pos_emb_width, pos_emb_time, hidden_dim) + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + x = self.proj(x).view(x.size(0), -1) + return self.pos_emb(x, grid_thws) + + +class K3VisionMLP(MLP): + """Bias-free MoonViT3d MLP2 (fc0 -> gelu_tanh -> fc1).""" + + def __init__( + self, model_config: ModelConfig, layer_idx: int, hidden_dim: int, mlp_dim: int + ) -> None: + super().__init__( + hidden_size=hidden_dim, + intermediate_size=mlp_dim, + bias=False, + activation=_gelu_tanh, + dtype=model_config.torch_dtype, + config=model_config, + layer_idx=layer_idx, + overridden_tp_size=1 if model_config.mapping.enable_attention_dp else None, + ) + + +class K3EncoderLayer(nn.Module): + """Single MoonViT3d encoder layer with RMSNorm + bias-free attention/MLP.""" + + def __init__( + self, + model_config: ModelConfig, + layer_idx: int, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + ) -> None: + super().__init__() + # Reference uses torch.nn.RMSNorm(hidden_dim) with default eps for the + # per-layer norms; match it exactly (created in fp32, cast with the rest + # of the vision tower to the model dtype in load_weights()). + self.norm0 = nn.RMSNorm(hidden_dim) + self.norm1 = nn.RMSNorm(hidden_dim) + # head_dim is taken from model_config.pretrained_config.head_dim, which + # KimiK3VisionModel sets to qkv_hidden_size // num_heads (128). + self.attn = KimiK25VisionAttention( + model_config, + hidden_dim=hidden_dim, + num_heads=num_heads, + layer_idx=layer_idx, + attn_bias=False, + ) + self.mlp = K3VisionMLP(model_config, layer_idx, hidden_dim, mlp_dim) + + def forward( + self, + x: torch.Tensor, + attn_metadata: AttentionMetadata, + freqs_cis: torch.Tensor, + ) -> torch.Tensor: + residual = x + x = self.norm0(x) + x = self.attn(x, attn_metadata, freqs_cis) + x = residual + x + residual = x + x = self.norm1(x) + x = residual + self.mlp(x) + return x + + +class K3MoonViT3dEncoder(MoonViT3dEncoder): + """MoonViT3d encoder stack with K3 deltas (RMSNorm + qkv-sized head_dim). + + Reuses :meth:`MoonViT3dEncoder.prepare_attn_metadata` and + :meth:`MoonViT3dEncoder.forward`; only the sub-module construction differs, + so ``__init__`` builds the K3 blocks directly and skips the K2.5 parent + ``__init__`` (which would build LayerNorm blocks with the wrong head_dim). + """ + + def __init__( + self, + model_config: ModelConfig, + hidden_dim: int, + num_layers: int, + num_heads: int, + mlp_dim: int, + head_dim: int, + ) -> None: + nn.Module.__init__(self) + self.rope_2d = Rope2D(head_dim) + self.blocks = nn.ModuleList( + [ + K3EncoderLayer( + model_config, + layer_idx=layer_idx, + num_heads=num_heads, + hidden_dim=hidden_dim, + mlp_dim=mlp_dim, + ) + for layer_idx in range(num_layers) + ] + ) + self.final_layernorm = nn.RMSNorm(hidden_dim) + self.metadata_cls = get_attention_backend(model_config.attn_backend).Metadata + self.attn_metadata: Optional[AttentionMetadata] = None + + +class PatchMergerMLPV2(nn.Module): + """K3 vision->text projector: view-merge -> Linear -> GELU -> Linear -> RMSNorm. + + Matches the reference ``PatchMergerMLPV2``: no ``pre_norm``, bias-free + projections, exact (erf) ``nn.GELU``, and a trailing ``RMSNorm`` over the + text hidden size. Reuses TRT-LLM ``Linear`` (tp_size==1 under attention_dp) + so the projector stays TP-composable like the K2.5 projector. + """ + + def __init__( + self, + model_config: ModelConfig, + mm_hidden_size: int, + text_hidden_size: int, + merge_kernel_size: Tuple[int, int] = (2, 2), + ln_eps: float = 1e-5, + ) -> None: + super().__init__() + kh, kw = merge_kernel_size + self.merged_dim = mm_hidden_size * kh * kw + mapping = _get_vision_tp_mapping(model_config) + self.proj = nn.Sequential( + Linear( + self.merged_dim, + self.merged_dim, + bias=False, + dtype=model_config.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + allreduce_strategy=model_config.allreduce_strategy, + ), + nn.GELU(), + Linear( + self.merged_dim, + text_hidden_size, + bias=False, + dtype=model_config.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + allreduce_strategy=model_config.allreduce_strategy, + ), + ) + self.post_norm = nn.RMSNorm(text_hidden_size, eps=ln_eps) + + def forward(self, x): + if isinstance(x, (list, tuple)): + lengths = [item.shape[0] for item in x] + merged = torch.cat([item.reshape(item.shape[0], -1) for item in x], dim=0) + out = self.post_norm(self.proj(merged)) + return list(torch.split(out, lengths, dim=0)) + batch = x.shape[0] + return self.post_norm(self.proj(x.view(batch, -1, self.merged_dim))) + + +# --------------------------------------------------------------------------- +# MoonViT3d Vision Encoder (top-level wrapper) +# --------------------------------------------------------------------------- + + +class KimiK3VisionModel(KimiK25VisionModel): + """Native MoonViT3d encoder + PatchMergerMLPV2 projector for Kimi K3. + + Reuses :meth:`KimiK25VisionModel.load_weights`, ``_extract_features`` and + ``forward`` (the HF weight names and the merge/projection pipeline are + identical); only the sub-module construction differs, so ``__init__`` builds + the K3 tower and skips the K2.5 parent ``__init__``. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + nn.Module.__init__(self) + self.model_config = copy.copy(model_config) + self.model_config.extra_attrs = copy.copy(model_config.extra_attrs) + self.model_config._frozen = False + # Vision tower is not quantized (checkpoint quant ignore list covers + # vision_tower.* / mm_projector.*): keep only the kv-cache quant algo. + self.model_config.quant_config = QuantConfig( + kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo + ) + self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) + pretrained_config = self.model_config.pretrained_config + model_dtype = ( + getattr(pretrained_config, "torch_dtype", None) + or getattr(pretrained_config, "dtype", None) + or torch.bfloat16 + ) + if isinstance(model_dtype, str): + model_dtype = getattr(torch, model_dtype, torch.bfloat16) + pretrained_config.torch_dtype = model_dtype + + vision_cfg = getattr(pretrained_config, "vision_config", {}) + if vision_cfg is None: + vision_cfg = {} + if not isinstance(vision_cfg, dict): + vision_cfg = ( + vision_cfg.to_dict() if hasattr(vision_cfg, "to_dict") else vars(vision_cfg) + ) + + hidden_dim = vision_cfg.get("vt_hidden_size", vision_cfg.get("hidden_size", 1024)) + num_layers = vision_cfg.get("vt_num_hidden_layers", vision_cfg.get("num_hidden_layers", 27)) + num_heads = vision_cfg.get( + "vt_num_attention_heads", vision_cfg.get("num_attention_heads", 12) + ) + # K3 delta: the attention head_dim is qkv_hidden_size // num_heads (128), + # NOT vt_hidden_size // num_heads. wqkv projects hidden_dim -> 3*qkv and + # wo projects qkv -> hidden_dim, so q/k/v live in the qkv space. + qkv_hidden_size = vision_cfg.get("qkv_hidden_size", hidden_dim) + head_dim = qkv_hidden_size // num_heads + self.model_config.pretrained_config.head_dim = head_dim + self.model_config._frozen = True + + mlp_dim = vision_cfg.get("vt_intermediate_size", vision_cfg.get("intermediate_size", 4096)) + mm_hidden_size = vision_cfg.get("mm_hidden_size", hidden_dim) + text_hidden_size = vision_cfg.get("text_hidden_size", 7168) + patch_size = vision_cfg.get("patch_size", 14) + ln_eps = vision_cfg.get("projector_ln_eps", 1e-5) + pos_h = vision_cfg.get("init_pos_emb_height", 64) + pos_w = vision_cfg.get("init_pos_emb_width", 64) + pos_t = vision_cfg.get("init_pos_emb_time", 4) + + merge_ks = vision_cfg.get("merge_kernel_size", [2, 2]) + if isinstance(merge_ks, int): + self.merge_kernel_size = (merge_ks, merge_ks) + elif isinstance(merge_ks, (list, tuple)): + self.merge_kernel_size = tuple(merge_ks) + else: + self.merge_kernel_size = (2, 2) + self.merge_type = vision_cfg.get("merge_type", "sd2_tpool") + + text_config = getattr(pretrained_config, "text_config", pretrained_config) + self.model_dtype = model_dtype + self.text_hidden_size = ( + text_config.get("hidden_size", text_hidden_size) + if isinstance(text_config, dict) + else getattr(text_config, "hidden_size", text_hidden_size) + ) + self.config = PretrainedConfig( + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + tie_word_embeddings=False, + ) + + self.patch_embed = K3PatchEmbed3d(hidden_dim, patch_size, pos_h, pos_w, pos_t) + self.encoder = K3MoonViT3dEncoder( + self.model_config, + hidden_dim, + num_layers, + num_heads, + mlp_dim, + head_dim, + ) + self.mm_projector = PatchMergerMLPV2( + self.model_config, + mm_hidden_size, + self.text_hidden_size, + self.merge_kernel_size, + ln_eps, + ) + + +# --------------------------------------------------------------------------- +# Input Processor +# --------------------------------------------------------------------------- + + +class KimiK3InputProcessor(KimiK25InputProcessor): + """Image-only input processor for Kimi K3. + + Reuses the K2.5 processor: the K3 ``AutoProcessor`` (``KimiK3Processor``, + loaded via ``trust_remote_code`` from the checkpoint) exposes the same + ``image_processor.media_tokens_calculator`` and ``(medias=, text=)`` call + contract, and returns ``grid_thws`` / ``pixel_values``. The framework + injects the K3 ``<|kimi_image_placeholder|>`` marker (see the registration + below), which ``KimiK3Processor.update_raw_text`` expands into the + ``<|media_pad|>`` (id 163605) run that ``call_with_text_prompt`` then + duplicates to ``(h // merge_kh) * (w // merge_kw)`` tokens per image. + """ + + +# --------------------------------------------------------------------------- +# Full VLM Model +# --------------------------------------------------------------------------- + + +@register_vision_encoder(KimiK3VisionModel) +@register_auto_model("KimiK3ForConditionalGeneration") +@register_input_processor( + KimiK3InputProcessor, + model_type="kimi_k3", + placeholder_metadata=MultimodalPlaceholderMetadata( + placeholder_map={ + "image": "<|kimi_image_placeholder|>", + }, + placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, + ), +) +class KimiK3ForConditionalGeneration(KimiK25ForConditionalGeneration): + """Kimi K3 vision-language model: MoonViT3d + KimiLinear text backbone. + + Reuses the K2.5 wrapper's spec-dec / weight-loading property forwarding and + :meth:`forward`; only ``__init__`` (vision encoder + text backbone classes) + and the deferred vision-encoder recreation in :meth:`load_weights` differ. + """ + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + *args, + **kwargs, + ) -> None: + config = model_config.pretrained_config + self._supports_sdpa = True + # Skip the K2.5 parent __init__ (it wires DeepSeek-V3 + K2.5 vision); + # initialize the HF PreTrainedModel machinery directly, then build the + # K3 components below. + PreTrainedModel.__init__(self, config) + + if hasattr(self, "llm"): + return + + self.model_config = model_config + self._vlm_pretrained_config = config + + # --- Vision encoder (deferred under MetaInitMode, recreated in + # load_weights) --- + self.mm_encoder = None + if not DISAGG: + try: + mm_encoder = KimiK3VisionModel(model_config) + if _has_meta_tensors(mm_encoder): + logger.info("Vision encoder deferred to load_weights() (MetaInitMode active)") + else: + self.mm_encoder = mm_encoder + except MetaInitException: + logger.info("Vision encoder deferred to load_weights() (MetaInitMode active)") + + text_model_config = copy.copy(model_config) + assert hasattr(config, "text_config"), "Kimi K3 config must have text_config" + text_model_config._frozen = False + text_model_config.pretrained_config = config.text_config + + # Remap quant exclude_modules: language_model.X -> model.X + if text_model_config.quant_config.exclude_modules: + text_model_config.quant_config = copy.copy(text_model_config.quant_config) + p = self._LANG_PREFIX + mapped = [] + for m in text_model_config.quant_config.exclude_modules: + if m.startswith(p): + rest = m[len(p) :] + if rest.startswith("layers."): + rest = "model." + rest + mapped.append(rest) + else: + mapped.append(m) + text_model_config.quant_config.exclude_modules = mapped + + if not text_model_config.skip_create_weights_in_init: + text_model_config.skip_create_weights_in_init = True + text_model_config._frozen = True + + self.llm = KimiLinearForCausalLM(text_model_config) + + self._media_placeholder_token_id = getattr( + config, "media_placeholder_token_id", _MEDIA_PLACEHOLDER_TOKEN_ID + ) + self._mm_token_ids = torch.tensor([self._media_placeholder_token_id], dtype=torch.int32) + + # Point model_config at the text_config so the executor reads generation + # params (eos_token_id, ...) from the text backbone, matching K2.5. + self.config = self.llm.config + model_config._frozen = False + model_config.pretrained_config = self.llm.config + model_config._frozen = True + + def load_weights(self, weights) -> None: + """Load vision + projector + KimiLinear text weights from checkpoint.""" + if self.mm_encoder is not None and _has_meta_tensors(self.mm_encoder): + logger.info("Recreating deferred vision encoder after MetaInitMode") + self.mm_encoder = None + + if self.mm_encoder is None and not DISAGG: + vision_model_config = copy.copy(self.model_config) + vision_model_config._frozen = False + vision_model_config.pretrained_config = self._vlm_pretrained_config + vision_model_config._frozen = True + self.mm_encoder = KimiK3VisionModel(vision_model_config) + if self.mm_encoder is not None: + self.mm_encoder.load_weights(weights) + + if any(k.startswith(self._LANG_PREFIX) for k in weights): + lm_weights = filter_weights("language_model", weights) + lm_weights = ConsumableWeightsDict(lm_weights) + else: + lm_weights = weights + self.llm.load_weights(lm_weights) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 7bf5867c3596..7f8edb57ce91 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2089,10 +2089,14 @@ def _materialize(value) -> torch.Tensor: 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).""" + """Kimi K3 text core (KDA + MLA + MoE). + + Serves text-only ``kimi_linear`` checkpoints directly, and is reused as the + text backbone by the multimodal ``KimiK3ForConditionalGeneration`` wrapper + (``modeling_kimi_k3_vl``). The composite ``KimiK3ForConditionalGeneration`` + architecture is registered by that wrapper, not here.""" def __init__(self, model_config: ModelConfig): cfg = _get_text_config(model_config.pretrained_config) diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 383bc3834e5e..31b07fc4c4a5 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -513,6 +513,26 @@ def _build_minicpmv4_6_config( return composite_config +def is_kimi_k3_multimodal_config(config_dict: dict) -> bool: + """Detect Kimi K3's composite multimodal VLM checkpoint config. + + The released Kimi K3 VLM checkpoint advertises ``model_type: kimi_k3`` with + nested ``text_config`` (the ``kimi_linear`` text core) and ``vision_config`` + (the MoonViT-3D tower + projector). When both sub-configs are present and + multimodal is not explicitly disabled, TRT-LLM keeps the composite + ``KimiK3Config`` and routes the checkpoint to + ``KimiK3ForConditionalGeneration``. Text-only ``kimi_linear`` checkpoints (or + a ``kimi_k3`` config with ``language_model_only: true`` / no vision_config) + fall through to the text-only flatten path. + """ + text_config = config_dict.get("text_config") + vision_config = config_dict.get("vision_config") + return (config_dict.get("model_type") == "kimi_k3" + and config_dict.get("language_model_only") is not True + and isinstance(text_config, dict) and bool(text_config) + and isinstance(vision_config, dict) and bool(vision_config)) + + # TODO: remove this once the transformers can support all of those models in _CONFIG_REGISTRY class LazyConfigDict(dict): @@ -597,11 +617,19 @@ 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 is_kimi_k3_multimodal_config(config_dict): + # Kimi K3 multimodal VLM: keep the composite KimiK3Config so the vision + # tower, projector, and multimodal token id remain available, and route + # to KimiK3ForConditionalGeneration. Must precede the text-only flatten + # branch below so vision_config isn't dropped. Built from the in-tree + # config classes (no trust_remote_code needed for the config). + from tensorrt_llm._torch.configs import KimiK3Config + model_config = KimiK3Config.from_dict(config_dict, **kwargs) + model_config.architectures = ["KimiK3ForConditionalGeneration"] 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). + # Kimi K3 text-only (or multimodal explicitly disabled): flatten to the + # in-tree KimiLinearConfig and run the text model only (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) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index 2beaed8a558a..f927729378ab 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -1004,18 +1004,21 @@ def evaluate(self, def command_harness(cls, ctx, **kwargs): llm: PyTorchLLM = ctx.obj - # Resolve the post-processor: accept a callable (already-bound) or the - # string key "strip_thinking_mmmu" coming from CLI flags. + # Resolve the post-processor: accept a callable (already-bound) or a + # string key coming from CLI flags. post_process_fn = kwargs.pop("post_process_fn", None) if isinstance(post_process_fn, str): if post_process_fn == "strip_thinking_mmmu": from .post_processing import \ strip_thinking_and_extract_mmmu_answer post_process_fn = strip_thinking_and_extract_mmmu_answer + elif post_process_fn == "kimi_k3_mmmu": + from .post_processing import extract_kimi_k3_mmmu_answer + post_process_fn = extract_kimi_k3_mmmu_answer else: raise click.BadParameter( - f"Unknown --post_process_fn={post_process_fn!r}; expected 'strip_thinking_mmmu'." - ) + f"Unknown --post_process_fn={post_process_fn!r}; expected " + "'strip_thinking_mmmu' or 'kimi_k3_mmmu'.") evaluator = cls( dataset_path=kwargs.pop("dataset_path", None), @@ -1614,12 +1617,16 @@ def __init__(self, **kwargs): "produce chain-of-thought before the answer).") @click.option( "--post_process_fn", - type=click.Choice(["strip_thinking_mmmu"]), + type=click.Choice(["strip_thinking_mmmu", "kimi_k3_mmmu"]), default=None, help="Per-sample post-processor. 'strip_thinking_mmmu' strips " "... and then runs the MMMU answer extractor — needed " "for thinking models (Kimi K2.5, Step3p7) whose CoT output the " - "default lm-eval regex cannot parse.") + "default lm-eval regex cannot parse. 'kimi_k3_mmmu' reads the answer " + "from Kimi K3's <|open|>response<|sep|>...<|close|>response channel " + "(its reasoning ends with <|close|>think<|sep|>, not , so the " + "strip_thinking path cannot see the answer) and falls back to the " + "strip_thinking cascade when no channel is present.") @click.pass_context @staticmethod def command(ctx, **kwargs) -> None: diff --git a/tensorrt_llm/evaluate/post_processing.py b/tensorrt_llm/evaluate/post_processing.py index 770fca467b5f..2a1914f8c45f 100644 --- a/tensorrt_llm/evaluate/post_processing.py +++ b/tensorrt_llm/evaluate/post_processing.py @@ -161,3 +161,62 @@ def strip_thinking_and_extract_mmmu_answer(text: str) -> str: answer extraction. """ return extract_mmmu_answer(strip_thinking(text)) + + +# --- Kimi K3 channel-structured output ---------------------------------------- +# +# Kimi K3 does NOT use Kimi K2.5's ``...`` markup. It emits a +# channel-structured chat format whose reasoning block is terminated by +# ``<|close|>think<|sep|>`` and whose final answer lives in an explicit response +# channel:: +# +# <|close|>think<|sep|><|open|>response<|sep|>C<|close|>response<|sep|><|close|>message<|sep|> +# +# Because there is no ````, :func:`strip_thinking` returns the whole blob +# and :func:`extract_mmmu_answer`'s cascade mis-scores it — the correct letter in +# the response channel is discarded. On MMMU val this silently drops ~6-7 points +# even though the model answered correctly. The extractor below reads the answer +# straight from the response channel and reuses the shared MMMU cascade on just +# that span. + +# Capture the content of a ``<|open|>response<|sep|> ... `` channel up to its +# closing marker (or end-of-text when the output was truncated right after the +# answer opened). DOTALL so multi-line channel content is captured. +_KIMI_K3_RESPONSE_CHANNEL_RE = re.compile( + r"<\|open\|>\s*response\s*<\|sep\|>(.*?)" + r"(?:<\|close\|>\s*response|<\|close\|>\s*message|<\|open\|>|\Z)", + re.DOTALL, +) + +# Residual Kimi special tokens of the form ``<|...|>`` to scrub from a channel +# span before answer extraction. +_KIMI_SPECIAL_TOKEN_RE = re.compile(r"<\|[^|]*\|>") + + +def extract_kimi_k3_mmmu_answer(text: str) -> str: + r"""Extract an MMMU letter answer from Kimi K3 channel-structured output. + + Kimi K3 wraps its final answer in an explicit response channel + (``<|open|>response<|sep|> ... <|close|>response<|sep|>``) and terminates + reasoning with ``<|close|>think<|sep|>`` rather than ````. This + extractor pulls the *last* response-channel span and reuses the shared + :func:`extract_mmmu_answer` cascade on just that span. + + When no response channel is present — a short direct answer, or a thinking + trace truncated by ``finish_reason=length`` before the channel opened — it + falls back to :func:`strip_thinking_and_extract_mmmu_answer`, so Kimi K2.5's + behavior, short-answer handling, and truncated-output handling are all + unchanged. The K2.5 ```` path in :func:`strip_thinking` is not + touched. + """ + if not text: + return "" + matches = _KIMI_K3_RESPONSE_CHANNEL_RE.findall(text) + # The last non-empty response channel is the final answer. + for span in reversed(matches): + cleaned = _KIMI_SPECIAL_TOKEN_RE.sub(" ", span).strip() + if cleaned: + return extract_mmmu_answer(cleaned) + # No usable response channel: fall back to the K2.5 path, which also handles + # bare-letter direct answers and truncated thinking traces correctly. + return strip_thinking_and_extract_mmmu_answer(text) diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index bb789a7a8bab..3e67b283f4bb 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -1620,3 +1620,114 @@ def test_e2e_windowed_matches_final_score(monkeypatch): ) score = results["results"]["toy_arith"]["exact_match,strict-match"] assert score == pytest.approx(0.75) + + +# =========================================================================== +# post_processing — Kimi K3 channel-structured MMMU answer extraction +# =========================================================================== +# +# Kimi K3 emits a channel-structured chat format whose reasoning ends with +# ``<|close|>think<|sep|>`` (NOT ````) and whose final answer lives in a +# ``<|open|>response<|sep|> X <|close|>response<|sep|>`` channel. The K2.5 +# strip_thinking() path keys on ```` and therefore cannot see the +# answer, silently dropping ~6-7 MMMU points even when the model is correct +# (observed on the real checkpoint: full mmmu_val 67.67 -> 74.11 by parsing the +# channel alone). These tests pin the new extractor and guard that the K2.5 +# path is unchanged. + +from tensorrt_llm.evaluate.post_processing import ( # noqa: E402 + extract_kimi_k3_mmmu_answer, strip_thinking, + strip_thinking_and_extract_mmmu_answer) + + +def _k3_output(thinking: str, answer: str) -> str: + """Build a well-formed Kimi K3 channel-structured output string.""" + return (f"{thinking}<|close|>think<|sep|>" + f"<|open|>response<|sep|>{answer}<|close|>response<|sep|>" + f"<|close|>message<|sep|>") + + +def test_k3_channel_bare_letter(): + """Answer is a bare option letter inside the response channel.""" + out = _k3_output("Reasoning about the options... I'll go with C.", "C") + assert extract_kimi_k3_mmmu_answer(out) == "C" + + +def test_k3_channel_real_samples_recovered(): + """Real committed mmmu_val samples the old parser scored wrong (channel has + the correct letter): accounting doc_id 7->C, 12->D, 21->A.""" + assert extract_kimi_k3_mmmu_answer( + _k3_output("...Total debits adjusted = 126,925. Final: C.", "C")) == "C" + assert extract_kimi_k3_mmmu_answer( + _k3_output("...Ending = 0. Option D. Final just D.", "D")) == "D" + assert extract_kimi_k3_mmmu_answer( + _k3_output("...commonly known by that name, so True. (A).", "A")) == "A" + + +def test_k3_channel_parenthesized_answer(): + """Channel content ``(C) Photos 2 & 3`` reduces to the option letter.""" + out = _k3_output("Both use negative space.", "(C) Photos 2 & 3") + assert extract_kimi_k3_mmmu_answer(out) == "C" + + +def test_k3_channel_answer_is_phrase(): + """Channel content ``The answer is (D).`` reduces to the option letter.""" + out = _k3_output("Long derivation here.", "The answer is (D).") + assert extract_kimi_k3_mmmu_answer(out) == "D" + + +def test_k3_truncated_after_channel_open(): + """Output truncated right after the response channel opened still yields the + letter (regex boundary falls back to end-of-text).""" + out = ("Some reasoning.<|close|>think<|sep|>" + "<|open|>response<|sep|>B") # cut off before <|close|>response + assert extract_kimi_k3_mmmu_answer(out) == "B" + + +def test_k3_last_channel_wins(): + """When multiple response channels are present, the final one is the answer.""" + out = ( + "r1<|close|>think<|sep|><|open|>response<|sep|>A<|close|>response<|sep|>" + "<|open|>response<|sep|>E<|close|>response<|sep|><|close|>message<|sep|>") + assert extract_kimi_k3_mmmu_answer(out) == "E" + + +def test_k3_no_channel_bare_letter_falls_back(): + """Short direct answers with no channel go through the K2.5 fallback.""" + assert extract_kimi_k3_mmmu_answer("C") == "C" + assert extract_kimi_k3_mmmu_answer("Answer: (B)") == "B" + + +def test_k3_no_channel_truncated_thinking_does_not_crash(): + """A thinking trace truncated before the channel opened (finish=length): + no channel to parse -> fall back; must not raise and must not fabricate.""" + truncated = ("Let me reason step by step about this very long problem " + "that never reaches a final answer channel " * 20) + # Should not raise; returns whatever the fallback cascade yields (the model + # genuinely did not emit an answer, so the exact value is not asserted). + out = extract_kimi_k3_mmmu_answer(truncated) + assert isinstance(out, str) + + +def test_k3_empty_input(): + assert extract_kimi_k3_mmmu_answer("") == "" + + +def test_k3_scrubs_residual_special_tokens(): + """Residual ``<|...|>`` tokens inside a channel span are scrubbed, not + returned as part of the answer.""" + out = ("t<|close|>think<|sep|><|open|>response<|sep|>" + "<|reserved|>D<|close|>response<|sep|><|close|>message<|sep|>") + assert extract_kimi_k3_mmmu_answer(out) == "D" + + +def test_k2_5_strip_thinking_path_unchanged(): + """Guard: the new K3 extractor must not alter the K2.5 behavior.""" + k25 = "chain of thought hereAnswer: (C)" + # K2.5 path still extracts C directly. + assert strip_thinking_and_extract_mmmu_answer(k25) == "C" + # strip_thinking still returns content after the last . + assert strip_thinking(k25) == "Answer: (C)" + # And the K3 extractor, given a blob with no K3 response channel, + # defers to the K2.5 cascade and returns the same answer. + assert extract_kimi_k3_mmmu_answer(k25) == "C" From 13d8308eaa3df6fc3626567d6627b128dd9528a5 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:41:08 -0700 Subject: [PATCH 2/9] Fix the issue about vision part can not work with TP16 Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- .../_torch/models/modeling_kimi_k25.py | 35 ++++++++++++++++++- .../_torch/models/modeling_kimi_k3_vl.py | 7 ++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index b0f31bf55893..9190ba1d3692 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -355,8 +355,34 @@ def _gelu_tanh(x: torch.Tensor) -> torch.Tensor: return F.gelu(x, approximate="tanh") +def _vision_requires_replication(model_config: ModelConfig) -> bool: + """Whether the MoonViT vision tower must run replicated (tp=1) rather than + tensor-parallel sharded across the attention-TP ranks. + + Replication is required when either: + + * attention data-parallelism is enabled (attention weights are already + replicated per rank), or + * the vision attention head count is not divisible by the attention-TP + degree (e.g. Kimi K3's 12 heads under TP16). Such a tower cannot be + TP-sharded at all, so it must be replicated instead of tripping the + ``num_heads % tp_size`` assertion in :class:`Attention`. + """ + mapping = model_config.mapping + if mapping.enable_attention_dp: + return True + vision_cfg = getattr(model_config.pretrained_config, "vision_config", + None) or {} + if not isinstance(vision_cfg, dict): + vision_cfg = (vision_cfg.to_dict() + if hasattr(vision_cfg, "to_dict") else vars(vision_cfg)) + num_heads = vision_cfg.get("vt_num_attention_heads", + vision_cfg.get("num_attention_heads", 12)) + return (num_heads % mapping.tp_size) != 0 + + def _get_vision_tp_mapping(model_config: ModelConfig) -> Mapping: - if not model_config.mapping.enable_attention_dp: + if not _vision_requires_replication(model_config): return model_config.mapping return Mapping( @@ -770,6 +796,13 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo ) self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) + # The MoonViT tower cannot be tensor-parallel sharded when its attention + # head count is not divisible by the attention-TP degree (e.g. Kimi K3's + # 12 heads under TP16); run the whole tower replicated (tp=1) in that + # case so module construction and weight loading agree. Attention-DP + # already replicates the vision tower via its own path, so leave it be. + if not model_config.mapping.enable_attention_dp: + self.model_config.mapping = _get_vision_tp_mapping(model_config) pretrained_config = self.model_config.pretrained_config model_dtype = ( getattr(pretrained_config, "torch_dtype", None) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index bca6a545e760..dba194e54276 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -292,6 +292,13 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo ) self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) + # The MoonViT tower cannot be tensor-parallel sharded when its attention + # head count is not divisible by the attention-TP degree (e.g. Kimi K3's + # 12 heads under TP16); run the whole tower replicated (tp=1) in that + # case so module construction and weight loading agree. Attention-DP + # already replicates the vision tower via its own path, so leave it be. + if not model_config.mapping.enable_attention_dp: + self.model_config.mapping = _get_vision_tp_mapping(model_config) pretrained_config = self.model_config.pretrained_config model_dtype = ( getattr(pretrained_config, "torch_dtype", None) From b66fd8b1cada596e18aa70c5248724494b0ed0d6 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 3 Aug 2026 23:42:44 -0700 Subject: [PATCH 3/9] TEP16/DEP16 sbatch script update Signed-off-by: Michal Guzek --- .../deployment-guide-for-kimi-k3-on-trtllm.md | 4 +- examples/kimi_k3/README.md | 36 +- examples/kimi_k3/perf_sweep/acc_sweep.sbatch | 2 +- examples/kimi_k3/run_dspark_acceptance.sbatch | 174 ++++++++++ examples/kimi_k3/run_eval_kimi_k3.sbatch | 327 ++++++++++++++++++ examples/kimi_k3/run_gsm8k_kimi_k3.sbatch | 177 ---------- 6 files changed, 519 insertions(+), 201 deletions(-) create mode 100644 examples/kimi_k3/run_dspark_acceptance.sbatch create mode 100644 examples/kimi_k3/run_eval_kimi_k3.sbatch delete mode 100644 examples/kimi_k3/run_gsm8k_kimi_k3.sbatch diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index af3d60c9e0e2..52b890f1cf11 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -247,12 +247,12 @@ The response should contain a `choices[0].message.content` field completing the The repository ships a ready-made multi-node GSM8K evaluation job built on `trtllm-eval` with the tested DEP16 settings: ```bash -sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ +sbatch examples/kimi_k3/run_eval_kimi_k3.sbatch \ --model /path/to/kimi-k3-checkpoint \ --image /path/to/tensorrt-llm-container.sqsh ``` -The job writes progress and results to `kimi-k3-gsm8k-.log` in the submission directory. If no local dataset path is configured, `trtllm-eval` downloads GSM8K from the Hugging Face Hub. The completed log contains a results table with the normalized GSM8K exact-match scores. With the tested checkpoint and the settings in this example, users should expect approximately: +The job writes progress and results to `kimi-k3-eval-.log` in the submission directory. If no local dataset path is configured, `trtllm-eval` downloads GSM8K from the Hugging Face Hub. The completed log contains a results table with the normalized GSM8K exact-match scores. With the tested checkpoint and the settings in this example, users should expect approximately: | Filter | Exact match | | :-- | --: | diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index e388c040f168..dc4e0b249da7 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -21,12 +21,9 @@ other GPU architectures may be added in a future release. [build from source](../../docs/source/installation/build-from-source.md) for details. - `build_wheel.py` creates the virtual environment at the repository - root, named after the container's Python version: `.venv-3.12` for the - current containers (Python 3.12). If your container ships a different - Python, substitute the matching `.venv-.` path in the - commands on this page. Adjust `--cuda_architectures` to the target - GPUs (`103-real` for GB300). + `build_wheel.py` creates the `.venv-3.12` virtual environment at the + repository root (named after the container's Python version). Adjust + `--cuda_architectures` to the target GPUs (`103-real` for GB300). - A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g. [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded from the Hugging Face Hub (the example scripts take a local filesystem @@ -70,7 +67,7 @@ other GPU architectures may be added in a future release. provides FlashInfer's runtime dependencies; `--no-deps` prevents pip from replacing its pinned PyTorch, Triton, CUDA, and CuTeDSL packages. Install FlashInfer last: TensorRT-LLM - currently pins `flashinfer-python==0.6.16`, so a later + currently pins `flashinfer-python==0.6.14`, so a later dependency-resolving TensorRT-LLM install can replace this source revision. ## Run the model @@ -93,13 +90,13 @@ should report `True` for all four checks. For a full GSM8K evaluation, submit: ```bash -sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ +sbatch examples/kimi_k3/run_eval_kimi_k3.sbatch \ --model /path/to/kimi-k3-checkpoint \ --image /path/to/tensorrt-llm-container.sqsh ``` This job writes progress and results to -`kimi-k3-gsm8k-.log`. If no local dataset path is configured, +`kimi-k3-eval-.log`. If no local dataset path is configured, `trtllm-eval` downloads GSM8K from the Hugging Face Hub. The completed log contains a results table with the normalized GSM8K exact-match scores. With the tested checkpoint and the settings in this example, users should expect @@ -116,7 +113,7 @@ possible with different checkpoint or dependency revisions. To evaluate with suffix-automaton (SA) speculative decoding, add `--sa`: ```bash -sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ +sbatch examples/kimi_k3/run_eval_kimi_k3.sbatch \ --model /path/to/kimi-k3-checkpoint \ --image /path/to/tensorrt-llm-container.sqsh \ --sa @@ -184,7 +181,7 @@ Request a longer allocation if the time is not sufficient for your environment: sbatch --time=02:00:00 examples/kimi_k3/quick_start_kimi_k3.sbatch \ --model MODEL --image IMAGE -sbatch --time=04:00:00 examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ +sbatch --time=04:00:00 examples/kimi_k3/run_eval_kimi_k3.sbatch \ --model MODEL --image IMAGE ``` @@ -194,24 +191,21 @@ Chunked prefill is supported and enabled by default in this example (`enable_chunked_prefill: true` in the quick start and in `eval_extra_llm_options.yaml`). -KV-cache block reuse is supported. The LLM API enables it by default -(`KvCacheConfig.enable_block_reuse` defaults to `true`), but the example -configurations here explicitly disable it and treat reuse as an opt-in: -set `kv_cache_config.enable_block_reuse: true`, or use the example -flags — `--enable-block-reuse` for the quick start and `--reuse` for the -GSM8K job (which selects `eval_extra_llm_options_reuse.yaml`): +KV-cache block reuse is supported as an opt-in: set +`kv_cache_config.enable_block_reuse: true`, or use the example flags — +`--enable-block-reuse` for the quick start and `--reuse` for the GSM8K +job (which selects `eval_extra_llm_options_reuse.yaml`): ```bash sbatch examples/kimi_k3/quick_start_kimi_k3.sbatch \ --model MODEL --image IMAGE --enable-block-reuse -sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ +sbatch examples/kimi_k3/run_eval_kimi_k3.sbatch \ --model MODEL --image IMAGE --reuse ``` -Unless one of the flags above is passed, these examples run with block -reuse disabled; the tested evaluation and serving configurations use the -default cache manager. +Block reuse stays off by default because suffix-automaton speculative +decoding requires the default cache manager, which cannot reuse blocks. ## Current limitations diff --git a/examples/kimi_k3/perf_sweep/acc_sweep.sbatch b/examples/kimi_k3/perf_sweep/acc_sweep.sbatch index 01d655adfd83..2b32a05cade8 100644 --- a/examples/kimi_k3/perf_sweep/acc_sweep.sbatch +++ b/examples/kimi_k3/perf_sweep/acc_sweep.sbatch @@ -149,7 +149,7 @@ srun --mpi=pmix --kill-on-bad-exit=1 \ export HF_HOME=\${HF_HOME:-$CACHE_DIR/hf_home} mkdir -p \"\$FLASHINFER_WORKSPACE_BASE\" \"\$FLASHINFER_CUBIN_DIR\" \"\$HF_HOME\" # Node-local HF remote-code modules cache: racing mkdir over shared - # \$HOME escapes pathlib's exist_ok on NFS (see run_gsm8k_kimi_k3). + # \$HOME escapes pathlib's exist_ok on NFS (see run_eval_kimi_k3). export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0} # Running partial score every N responses (early failure signal). export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\" diff --git a/examples/kimi_k3/run_dspark_acceptance.sbatch b/examples/kimi_k3/run_dspark_acceptance.sbatch new file mode 100644 index 000000000000..6b3c56de2401 --- /dev/null +++ b/examples/kimi_k3/run_dspark_acceptance.sbatch @@ -0,0 +1,174 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 DSpark acceptance / speedup A/B on 16 NVIDIA Blackwell GPUs. +# +# One submission runs BOTH legs of the weights-day A/B back to back in the +# same allocation (identical nodes, so the TPOT comparison is clean): +# (a) spec-off TPOT reference (recorder off), +# (b) DSpark-on measurement (AL, per-position AR, histogram, and — with a +# confidence head in the drafter — calibration data). +# +# sbatch examples/kimi_k3/run_dspark_acceptance.sbatch \ +# --model /path/to/kimi-k3-checkpoint \ +# --drafter /path/to/dspark-drafter \ +# --image /path/to/tensorrt-llm-container.sqsh \ +# --outdir /path/to/results +# +# Optional: --num-prompts N (default 64), --max-tokens N (default 256), +# --confidence-threshold T --confidence-policy P (leave unset for +# calibration runs; see measure_dspark_acceptance.py), --skip-baseline +# (rerun only the DSpark leg, e.g. for a threshold sweep). +# +# Worktree submits: export REPO=, EXTRA_MOUNTS="$MAIN:$MAIN:rw", +# and submit with --export=ALL (mirrors run_eval_kimi_k3.sbatch). +# +#SBATCH --job-name=kimi-k3-dspark-accept +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=04:00:00 +#SBATCH --output=kimi-k3-dspark-accept-%j.log + +set -euo pipefail + +usage() { + echo "Usage: sbatch $0 --model PATH --drafter PATH --image PATH [--outdir DIR]" \ + "[--num-prompts N] [--max-tokens N] [--confidence-threshold T]" \ + "[--confidence-policy first_below|cumulative] [--skip-baseline]" +} + +MODEL="" +DRAFTER="" +CONTAINER_IMAGE="" +OUTDIR="" +NUM_PROMPTS=64 +MAX_TOKENS=256 +CONF_THRESHOLD="" +CONF_POLICY=first_below +SKIP_BASELINE=0 +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL=$2; shift 2 ;; + --model=*) MODEL=${1#*=}; shift ;; + --drafter) DRAFTER=$2; shift 2 ;; + --drafter=*) DRAFTER=${1#*=}; shift ;; + --image) CONTAINER_IMAGE=$2; shift 2 ;; + --image=*) CONTAINER_IMAGE=${1#*=}; shift ;; + --outdir) OUTDIR=$2; shift 2 ;; + --outdir=*) OUTDIR=${1#*=}; shift ;; + --num-prompts) NUM_PROMPTS=$2; shift 2 ;; + --max-tokens) MAX_TOKENS=$2; shift 2 ;; + --confidence-threshold) CONF_THRESHOLD=$2; shift 2 ;; + --confidence-policy) CONF_POLICY=$2; shift 2 ;; + --skip-baseline) SKIP_BASELINE=1; shift ;; + -h|--help) usage; exit 0 ;; + *) echo "error: unknown argument: $1" >&2; usage >&2; exit 2 ;; + esac +done + +[[ -n "$MODEL" && -n "$DRAFTER" && -n "$CONTAINER_IMAGE" ]] || { usage >&2; exit 2; } +[[ -e "$MODEL" ]] || { echo "error: model path does not exist: $MODEL" >&2; exit 2; } +[[ -e "$DRAFTER" ]] || { echo "error: drafter path does not exist: $DRAFTER" >&2; exit 2; } +[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } + +REPO=${REPO:-$SLURM_SUBMIT_DIR} +OUTDIR=${OUTDIR:-$REPO/dspark-accept-$SLURM_JOB_ID} +mkdir -p "$OUTDIR" + +CONF_ARGS="" +[[ -n "$CONF_THRESHOLD" ]] && \ + CONF_ARGS="--confidence-threshold $CONF_THRESHOLD --confidence-policy $CONF_POLICY" + +# Mount every input at its host path so the paths are valid inside too. +# EXTRA_MOUNTS (comma-separated src:dst:flags) lets git-worktree submits +# also mount the built main checkout that the worktree's artifact +# symlinks resolve into (submit with --export=ALL). +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro,$DRAFTER:$DRAFTER:ro,$OUTDIR:$OUTDIR:rw" +if [ -n "${EXTRA_MOUNTS:-}" ]; then + MOUNTS+=",$EXTRA_MOUNTS" +fi + +# User cache (GSM8K dataset under ~/.cache/huggingface): +# --container-mount-home mounts $HOME at /root, but when ~/.cache is a +# symlink onto a shared filesystem its target dead-ends inside the container +# (FileNotFoundError /root/.cache/huggingface). Mount the +# resolved cache root at its host path — the /root/.cache symlink chain +# then works too — and point HF_HOME at it explicitly. flashinfer/triton +# JIT caches stay node-local via the /tmp overrides below. +CACHE_HOST=$(readlink -f "$HOME/.cache" 2>/dev/null || echo "$HOME/.cache") +if [[ -d "$CACHE_HOST" && "$CACHE_HOST" != "$HOME/.cache" ]]; then + MOUNTS+=",$CACHE_HOST:$CACHE_HOST:rw" + # Also mount at the symlink's literal target so the in-container + # /root/.cache -> chain resolves (the target's parent dirs + # do not otherwise exist in the container). + CACHE_LINK_TARGET=$(readlink "$HOME/.cache") + if [[ -n "$CACHE_LINK_TARGET" && "$CACHE_LINK_TARGET" != "$CACHE_HOST" ]]; then + MOUNTS+=",$CACHE_HOST:$CACHE_LINK_TARGET:rw" + fi +fi +HF_CACHE_HOST="$CACHE_HOST/huggingface" + +run_leg() { + local leg=$1; shift + srun --mpi=pmix \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + bash -c " + set -x + ulimit -n 65536 + # Node-local JIT caches (shared-NFS races; see run_eval_kimi_k3.sbatch) + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} + export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0} + export HF_HOME='$HF_CACHE_HOST' + # Recorder env must be exported HERE (per rank), not inside the + # python script: trtllm-llmapi-launch pre-spawns the MPI worker + # ranks, so os.environ changes in the driver never reach the + # workers where DFlashWorker lives (symptom: no accept-stats + # files). Empty for the spec-off leg (recorder off = clean TPOT). + ${LEG_ENV_EXPORT:-true} + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + # Import tensorrt_llm from \$REPO, not from wherever the venv's + # in-place install points (worktree submits differ). + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ + '$REPO/examples/kimi_k3/measure_dspark_acceptance.py' \ + --model '$MODEL' --tp-size 16 \ + --num-prompts $NUM_PROMPTS --max-tokens $MAX_TOKENS \ + $* + " 2>&1 | tee "$OUTDIR/$leg.log" +} + +if [[ "$SKIP_BASELINE" -eq 0 ]]; then + echo "=== leg (a): spec-off TPOT reference ===" + LEG_ENV_EXPORT="" \ + run_leg spec_off --spec-off --output-json "$OUTDIR/results_spec_off.json" +fi + +echo "=== leg (b): DSpark-on measurement ===" +LEG_ENV_EXPORT="export TLLM_DFLASH_ACCEPT_STATS_DIR='$OUTDIR/accept-stats'" \ + run_leg dspark --drafter "$DRAFTER" $CONF_ARGS \ + --stats-dir "$OUTDIR/accept-stats" \ + --output-json "$OUTDIR/results_dspark.json" + +# Combined speedup summary (host-side python3, stdlib only) +python3 - "$OUTDIR" <<'EOF' +import json, os, sys +outdir = sys.argv[1] +def load(name): + p = os.path.join(outdir, name) + return json.load(open(p)) if os.path.exists(p) else None +ref, spec = load("results_spec_off.json"), load("results_dspark.json") +summary = {"spec_off": ref, "dspark": spec} +if ref and spec and ref.get("tpot_proxy_ms") and spec.get("tpot_proxy_ms"): + summary["e2e_speedup"] = ref["tpot_proxy_ms"] / spec["tpot_proxy_ms"] + print(f"E2E speedup (spec-off TPOT / dspark TPOT): {summary['e2e_speedup']:.3f}x") +json.dump(summary, open(os.path.join(outdir, "ab_summary.json"), "w"), indent=2) +print(f"A/B summary: {os.path.join(outdir, 'ab_summary.json')}") +EOF diff --git a/examples/kimi_k3/run_eval_kimi_k3.sbatch b/examples/kimi_k3/run_eval_kimi_k3.sbatch new file mode 100644 index 000000000000..5ea6ead2ddd4 --- /dev/null +++ b/examples/kimi_k3/run_eval_kimi_k3.sbatch @@ -0,0 +1,327 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Kimi K3 accuracy evaluation (GSM8K or MMMU) on 16 NVIDIA Blackwell GPUs. +# +# Submit the job with the checkpoint and container paths: +# +# sbatch examples/kimi_k3/run_eval_kimi_k3.sbatch \ +# --model /path/to/kimi-k3-checkpoint \ +# --image /path/to/tensorrt-llm-container.sqsh \ +# [--task gsm8k|mmmu] [--parallel dep|tep] +# +# --task selects the benchmark (default: gsm8k, identical behavior to the +# former run_gsm8k_kimi_k3.sbatch). +# +# --parallel selects the 16-GPU attention layout (default: dep). +# +# * dep: attention data-parallel + MoE EP16 (DEP16) — the layout the +# checked-in eval YAMLs pin and the reference scores were measured with. +# * tep: attention tensor-parallel + MoE EP16 (TEP16). Rewrites a per-job +# copy of the selected YAML with enable_attention_dp: false. Every rank +# then serves the same global batch instead of its own, so in the default +# mode max_batch_size is raised 32 -> 128 to recover eval concurrency +# (the spec-dec modes keep their own batch limits). Note: with +# attention-DP off the MoE router gate defaults to its bf16 fast path, +# which can flip borderline expert picks; export KIMI_K3_ROUTER_BF16=0 +# for apples-to-apples accuracy comparison against DEP16 references. +# +# --task mmmu evaluates the multimodal path (KimiK3ForConditionalGeneration +# with the vision tower; point --model at a K3 VL checkpoint). It passes +# --post_process_fn kimi_k3_mmmu so the MMMU letter answer is read from K3's +# <|open|>response<|sep|> channel, and keeps a 16384-token generation budget +# via --preserve_caller_max_tokens: lm-eval's MMMU default of 512 tokens +# truncates the chain-of-thought before the response channel opens, silently +# degrading the score. --max_seq_len is raised to 24576 (8192 input + 16384 +# output) for this task; explicitly passed trtllm-eval CLI flags override the +# --config YAML's max_seq_len. +# +# MMMU downloads its images from the HF hub on first use, so compute nodes +# need network access — or pre-fetch the dataset into a shared cache and +# export HF_HOME before submitting (sbatch propagates the environment): +# +# HF_HOME=/shared/hf_home python3 -c "from huggingface_hub import \ +# snapshot_download; snapshot_download('MMMU/MMMU', repo_type='dataset')" +# export HF_HOME=/shared/hf_home +# +# Pass --sa to evaluate with suffix-automaton speculative decoding +# (selects eval_extra_llm_options_sa.yaml and the SA-required +# max_batch_size of 8). +# +# Pass --reuse to additionally enable KV-cache block reuse +# (eval_extra_llm_options_reuse.yaml); chunked prefill is enabled by +# default. +# +# Pass --dflash /path/to/drafter to evaluate with DFlash speculative +# decoding (selects eval_extra_llm_options_dflash.yaml with the drafter +# path substituted in). With a synthetic drafter +# (make_synthetic_dflash_drafter.py) this is a wiring smoke only. +# +# --sa, --reuse, and --dflash are mutually exclusive. +# +# REPO defaults to the submit directory (override by exporting REPO), and +# must contain the built TensorRT-LLM checkout. Provide scheduler-specific +# options on the sbatch command line as needed. +# +#SBATCH --job-name=kimi-k3-eval +#SBATCH --partition=batch +#SBATCH --account=${account} +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-node=4 +#SBATCH --time=02:00:00 +#SBATCH --output=kimi-k3-eval-%j.log + +set -euo pipefail + +usage() { + echo "Usage: sbatch $0 --model PATH --image PATH [--task gsm8k|mmmu] [--parallel dep|tep] [--sa | --reuse | --dflash DRAFTER_PATH]" +} + +MODEL="" +CONTAINER_IMAGE="" +TASK=gsm8k +PARALLEL=dep +MODE=default +SPEC_STATS_DEFAULT=0 +DFLASH_DRAFTER="" +while [[ $# -gt 0 ]]; do + case "$1" in + --sa|--reuse) + [[ "$MODE" == default ]] || { echo "error: --sa, --reuse, and --dflash are mutually exclusive" >&2; usage >&2; exit 2; } + MODE=${1#--} + shift + ;; + --dflash) + [[ $# -ge 2 ]] || { echo "error: --dflash requires a drafter checkpoint path" >&2; usage >&2; exit 2; } + [[ "$MODE" == default ]] || { echo "error: --sa, --reuse, and --dflash are mutually exclusive" >&2; usage >&2; exit 2; } + MODE=dflash + DFLASH_DRAFTER=$2 + shift 2 + ;; + --dflash=*) + [[ "$MODE" == default ]] || { echo "error: --sa, --reuse, and --dflash are mutually exclusive" >&2; usage >&2; exit 2; } + MODE=dflash + DFLASH_DRAFTER=${1#*=} + shift + ;; + --task) + [[ $# -ge 2 ]] || { echo "error: --task requires a value" >&2; usage >&2; exit 2; } + TASK=$2 + shift 2 + ;; + --task=*) + TASK=${1#*=} + shift + ;; + --parallel) + [[ $# -ge 2 ]] || { echo "error: --parallel requires a value" >&2; usage >&2; exit 2; } + PARALLEL=$2 + shift 2 + ;; + --parallel=*) + PARALLEL=${1#*=} + shift + ;; + --model) + [[ $# -ge 2 ]] || { echo "error: --model requires a value" >&2; usage >&2; exit 2; } + MODEL=$2 + shift 2 + ;; + --model=*) + MODEL=${1#*=} + shift + ;; + --image) + [[ $# -ge 2 ]] || { echo "error: --image requires a value" >&2; usage >&2; exit 2; } + CONTAINER_IMAGE=$2 + shift 2 + ;; + --image=*) + CONTAINER_IMAGE=${1#*=} + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "error: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +[[ -n "$MODEL" ]] || { echo "error: --model is required" >&2; usage >&2; exit 2; } +[[ -n "$CONTAINER_IMAGE" ]] || { echo "error: --image is required" >&2; usage >&2; exit 2; } +[[ -e "$MODEL" ]] || { echo "error: model path does not exist: $MODEL" >&2; exit 2; } +[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } +case "$TASK" in + gsm8k|mmmu) ;; + *) echo "error: unknown --task: $TASK (expected gsm8k or mmmu)" >&2; usage >&2; exit 2 ;; +esac +case "$PARALLEL" in + dep|tep) ;; + *) echo "error: unknown --parallel: $PARALLEL (expected dep or tep)" >&2; usage >&2; exit 2 ;; +esac + +REPO=${REPO:-$SLURM_SUBMIT_DIR} + +# Per-mode LLM options. +# - sa: SA speculative decoding needs its own LLM options (eager, plain-EP; +# see the yaml) and max_batch_size <= 8 (SpeculativeState buffers). +# - reuse: KV-cache block reuse; drops --disable_kv_cache_reuse so the +# yaml's enable_block_reuse takes effect. +# - dflash: DFlash speculative decoding against the --dflash drafter; the +# drafter path is substituted into a per-job copy of the DFlash yaml +# (the checked-in file carries a placeholder). Smaller warmup/activation +# footprint and KV budget: drafter weights, capture buffer, and DFlash +# context-KV slots leave less headroom than the SA config assumes +# (0.25/8192 OOM'd in warmup on GB300). +MAX_BATCH_SIZE=32 +MAX_NUM_TOKENS=8192 +KV_FRAC=0.25 +KV_REUSE_ARG="--disable_kv_cache_reuse" +case "$MODE" in + sa) + EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_sa.yaml + MAX_BATCH_SIZE=8 + SPEC_STATS_DEFAULT=1 + ;; + reuse) + EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_reuse.yaml + KV_REUSE_ARG="" + ;; + dflash) + [[ -e "$DFLASH_DRAFTER" ]] || { echo "error: drafter path does not exist: $DFLASH_DRAFTER" >&2; exit 2; } + EVAL_CONFIG=$REPO/examples/kimi_k3/.eval_dflash_runtime.$SLURM_JOB_ID.yaml + sed "s|speculative_model: .*|speculative_model: $DFLASH_DRAFTER|" \ + "$REPO/examples/kimi_k3/eval_extra_llm_options_dflash.yaml" > "$EVAL_CONFIG" + MAX_BATCH_SIZE=8 + MAX_NUM_TOKENS=4096 + KV_FRAC=0.20 + SPEC_STATS_DEFAULT=1 + ;; + *) + EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options.yaml + ;; +esac + +# Per-task eval settings (see the header for the mmmu rationale). The task +# subcommand string is expanded into the container command below, so keep it +# on a single line. Per-sample MMMU outputs land next to the job log. +MAX_SEQ_LEN=8192 +case "$TASK" in + gsm8k) + TASK_CMD="gsm8k" + ;; + mmmu) + MAX_SEQ_LEN=24576 + TASK_CMD="mmmu --post_process_fn kimi_k3_mmmu --max_input_length 8192 --max_output_length 16384 --preserve_caller_max_tokens --output_path '$REPO/kimi-k3-mmmu-results-$SLURM_JOB_ID.json'" + ;; +esac + +# Per-parallel-layout adjustments (see the header). The TEP rewrite happens on +# a per-job copy so the checked-in YAMLs stay canonical for DEP16. In the +# default mode the batch bump also rewrites the YAML's own max_batch_size and +# the nested cuda_graph_config max_batch_size so graph coverage follows. +if [[ "$PARALLEL" == tep ]]; then + TEP_CONFIG=$REPO/examples/kimi_k3/.eval_tep_runtime.$SLURM_JOB_ID.yaml + if [[ "$MODE" == default ]]; then + MAX_BATCH_SIZE=128 + sed -e 's/^enable_attention_dp: true$/enable_attention_dp: false/' \ + -e "s/^max_batch_size: .*/max_batch_size: $MAX_BATCH_SIZE/" \ + -e "s/^ max_batch_size: .*/ max_batch_size: $MAX_BATCH_SIZE/" \ + "$EVAL_CONFIG" > "$TEP_CONFIG" + else + sed -e 's/^enable_attention_dp: true$/enable_attention_dp: false/' \ + "$EVAL_CONFIG" > "$TEP_CONFIG" + fi + EVAL_CONFIG=$TEP_CONFIG +fi + +# Mount every input at its host path so the paths are valid inside too. +# EXTRA_MOUNTS (comma-separated src:dst:flags) lets git-worktree submits +# also mount the built main checkout that the worktree's artifact +# symlinks resolve into (submit with --export=ALL). +MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro" +[[ -n "$DFLASH_DRAFTER" ]] && MOUNTS="$MOUNTS,$DFLASH_DRAFTER:$DFLASH_DRAFTER:ro" +# Optional extra container mounts (srun --container-mounts syntax). Needed +# e.g. when $REPO is a git worktree whose build artifacts (.venv-3.12, +# bindings *.so) are symlinks into the main checkout. +if [ -n "${EXTRA_MOUNTS:-}" ]; then + MOUNTS+=",$EXTRA_MOUNTS" +fi + +srun --mpi=pmix \ + --container-image="$CONTAINER_IMAGE" \ + --container-mount-home \ + --container-mounts="$MOUNTS" \ + bash -c " + set -x + ulimit -n 65536 + + # Ignore \$HOME/.local site-packages, which --container-mount-home + # leaks into the container: a user-site torch there shadows the + # container's torch and breaks bindings symbol resolution + # (libth_common.so: undefined symbol). + export PYTHONNOUSERSITE=1 + + # Node-local Triton cache: ~/.triton on shared NFS races across + # ranks during autotune JIT (stale file handles). + export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} + mkdir -p \"\$TRITON_CACHE_DIR\" + + # Node-local flashinfer cache: the default (\$HOME) is shared NFS and + # races across ranks during cubin download (stale file handles). + # These caches are cold on every run (~minutes of JIT per rank). + export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} + + # Node-local HF remote-code modules cache: all ranks share \$HOME + # (mounted at /root), and racing mkdir of + # ~/.cache/huggingface/modules at startup escapes pathlib's + # exist_ok on NFS (stale attribute cache fails the is_dir() + # recheck after EEXIST) -> 'FileExistsError: /root/.cache' spam + # from every rank. Only the trust_remote_code modules move + # node-local; the dataset/hub cache stays shared for reuse. + export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0} + + # HF hub/datasets cache. Defaults to the shared \$HOME cache provided + # by --container-mount-home; export HF_HOME before submitting to use + # a pre-fetched location instead (recommended for the MMMU images — + # see the header). + export HF_HOME=\"${HF_HOME:-\$HOME/.cache/huggingface}\" + + # Run trtllm-eval from the in-place installation created in the + # README Prerequisites section (build_wheel.py creates .venv-3.12). + export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + + # Log a running partial score every N completed responses (0 = off). + export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\" + # Cap in-flight requests: may cost some throughput vs submit-all, but yields steady partial scores (early failure signal) instead of one burst at the end (0 = off). + export TLLM_EVAL_MAX_IN_FLIGHT=\"\${TLLM_EVAL_MAX_IN_FLIGHT:-0}\" + # Spec-dec acceptance summary (AL/AR) at eval end; default ON for the + # sa/dflash modes, off otherwise (explicit env still wins either way). + export TLLM_EVAL_SPEC_STATS=\"\${TLLM_EVAL_SPEC_STATS:-$SPEC_STATS_DEFAULT}\" + + # Import tensorrt_llm from $REPO, not from wherever the venv's + # in-place install points (for git-worktree submits those differ; + # without this the job silently tests the main checkout's code). + export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" + + exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ + '$REPO/.venv-3.12/bin/trtllm-eval' \ + --model \"$MODEL\" \ + --backend pytorch \ + --tp_size 16 \ + --max_batch_size $MAX_BATCH_SIZE \ + --max_seq_len $MAX_SEQ_LEN \ + --max_num_tokens $MAX_NUM_TOKENS \ + --kv_cache_free_gpu_memory_fraction $KV_FRAC \ + $KV_REUSE_ARG \ + --trust_remote_code \ + --config '$EVAL_CONFIG' \ + $TASK_CMD + " diff --git a/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch b/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch deleted file mode 100644 index 441247201612..000000000000 --- a/examples/kimi_k3/run_gsm8k_kimi_k3.sbatch +++ /dev/null @@ -1,177 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Kimi K3 GSM8K evaluation on 16 NVIDIA Blackwell GPUs. -# -# Submit the job with the checkpoint and container paths: -# -# sbatch examples/kimi_k3/run_gsm8k_kimi_k3.sbatch \ -# --model /path/to/kimi-k3-checkpoint \ -# --image /path/to/tensorrt-llm-container.sqsh -# -# Pass --sa to evaluate with suffix-automaton (SA) speculative decoding -# (selects eval_extra_llm_options_sa.yaml, which carries the SA-required -# max_batch_size of 8). -# -# Pass --reuse to additionally enable KV-cache block reuse -# (eval_extra_llm_options_reuse.yaml); chunked prefill is enabled by -# default. -# -# --sa and --reuse are mutually exclusive. -# -# REPO defaults to the submit directory (override by exporting REPO), and -# must contain the built TensorRT-LLM checkout. The #SBATCH partition and -# account below are placeholders — edit them or override on the sbatch -# command line (sbatch --partition=... --account=...), along with any -# other scheduler-specific options your cluster requires. -# -#SBATCH --job-name=kimi-k3-gsm8k -#SBATCH --partition=batch -#SBATCH --account=${account} -#SBATCH --nodes=4 -#SBATCH --ntasks-per-node=4 -#SBATCH --gpus-per-node=4 -#SBATCH --time=02:00:00 -#SBATCH --output=kimi-k3-gsm8k-%j.log - -set -euo pipefail - -usage() { - echo "Usage: sbatch $0 --model PATH --image PATH [--sa | --reuse]" -} - -MODEL="" -CONTAINER_IMAGE="" -MODE=default -# Speculative-decoding acceptance stats: on for --sa, off otherwise. -SPEC_STATS_DEFAULT=0 -while [[ $# -gt 0 ]]; do - case "$1" in - --sa|--reuse) - [[ "$MODE" == default ]] || { echo "error: --sa and --reuse are mutually exclusive" >&2; usage >&2; exit 2; } - MODE=${1#--} - shift - ;; - --model) - [[ $# -ge 2 ]] || { echo "error: --model requires a value" >&2; usage >&2; exit 2; } - MODEL=$2 - shift 2 - ;; - --model=*) - MODEL=${1#*=} - shift - ;; - --image) - [[ $# -ge 2 ]] || { echo "error: --image requires a value" >&2; usage >&2; exit 2; } - CONTAINER_IMAGE=$2 - shift 2 - ;; - --image=*) - CONTAINER_IMAGE=${1#*=} - shift - ;; - -h|--help) - usage - exit 0 - ;; - *) - echo "error: unknown argument: $1" >&2 - usage >&2 - exit 2 - ;; - esac -done - -[[ -n "$MODEL" ]] || { echo "error: --model is required" >&2; usage >&2; exit 2; } -[[ -n "$CONTAINER_IMAGE" ]] || { echo "error: --image is required" >&2; usage >&2; exit 2; } -[[ -e "$MODEL" ]] || { echo "error: model path does not exist: $MODEL" >&2; exit 2; } -[[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } - -REPO=${REPO:-$SLURM_SUBMIT_DIR} - -# Per-mode LLM options. All engine options live in the selected yaml -# (single source of truth); trtllm-eval lets explicit CLI flags override -# --config, so none are duplicated here. -# - sa: suffix-automaton speculative decoding, which needs its own LLM -# options (overlap scheduler off, max_batch_size 8; see the yaml). -# - reuse: KV-cache block reuse (enable_block_reuse: true in the yaml). -case "$MODE" in - sa) - EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_sa.yaml - SPEC_STATS_DEFAULT=1 - ;; - reuse) - EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options_reuse.yaml - ;; - *) - EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options.yaml - ;; -esac - -# Mount every input at its host path so the paths are valid inside too. -# EXTRA_MOUNTS (comma-separated src:dst:flags) lets git-worktree submits -# also mount the built main checkout that the worktree's artifact -# symlinks resolve into (submit with --export=ALL). -MOUNTS="$REPO:$REPO:rw,$MODEL:$MODEL:ro" -# Optional extra container mounts (srun --container-mounts syntax). Needed -# e.g. when $REPO is a git worktree whose build artifacts (.venv-3.12, -# bindings *.so) are symlinks into the main checkout. -if [ -n "${EXTRA_MOUNTS:-}" ]; then - MOUNTS+=",$EXTRA_MOUNTS" -fi - -srun --mpi=pmix \ - --container-image="$CONTAINER_IMAGE" \ - --container-mount-home \ - --container-mounts="$MOUNTS" \ - bash -c " - set -x - ulimit -n 65536 - - # Node-local Triton cache: ~/.triton on shared NFS races across - # ranks during autotune JIT (stale file handles). - export TRITON_CACHE_DIR=/tmp/triton-cache-rank\${SLURM_PROCID:-0} - mkdir -p \"\$TRITON_CACHE_DIR\" - - # Node-local flashinfer cache: the default (\$HOME) is shared NFS and - # races across ranks during cubin download (stale file handles). - # These caches are cold on every run (~minutes of JIT per rank); - # the GSM8K dataset cache is \$HOME/.cache/huggingface, provided by - # --container-mount-home above. - export FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-rank\${SLURM_PROCID:-0} - - # Node-local HF remote-code modules cache: all ranks share \$HOME - # (mounted at /root), and racing mkdir of - # ~/.cache/huggingface/modules at startup escapes pathlib's - # exist_ok on NFS (stale attribute cache fails the is_dir() - # recheck after EEXIST) -> 'FileExistsError: /root/.cache' spam - # from every rank. Only the trust_remote_code modules move - # node-local; the dataset/hub cache stays shared for reuse. - export HF_MODULES_CACHE=/tmp/hf-modules-rank\${SLURM_PROCID:-0} - - # Run trtllm-eval from the in-place installation created in the - # README Prerequisites section (build_wheel.py creates .venv-3.12). - export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" - - # Log a running partial score every N completed responses (0 = off). - export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\" - # Cap in-flight requests: may cost some throughput vs submit-all, but yields steady partial scores (early failure signal) instead of one burst at the end (0 = off). - export TLLM_EVAL_MAX_IN_FLIGHT=\"\${TLLM_EVAL_MAX_IN_FLIGHT:-0}\" - # Speculative-decoding acceptance summary (acceptance length / - # rate) at eval end; on by default for --sa, off otherwise (an - # explicit env value still wins either way). - export TLLM_EVAL_SPEC_STATS=\"\${TLLM_EVAL_SPEC_STATS:-$SPEC_STATS_DEFAULT}\" - - # Import tensorrt_llm from $REPO, not from wherever the venv's - # in-place install points (for git-worktree submits those differ; - # without this the job silently tests the main checkout's code). - export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" - - exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ - '$REPO/.venv-3.12/bin/trtllm-eval' \ - --model \"$MODEL\" \ - --backend pytorch \ - --config '$EVAL_CONFIG' \ - gsm8k - " From 3fb33116752ba56f1bdf2beaddf6ec68d33b0af8 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Fri, 7 Aug 2026 15:13:32 -0700 Subject: [PATCH 4/9] [None][fix] Address K3 multimodal review comments - Pass the vision tower head count from both KimiK*VisionModel inits and the projector MLPs into _vision_requires_replication / _get_vision_tp_mapping instead of re-deriving it there with a hardcoded K3 default (a K2.5 config omitting both head-count keys could silently force a shardable 16-head tower to tp=1). - Document the runtime-dtype-dependent eps of the reference-matching default-eps RMSNorms in the K3 vision tower (norm0/norm1 and final_layernorm). - Clean up the per-job TEP runtime yaml: remove it via trap on batch script exit and gitignore the .eval_tep_runtime/.eval_dflash_runtime patterns. - Unit-test is_kimi_k3_multimodal_config (composite, language_model_only opt-out, missing/empty/non-dict sub-configs, text-only kimi_linear, wrong model_type) so released-config field renames fail loudly. Signed-off-by: Michal Guzek --- .gitignore | 5 ++ examples/kimi_k3/run_eval_kimi_k3.sbatch | 7 ++ .../_torch/models/modeling_kimi_k25.py | 55 +++++++------ .../_torch/models/modeling_kimi_k3_vl.py | 45 +++++++---- .../modeling/test_kimi_k3_config_routing.py | 77 +++++++++++++++++++ 5 files changed, 144 insertions(+), 45 deletions(-) create mode 100644 tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py diff --git a/.gitignore b/.gitignore index 9417ba9a7ebc..e5365b7e68fa 100644 --- a/.gitignore +++ b/.gitignore @@ -124,3 +124,8 @@ tests/integration/defs/stress_test/artifacts/ .claude/settings.json .plans/ + +# Kimi K3 eval harness per-job runtime configs +# (run_eval_kimi_k3.sbatch --parallel tep and --dflash rewrites) +examples/kimi_k3/.eval_tep_runtime.*.yaml +examples/kimi_k3/.eval_dflash_runtime.*.yaml diff --git a/examples/kimi_k3/run_eval_kimi_k3.sbatch b/examples/kimi_k3/run_eval_kimi_k3.sbatch index 5ea6ead2ddd4..60299481aa0e 100644 --- a/examples/kimi_k3/run_eval_kimi_k3.sbatch +++ b/examples/kimi_k3/run_eval_kimi_k3.sbatch @@ -228,7 +228,14 @@ esac # default mode the batch bump also rewrites the YAML's own max_batch_size and # the nested cuda_graph_config max_batch_size so graph coverage follows. if [[ "$PARALLEL" == tep ]]; then + # The per-job copy must live under $REPO: that is the only path guaranteed + # to be container-mounted at the same location on every rank + # ($SLURM_SUBMIT_DIR only coincides with it by default, and a node-local + # mktemp file would be invisible to the other nodes). Removed when this + # batch script exits; also gitignored in case the job dies before the + # trap fires. TEP_CONFIG=$REPO/examples/kimi_k3/.eval_tep_runtime.$SLURM_JOB_ID.yaml + trap 'rm -f "$TEP_CONFIG"' EXIT if [[ "$MODE" == default ]]; then MAX_BATCH_SIZE=128 sed -e 's/^enable_attention_dp: true$/enable_attention_dp: false/' \ diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 9190ba1d3692..50aaada5c226 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -355,7 +355,7 @@ def _gelu_tanh(x: torch.Tensor) -> torch.Tensor: return F.gelu(x, approximate="tanh") -def _vision_requires_replication(model_config: ModelConfig) -> bool: +def _vision_requires_replication(model_config: ModelConfig, num_heads: int) -> bool: """Whether the MoonViT vision tower must run replicated (tp=1) rather than tensor-parallel sharded across the attention-TP ranks. @@ -363,26 +363,20 @@ def _vision_requires_replication(model_config: ModelConfig) -> bool: * attention data-parallelism is enabled (attention weights are already replicated per rank), or - * the vision attention head count is not divisible by the attention-TP - degree (e.g. Kimi K3's 12 heads under TP16). Such a tower cannot be - TP-sharded at all, so it must be replicated instead of tripping the - ``num_heads % tp_size`` assertion in :class:`Attention`. + * ``num_heads`` — the tower's attention head count, resolved by the caller + with its per-model default (16 for K2.5, 12 for K3) — is not divisible + by the attention-TP degree (e.g. Kimi K3's 12 heads under TP16). Such a + tower cannot be TP-sharded at all, so it must be replicated instead of + tripping the ``num_heads % tp_size`` assertion in :class:`Attention`. """ mapping = model_config.mapping if mapping.enable_attention_dp: return True - vision_cfg = getattr(model_config.pretrained_config, "vision_config", - None) or {} - if not isinstance(vision_cfg, dict): - vision_cfg = (vision_cfg.to_dict() - if hasattr(vision_cfg, "to_dict") else vars(vision_cfg)) - num_heads = vision_cfg.get("vt_num_attention_heads", - vision_cfg.get("num_attention_heads", 12)) return (num_heads % mapping.tp_size) != 0 -def _get_vision_tp_mapping(model_config: ModelConfig) -> Mapping: - if not _vision_requires_replication(model_config): +def _get_vision_tp_mapping(model_config: ModelConfig, num_heads: int) -> Mapping: + if not _vision_requires_replication(model_config, num_heads): return model_config.mapping return Mapping( @@ -720,6 +714,7 @@ def __init__( model_config: ModelConfig, mm_hidden_size: int, text_hidden_size: int, + num_heads: int, merge_kernel_size: Tuple[int, int] = (2, 2), ln_eps: float = 1e-5, ) -> None: @@ -731,7 +726,7 @@ def __init__( eps=ln_eps, dtype=model_config.torch_dtype, ) - mapping = _get_vision_tp_mapping(model_config) + mapping = _get_vision_tp_mapping(model_config, num_heads) self.proj = nn.Sequential( Linear( self.merged_dim, @@ -796,13 +791,28 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo ) self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) + + # Extract vision config dict (num_heads is resolved here, with this + # model's default of 16 heads, because the TP-replication decision + # below must use the same head count the tower is built with). + vision_cfg = getattr(self.model_config.pretrained_config, "vision_config", {}) + if vision_cfg is None: + vision_cfg = {} + if not isinstance(vision_cfg, dict): + vision_cfg = ( + vision_cfg.to_dict() if hasattr(vision_cfg, "to_dict") else vars(vision_cfg) + ) + num_heads = vision_cfg.get( + "vt_num_attention_heads", vision_cfg.get("num_attention_heads", 16) + ) + # The MoonViT tower cannot be tensor-parallel sharded when its attention # head count is not divisible by the attention-TP degree (e.g. Kimi K3's # 12 heads under TP16); run the whole tower replicated (tp=1) in that # case so module construction and weight loading agree. Attention-DP # already replicates the vision tower via its own path, so leave it be. if not model_config.mapping.enable_attention_dp: - self.model_config.mapping = _get_vision_tp_mapping(model_config) + self.model_config.mapping = _get_vision_tp_mapping(model_config, num_heads) pretrained_config = self.model_config.pretrained_config model_dtype = ( getattr(pretrained_config, "torch_dtype", None) @@ -813,21 +823,9 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: model_dtype = getattr(torch, model_dtype, torch.bfloat16) pretrained_config.torch_dtype = model_dtype - # Extract vision config dict - vision_cfg = getattr(pretrained_config, "vision_config", {}) - if vision_cfg is None: - vision_cfg = {} - if not isinstance(vision_cfg, dict): - vision_cfg = ( - vision_cfg.to_dict() if hasattr(vision_cfg, "to_dict") else vars(vision_cfg) - ) - # Read HF-prefixed names with unprefixed fallback hidden_dim = vision_cfg.get("vt_hidden_size", vision_cfg.get("hidden_size", 1152)) num_layers = vision_cfg.get("vt_num_hidden_layers", vision_cfg.get("num_hidden_layers", 27)) - num_heads = vision_cfg.get( - "vt_num_attention_heads", vision_cfg.get("num_attention_heads", 16) - ) self.model_config.pretrained_config.head_dim = hidden_dim // num_heads self.model_config._frozen = True mlp_dim = vision_cfg.get("vt_intermediate_size", vision_cfg.get("intermediate_size", 4304)) @@ -878,6 +876,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: self.model_config, mm_hidden_size, self.text_hidden_size, + num_heads, self.merge_kernel_size, ln_eps, ) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index dba194e54276..25c8af1bed2a 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -48,6 +48,7 @@ MultimodalPlaceholderPlacement, register_input_processor, ) +from ...logger import logger # noqa: E402 from ..attention_backend import AttentionMetadata from ..attention_backend.utils import get_attention_backend from ..model_config import ModelConfig @@ -77,9 +78,6 @@ register_vision_encoder, ) -from ...logger import logger # noqa: E402 - - # --------------------------------------------------------------------------- # Native MoonViT3d Vision Encoder Components (K3 deltas) # --------------------------------------------------------------------------- @@ -142,7 +140,12 @@ def __init__( super().__init__() # Reference uses torch.nn.RMSNorm(hidden_dim) with default eps for the # per-layer norms; match it exactly (created in fp32, cast with the rest - # of the vision tower to the model dtype in load_weights()). + # of the vision tower to the model dtype in load_weights()). NOTE: + # eps=None resolves at *runtime* to torch.finfo(input.dtype).eps + # (~7.8e-3 for bf16 vs ~1.2e-7 for fp32), so exact parity with the + # reference holds only while both towers run in the same dtype (bf16 + # today). If the vision tower dtype ever changes, pin eps explicitly on + # both sides together. Same applies to final_layernorm below. self.norm0 = nn.RMSNorm(hidden_dim) self.norm1 = nn.RMSNorm(hidden_dim) # head_dim is taken from model_config.pretrained_config.head_dim, which @@ -204,6 +207,8 @@ def __init__( for layer_idx in range(num_layers) ] ) + # Default-eps RMSNorm to match the reference; the eps value is + # runtime-dtype-dependent — see the note on K3EncoderLayer.norm0. self.final_layernorm = nn.RMSNorm(hidden_dim) self.metadata_cls = get_attention_backend(model_config.attn_backend).Metadata self.attn_metadata: Optional[AttentionMetadata] = None @@ -223,13 +228,14 @@ def __init__( model_config: ModelConfig, mm_hidden_size: int, text_hidden_size: int, + num_heads: int, merge_kernel_size: Tuple[int, int] = (2, 2), ln_eps: float = 1e-5, ) -> None: super().__init__() kh, kw = merge_kernel_size self.merged_dim = mm_hidden_size * kh * kw - mapping = _get_vision_tp_mapping(model_config) + mapping = _get_vision_tp_mapping(model_config, num_heads) self.proj = nn.Sequential( Linear( self.merged_dim, @@ -292,13 +298,28 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo ) self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) + + # Extract vision config dict (num_heads is resolved here, with this + # model's default of 12 heads, because the TP-replication decision + # below must use the same head count the tower is built with). + vision_cfg = getattr(self.model_config.pretrained_config, "vision_config", {}) + if vision_cfg is None: + vision_cfg = {} + if not isinstance(vision_cfg, dict): + vision_cfg = ( + vision_cfg.to_dict() if hasattr(vision_cfg, "to_dict") else vars(vision_cfg) + ) + num_heads = vision_cfg.get( + "vt_num_attention_heads", vision_cfg.get("num_attention_heads", 12) + ) + # The MoonViT tower cannot be tensor-parallel sharded when its attention # head count is not divisible by the attention-TP degree (e.g. Kimi K3's # 12 heads under TP16); run the whole tower replicated (tp=1) in that # case so module construction and weight loading agree. Attention-DP # already replicates the vision tower via its own path, so leave it be. if not model_config.mapping.enable_attention_dp: - self.model_config.mapping = _get_vision_tp_mapping(model_config) + self.model_config.mapping = _get_vision_tp_mapping(model_config, num_heads) pretrained_config = self.model_config.pretrained_config model_dtype = ( getattr(pretrained_config, "torch_dtype", None) @@ -309,19 +330,8 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: model_dtype = getattr(torch, model_dtype, torch.bfloat16) pretrained_config.torch_dtype = model_dtype - vision_cfg = getattr(pretrained_config, "vision_config", {}) - if vision_cfg is None: - vision_cfg = {} - if not isinstance(vision_cfg, dict): - vision_cfg = ( - vision_cfg.to_dict() if hasattr(vision_cfg, "to_dict") else vars(vision_cfg) - ) - hidden_dim = vision_cfg.get("vt_hidden_size", vision_cfg.get("hidden_size", 1024)) num_layers = vision_cfg.get("vt_num_hidden_layers", vision_cfg.get("num_hidden_layers", 27)) - num_heads = vision_cfg.get( - "vt_num_attention_heads", vision_cfg.get("num_attention_heads", 12) - ) # K3 delta: the attention head_dim is qkv_hidden_size // num_heads (128), # NOT vt_hidden_size // num_heads. wqkv projects hidden_dim -> 3*qkv and # wo projects qkv -> hidden_dim, so q/k/v live in the qkv space. @@ -374,6 +384,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: self.model_config, mm_hidden_size, self.text_hidden_size, + num_heads, self.merge_kernel_size, ln_eps, ) diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py new file mode 100644 index 000000000000..a96c474c28ea --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Routing-predicate tests for the Kimi K3 checkpoint family. + +``is_kimi_k3_multimodal_config`` decides, for every ``model_type: kimi_k3`` +checkpoint, whether TRT-LLM keeps the composite VLM config (and routes to +``KimiK3ForConditionalGeneration``) or flattens to the text-only +``KimiLinearConfig`` path. These cases pin that contract so a field rename in +a released ``config.json`` fails loudly here instead of silently rerouting +the checkpoint. +""" + +import unittest + +from tensorrt_llm._torch.pyexecutor.config_utils import is_kimi_k3_multimodal_config + + +def _composite_config(): + """Minimal shape of the released K3 VL checkpoint config.""" + return { + "model_type": "kimi_k3", + "text_config": {"model_type": "kimi_linear", "hidden_size": 7168}, + "vision_config": {"vt_num_attention_heads": 12, "vt_hidden_size": 1024}, + } + + +class TestIsKimiK3MultimodalConfig(unittest.TestCase): + def test_composite_vlm_config_routes_multimodal(self): + self.assertTrue(is_kimi_k3_multimodal_config(_composite_config())) + + def test_language_model_only_opts_out(self): + cfg = _composite_config() + cfg["language_model_only"] = True + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + def test_language_model_only_false_stays_multimodal(self): + cfg = _composite_config() + cfg["language_model_only"] = False + self.assertTrue(is_kimi_k3_multimodal_config(cfg)) + + def test_missing_vision_config_is_text_only(self): + cfg = _composite_config() + del cfg["vision_config"] + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + def test_missing_text_config_is_not_multimodal(self): + cfg = _composite_config() + del cfg["text_config"] + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + def test_empty_dict_subconfigs_are_text_only(self): + cfg = _composite_config() + cfg["vision_config"] = {} + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + cfg = _composite_config() + cfg["text_config"] = {} + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + def test_non_dict_subconfigs_are_text_only(self): + cfg = _composite_config() + cfg["vision_config"] = "not-a-dict" + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + def test_text_only_kimi_linear_checkpoint(self): + self.assertFalse( + is_kimi_k3_multimodal_config({"model_type": "kimi_linear", "hidden_size": 7168}) + ) + + def test_other_model_type_with_subconfigs(self): + cfg = _composite_config() + cfg["model_type"] = "qwen3_5" + self.assertFalse(is_kimi_k3_multimodal_config(cfg)) + + +if __name__ == "__main__": + unittest.main() From 02392d4d18c1e0f27f6475e5ee5979de4c4f856d Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 10 Aug 2026 15:12:37 -0700 Subject: [PATCH 5/9] [None][fix] Address K3 VL review comments from 2ez4bz - Align the K2.5-vs-K3 delta table in the module docstring. - Stop propagating the text model's kv_cache_quant_algo into the vision tower's QuantConfig: the vision encoder layers keep no KV cache, so the inherited value could only steer attention kernel selection for cache-less attention (see #12851 for a past instance). - Replace the hand-rolled torch_dtype/dtype/string normalization with the shared resolve_hf_torch_dtype helper and document the write-back contract that makes model_config.torch_dtype safe everywhere after __init__. - Document why the wrapper declares _supports_sdpa and why a second __init__ on a built instance must early-return. Signed-off-by: Michal Guzek --- .../_torch/models/modeling_kimi_k3_vl.py | 48 +++++++++++-------- 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index 25c8af1bed2a..66bb5eb3a666 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -26,14 +26,14 @@ (``modeling_kimi_k25.py``); this file only carries the K3-specific deltas and reuses everything numerically identical from the K2.5 implementation: - delta | K2.5 | K3 - ---------------------+----------------------------+--------------------------- - vision norms | LayerNorm | RMSNorm (torch.nn.RMSNorm) - attention head_dim | vt_hidden_size // heads | qkv_hidden_size // heads - patch-embed conv bias| True | False - vision qkv/o/MLP bias| True | False - projector | PatchMergerMLP (pre_norm) | PatchMergerMLPV2 (post_norm) - text backbone | DeepseekV3ForCausalLM | KimiLinearForCausalLM + delta | K2.5 | K3 + ----------------------+----------------------------+----------------------------- + vision norms | LayerNorm | RMSNorm (torch.nn.RMSNorm) + attention head_dim | vt_hidden_size // heads | qkv_hidden_size // heads + patch-embed conv bias | True | False + vision qkv/o/MLP bias | True | False + projector | PatchMergerMLP (pre_norm) | PatchMergerMLPV2 (post_norm) + text backbone | DeepseekV3ForCausalLM | KimiLinearForCausalLM """ import copy @@ -54,6 +54,7 @@ from ..model_config import ModelConfig from ..modules.linear import Linear, TensorParallelMode from ..modules.mlp import MLP +from ..pyexecutor.config_utils import resolve_hf_torch_dtype from .checkpoints.base_weight_loader import ConsumableWeightsDict from .modeling_kimi_k25 import ( _MEDIA_PLACEHOLDER_TOKEN_ID, @@ -293,10 +294,11 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: self.model_config.extra_attrs = copy.copy(model_config.extra_attrs) self.model_config._frozen = False # Vision tower is not quantized (checkpoint quant ignore list covers - # vision_tower.* / mm_projector.*): keep only the kv-cache quant algo. - self.model_config.quant_config = QuantConfig( - kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo - ) + # vision_tower.* / mm_projector.*), and its encoder layers keep no KV + # cache — so do not carry the text model's kv_cache_quant_algo over: + # for cache-less vision attention it could only steer kernel selection + # (a past instance of that failure mode is #12851). + self.model_config.quant_config = QuantConfig() self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) # Extract vision config dict (num_heads is resolved here, with this @@ -321,13 +323,12 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: if not model_config.mapping.enable_attention_dp: self.model_config.mapping = _get_vision_tp_mapping(model_config, num_heads) pretrained_config = self.model_config.pretrained_config - model_dtype = ( - getattr(pretrained_config, "torch_dtype", None) - or getattr(pretrained_config, "dtype", None) - or torch.bfloat16 - ) - if isinstance(model_dtype, str): - model_dtype = getattr(torch, model_dtype, torch.bfloat16) + # Normalize the checkpoint dtype once (covers both the modern `dtype` + # and legacy `torch_dtype` field names, string forms, and "auto") and + # write the concrete torch.dtype back — after this, every access in + # this file can simply read model_config.torch_dtype, whose property + # assumes pretrained_config.torch_dtype is already concrete. + model_dtype = resolve_hf_torch_dtype(pretrained_config) or torch.bfloat16 pretrained_config.torch_dtype = model_dtype hidden_dim = vision_cfg.get("vt_hidden_size", vision_cfg.get("hidden_size", 1024)) @@ -441,12 +442,21 @@ def __init__( **kwargs, ) -> None: config = model_config.pretrained_config + # PreTrainedModel.__init__ resolves config._attn_implementation and + # rejects the default "sdpa" unless the class declares support. This + # wrapper never runs HF attention itself (TRT-LLM attention is wired + # inside the sub-modules), so declare support to keep HF init happy. self._supports_sdpa = True # Skip the K2.5 parent __init__ (it wires DeepSeek-V3 + K2.5 vision); # initialize the HF PreTrainedModel machinery directly, then build the # K3 components below. PreTrainedModel.__init__(self, config) + # Re-entry guard (mirrors the K2.5 wrapper): the tail of this __init__ + # repoints model_config.pretrained_config at the text_config and remaps + # the quant exclude list, so running the body twice on the same + # instance would fail the text_config assert below. If the text + # backbone already exists this is a re-init on a built instance: no-op. if hasattr(self, "llm"): return From c1bbaf66419503ecc8de2dd6231495eba136e4c0 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 10 Aug 2026 15:17:59 -0700 Subject: [PATCH 6/9] [None][fix] Address K3 multimodal review comments from brnguyen2 - run_eval_kimi_k3.sbatch: collect every per-job runtime config (the --dflash rewrite included) in a CLEANUP_FILES array removed by a single EXIT trap, since a second bare 'trap ... EXIT' would replace the first rather than add to it. - Deduplicate the VLM load_weights override: the K2.5 wrapper now builds its tower through a _VISION_MODEL_CLS class attribute, Kimi K3 overrides only that attribute and inherits load_weights (and the MetaInitMode deferral/recreation logic) unchanged. - extract_kimi_k3_mmmu_answer: guard on the extraction succeeding rather than the channel span being non-empty, so channel content that reduces to nothing (e.g. markdown-bold whitespace after the cascade's bold-stripping) keeps falling back instead of returning ''; add a regression test. - Unit-test _vision_requires_replication (12 heads/tp16 -> replicate, 16 heads/tp8 -> shard, attention-DP -> always replicate). - Fix a D205 docstring in the K3 lm_eval tests flagged by main's ruff coverage. Signed-off-by: Michal Guzek --- examples/kimi_k3/run_eval_kimi_k3.sbatch | 16 +++- .../_torch/models/modeling_kimi_k25.py | 8 +- .../_torch/models/modeling_kimi_k3_vl.py | 37 ++------ tensorrt_llm/evaluate/post_processing.py | 13 ++- .../modeling/test_kimi_k3_config_routing.py | 27 ++++++ tests/unittest/others/test_lm_eval.py | 93 +++++++++++++------ 6 files changed, 132 insertions(+), 62 deletions(-) diff --git a/examples/kimi_k3/run_eval_kimi_k3.sbatch b/examples/kimi_k3/run_eval_kimi_k3.sbatch index 60299481aa0e..cb23693ecf2c 100644 --- a/examples/kimi_k3/run_eval_kimi_k3.sbatch +++ b/examples/kimi_k3/run_eval_kimi_k3.sbatch @@ -75,6 +75,13 @@ set -euo pipefail +# Per-job runtime config copies (the --dflash and --parallel tep rewrites) are +# removed when this batch script exits. Every branch that writes one appends +# it here; a single trap covers them all because a second bare `trap ... EXIT` +# would *replace* the first, not add to it. +CLEANUP_FILES=() +trap '[ ${#CLEANUP_FILES[@]} -eq 0 ] || rm -f "${CLEANUP_FILES[@]}"' EXIT + usage() { echo "Usage: sbatch $0 --model PATH --image PATH [--task gsm8k|mmmu] [--parallel dep|tep] [--sa | --reuse | --dflash DRAFTER_PATH]" } @@ -199,6 +206,7 @@ case "$MODE" in EVAL_CONFIG=$REPO/examples/kimi_k3/.eval_dflash_runtime.$SLURM_JOB_ID.yaml sed "s|speculative_model: .*|speculative_model: $DFLASH_DRAFTER|" \ "$REPO/examples/kimi_k3/eval_extra_llm_options_dflash.yaml" > "$EVAL_CONFIG" + CLEANUP_FILES+=("$EVAL_CONFIG") MAX_BATCH_SIZE=8 MAX_NUM_TOKENS=4096 KV_FRAC=0.20 @@ -231,11 +239,11 @@ if [[ "$PARALLEL" == tep ]]; then # The per-job copy must live under $REPO: that is the only path guaranteed # to be container-mounted at the same location on every rank # ($SLURM_SUBMIT_DIR only coincides with it by default, and a node-local - # mktemp file would be invisible to the other nodes). Removed when this - # batch script exits; also gitignored in case the job dies before the - # trap fires. + # mktemp file would be invisible to the other nodes). Removed by the + # CLEANUP_FILES exit trap; also gitignored in case the job dies before + # the trap fires. TEP_CONFIG=$REPO/examples/kimi_k3/.eval_tep_runtime.$SLURM_JOB_ID.yaml - trap 'rm -f "$TEP_CONFIG"' EXIT + CLEANUP_FILES+=("$TEP_CONFIG") if [[ "$MODE" == default ]]; then MAX_BATCH_SIZE=128 sed -e 's/^enable_attention_dp: true$/enable_attention_dp: false/' \ diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 50aaada5c226..22484c533f2b 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -1559,6 +1559,10 @@ class KimiK25ForConditionalGeneration(PreTrainedModel): """ _LANG_PREFIX = "language_model." + # Vision tower class. Subclasses (Kimi K3) override this so the shared + # __init__/load_weights MetaInitMode deferral-and-recreation logic + # constructs their tower without re-implementing either method. + _VISION_MODEL_CLS = KimiK25VisionModel @classmethod def get_preferred_kv_cache_manager_version( @@ -1610,7 +1614,7 @@ def __init__( self.mm_encoder = None if not DISAGG: try: - mm_encoder = KimiK25VisionModel(model_config) + mm_encoder = self._VISION_MODEL_CLS(model_config) if _has_meta_tensors(mm_encoder): logger.info("Vision encoder deferred to load_weights() (MetaInitMode active)") else: @@ -1722,7 +1726,7 @@ def load_weights(self, weights) -> None: vision_model_config._frozen = False vision_model_config.pretrained_config = self._vlm_pretrained_config vision_model_config._frozen = True - self.mm_encoder = KimiK25VisionModel(vision_model_config) + self.mm_encoder = self._VISION_MODEL_CLS(vision_model_config) if self.mm_encoder is not None: self.mm_encoder.load_weights(weights) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index 66bb5eb3a666..7c72eb47a607 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -55,7 +55,6 @@ from ..modules.linear import Linear, TensorParallelMode from ..modules.mlp import MLP from ..pyexecutor.config_utils import resolve_hf_torch_dtype -from .checkpoints.base_weight_loader import ConsumableWeightsDict from .modeling_kimi_k25 import ( _MEDIA_PLACEHOLDER_TOKEN_ID, DISAGG, @@ -74,7 +73,6 @@ from .modeling_utils import ( MetaInitException, QuantConfig, - filter_weights, register_auto_model, register_vision_encoder, ) @@ -430,11 +428,14 @@ class KimiK3InputProcessor(KimiK25InputProcessor): class KimiK3ForConditionalGeneration(KimiK25ForConditionalGeneration): """Kimi K3 vision-language model: MoonViT3d + KimiLinear text backbone. - Reuses the K2.5 wrapper's spec-dec / weight-loading property forwarding and - :meth:`forward`; only ``__init__`` (vision encoder + text backbone classes) - and the deferred vision-encoder recreation in :meth:`load_weights` differ. + Reuses the K2.5 wrapper's spec-dec / weight-loading property forwarding, + :meth:`forward`, and :meth:`load_weights` (which builds the tower through + ``_VISION_MODEL_CLS``); only ``__init__`` differs (vision encoder + text + backbone classes). """ + _VISION_MODEL_CLS = KimiK3VisionModel + def __init__( self, model_config: ModelConfig[PretrainedConfig], @@ -468,7 +469,7 @@ def __init__( self.mm_encoder = None if not DISAGG: try: - mm_encoder = KimiK3VisionModel(model_config) + mm_encoder = self._VISION_MODEL_CLS(model_config) if _has_meta_tensors(mm_encoder): logger.info("Vision encoder deferred to load_weights() (MetaInitMode active)") else: @@ -514,24 +515,6 @@ def __init__( model_config.pretrained_config = self.llm.config model_config._frozen = True - def load_weights(self, weights) -> None: - """Load vision + projector + KimiLinear text weights from checkpoint.""" - if self.mm_encoder is not None and _has_meta_tensors(self.mm_encoder): - logger.info("Recreating deferred vision encoder after MetaInitMode") - self.mm_encoder = None - - if self.mm_encoder is None and not DISAGG: - vision_model_config = copy.copy(self.model_config) - vision_model_config._frozen = False - vision_model_config.pretrained_config = self._vlm_pretrained_config - vision_model_config._frozen = True - self.mm_encoder = KimiK3VisionModel(vision_model_config) - if self.mm_encoder is not None: - self.mm_encoder.load_weights(weights) - - if any(k.startswith(self._LANG_PREFIX) for k in weights): - lm_weights = filter_weights("language_model", weights) - lm_weights = ConsumableWeightsDict(lm_weights) - else: - lm_weights = weights - self.llm.load_weights(lm_weights) + # load_weights is inherited from KimiK25ForConditionalGeneration: the + # MetaInitMode deferral/recreation logic is identical and constructs the + # tower via _VISION_MODEL_CLS, which this class overrides above. diff --git a/tensorrt_llm/evaluate/post_processing.py b/tensorrt_llm/evaluate/post_processing.py index 2a1914f8c45f..eb3237d1f045 100644 --- a/tensorrt_llm/evaluate/post_processing.py +++ b/tensorrt_llm/evaluate/post_processing.py @@ -212,11 +212,18 @@ def extract_kimi_k3_mmmu_answer(text: str) -> str: if not text: return "" matches = _KIMI_K3_RESPONSE_CHANNEL_RE.findall(text) - # The last non-empty response channel is the final answer. + # Prefer the most recent response channel whose content actually yields an + # answer. Guard on the *extraction* succeeding, not just the span being + # non-empty: a channel can hold content the cascade reduces to nothing + # (e.g. markdown-bold whitespace after Step-2 stripping), and returning "" + # would mask a recoverable answer in the reasoning text before the channel. for span in reversed(matches): cleaned = _KIMI_SPECIAL_TOKEN_RE.sub(" ", span).strip() - if cleaned: - return extract_mmmu_answer(cleaned) + if not cleaned: + continue + answer = extract_mmmu_answer(cleaned) + if answer: + return answer # No usable response channel: fall back to the K2.5 path, which also handles # bare-letter direct answers and truncated thinking traces correctly. return strip_thinking_and_extract_mmmu_answer(text) diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py index a96c474c28ea..a7bc1c63eff6 100644 --- a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py +++ b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py @@ -11,7 +11,9 @@ """ import unittest +from types import SimpleNamespace +from tensorrt_llm._torch.models.modeling_kimi_k25 import _vision_requires_replication from tensorrt_llm._torch.pyexecutor.config_utils import is_kimi_k3_multimodal_config @@ -73,5 +75,30 @@ def test_other_model_type_with_subconfigs(self): self.assertFalse(is_kimi_k3_multimodal_config(cfg)) +def _vision_model_config(tp_size, enable_attention_dp): + """Minimal stand-in exposing the mapping fields the predicate reads.""" + return SimpleNamespace( + mapping=SimpleNamespace(tp_size=tp_size, enable_attention_dp=enable_attention_dp) + ) + + +class TestVisionRequiresReplication(unittest.TestCase): + """Pin the (mapping, num_heads) contract of _vision_requires_replication. + + The predicate gates whether the MoonViT tower TP-shards or runs replicated + (tp=1); a silent change to the divisibility rule or the attention-DP + short-circuit would otherwise only fail at multi-GPU runtime. + """ + + def test_k3_12_heads_under_tp16_requires_replication(self): + self.assertTrue(_vision_requires_replication(_vision_model_config(16, False), num_heads=12)) + + def test_k25_16_heads_under_tp8_shards(self): + self.assertFalse(_vision_requires_replication(_vision_model_config(8, False), num_heads=16)) + + def test_attention_dp_always_replicates(self): + self.assertTrue(_vision_requires_replication(_vision_model_config(16, True), num_heads=16)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/others/test_lm_eval.py b/tests/unittest/others/test_lm_eval.py index 3e67b283f4bb..07f8dd6bd30d 100644 --- a/tests/unittest/others/test_lm_eval.py +++ b/tests/unittest/others/test_lm_eval.py @@ -1636,15 +1636,19 @@ def test_e2e_windowed_matches_final_score(monkeypatch): # path is unchanged. from tensorrt_llm.evaluate.post_processing import ( # noqa: E402 - extract_kimi_k3_mmmu_answer, strip_thinking, - strip_thinking_and_extract_mmmu_answer) + extract_kimi_k3_mmmu_answer, + strip_thinking, + strip_thinking_and_extract_mmmu_answer, +) def _k3_output(thinking: str, answer: str) -> str: """Build a well-formed Kimi K3 channel-structured output string.""" - return (f"{thinking}<|close|>think<|sep|>" - f"<|open|>response<|sep|>{answer}<|close|>response<|sep|>" - f"<|close|>message<|sep|>") + return ( + f"{thinking}<|close|>think<|sep|>" + f"<|open|>response<|sep|>{answer}<|close|>response<|sep|>" + f"<|close|>message<|sep|>" + ) def test_k3_channel_bare_letter(): @@ -1654,14 +1658,26 @@ def test_k3_channel_bare_letter(): def test_k3_channel_real_samples_recovered(): - """Real committed mmmu_val samples the old parser scored wrong (channel has - the correct letter): accounting doc_id 7->C, 12->D, 21->A.""" - assert extract_kimi_k3_mmmu_answer( - _k3_output("...Total debits adjusted = 126,925. Final: C.", "C")) == "C" - assert extract_kimi_k3_mmmu_answer( - _k3_output("...Ending = 0. Option D. Final just D.", "D")) == "D" - assert extract_kimi_k3_mmmu_answer( - _k3_output("...commonly known by that name, so True. (A).", "A")) == "A" + """Real committed mmmu_val samples the old parser scored wrong. + + The channel holds the correct letter: accounting doc_id 7->C, 12->D, 21->A. + """ + assert ( + extract_kimi_k3_mmmu_answer( + _k3_output("...Total debits adjusted = 126,925. Final: C.", "C") + ) + == "C" + ) + assert ( + extract_kimi_k3_mmmu_answer(_k3_output("...Ending = 0. Option D. Final just D.", "D")) + == "D" + ) + assert ( + extract_kimi_k3_mmmu_answer( + _k3_output("...commonly known by that name, so True. (A).", "A") + ) + == "A" + ) def test_k3_channel_parenthesized_answer(): @@ -1677,10 +1693,14 @@ def test_k3_channel_answer_is_phrase(): def test_k3_truncated_after_channel_open(): - """Output truncated right after the response channel opened still yields the - letter (regex boundary falls back to end-of-text).""" - out = ("Some reasoning.<|close|>think<|sep|>" - "<|open|>response<|sep|>B") # cut off before <|close|>response + """Truncation right after the channel opened still yields the letter. + + The regex boundary falls back to end-of-text for the unterminated span. + """ + out = ( + "Some reasoning.<|close|>think<|sep|>" + "<|open|>response<|sep|>B" + ) # cut off before <|close|>response assert extract_kimi_k3_mmmu_answer(out) == "B" @@ -1688,7 +1708,8 @@ def test_k3_last_channel_wins(): """When multiple response channels are present, the final one is the answer.""" out = ( "r1<|close|>think<|sep|><|open|>response<|sep|>A<|close|>response<|sep|>" - "<|open|>response<|sep|>E<|close|>response<|sep|><|close|>message<|sep|>") + "<|open|>response<|sep|>E<|close|>response<|sep|><|close|>message<|sep|>" + ) assert extract_kimi_k3_mmmu_answer(out) == "E" @@ -1699,10 +1720,14 @@ def test_k3_no_channel_bare_letter_falls_back(): def test_k3_no_channel_truncated_thinking_does_not_crash(): - """A thinking trace truncated before the channel opened (finish=length): - no channel to parse -> fall back; must not raise and must not fabricate.""" - truncated = ("Let me reason step by step about this very long problem " - "that never reaches a final answer channel " * 20) + """Thinking truncated before the channel opened (finish=length). + + No channel to parse -> fall back; must not raise and must not fabricate. + """ + truncated = ( + "Let me reason step by step about this very long problem " + "that never reaches a final answer channel " * 20 + ) # Should not raise; returns whatever the fallback cascade yields (the model # genuinely did not emit an answer, so the exact value is not asserted). out = extract_kimi_k3_mmmu_answer(truncated) @@ -1714,10 +1739,14 @@ def test_k3_empty_input(): def test_k3_scrubs_residual_special_tokens(): - """Residual ``<|...|>`` tokens inside a channel span are scrubbed, not - returned as part of the answer.""" - out = ("t<|close|>think<|sep|><|open|>response<|sep|>" - "<|reserved|>D<|close|>response<|sep|><|close|>message<|sep|>") + """Residual ``<|...|>`` tokens inside a channel span are scrubbed. + + They must not be returned as part of the answer. + """ + out = ( + "t<|close|>think<|sep|><|open|>response<|sep|>" + "<|reserved|>D<|close|>response<|sep|><|close|>message<|sep|>" + ) assert extract_kimi_k3_mmmu_answer(out) == "D" @@ -1731,3 +1760,15 @@ def test_k2_5_strip_thinking_path_unchanged(): # And the K3 extractor, given a blob with no K3 response channel, # defers to the K2.5 cascade and returns the same answer. assert extract_kimi_k3_mmmu_answer(k25) == "C" + + +def test_k3_channel_extracting_to_nothing_falls_back_to_cascade(): + """A channel whose content extracts to nothing must not mask the fallback. + + "** **" survives the special-token scrub as non-empty text, but the + cascade's markdown-bold stripping reduces it to "" — the extractor must + keep scanning and recover the letter from the reasoning text instead of + returning the empty string. + """ + out = _k3_output("Elimination shows the answer is (B).", "** **") + assert extract_kimi_k3_mmmu_answer(out) == "B" From f3e5868ac0e1a9d8b21b8e9636c29b2a20c6dac0 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 10 Aug 2026 15:49:27 -0700 Subject: [PATCH 7/9] [None][fix] Address CodeRabbit review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - KimiK3VisionConfig: route the flash_attention_2 default through the attn_implementation kwarg — transformers v5 PretrainedConfig.__init__ overwrites a directly-assigned _attn_implementation with the kwarg's None default; explicit caller overrides still win. - KimiK3VisionModel: validate the vision-config switches the tower hardcodes (norm_type/mlp_type/activation_func/pos_emb_type and the three bias flags) so an unsupported checkpoint variant fails loudly at construction instead of silently building a different architecture. - run_eval_kimi_k3.sbatch: escape sed replacement metacharacters in the --dflash drafter path so paths containing \, & or | substitute literally into the per-job YAML. - run_dspark_acceptance.sbatch: validate that value-taking options have a value before reading $2 (set -u aborted with an unbound-variable error instead of a usage message). - README: note that .venv-3.12 should be substituted with .venv-. when the container ships a different Python. Signed-off-by: Michal Guzek --- examples/kimi_k3/README.md | 4 +++- examples/kimi_k3/run_dspark_acceptance.sbatch | 23 ++++++++++++------- examples/kimi_k3/run_eval_kimi_k3.sbatch | 5 +++- tensorrt_llm/_torch/configs/kimi_k3.py | 10 +++++--- .../_torch/models/modeling_kimi_k3_vl.py | 20 ++++++++++++++++ 5 files changed, 49 insertions(+), 13 deletions(-) diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index dc4e0b249da7..67055748fc2a 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -22,7 +22,9 @@ other GPU architectures may be added in a future release. for details. `build_wheel.py` creates the `.venv-3.12` virtual environment at the - repository root (named after the container's Python version). Adjust + repository root (named after the container's Python version). If your + container ships a different Python, substitute `.venv-.` + for `.venv-3.12` in every command below. Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300). - A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g. [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded diff --git a/examples/kimi_k3/run_dspark_acceptance.sbatch b/examples/kimi_k3/run_dspark_acceptance.sbatch index 6b3c56de2401..ba07b6ba377a 100644 --- a/examples/kimi_k3/run_dspark_acceptance.sbatch +++ b/examples/kimi_k3/run_dspark_acceptance.sbatch @@ -50,20 +50,27 @@ MAX_TOKENS=256 CONF_THRESHOLD="" CONF_POLICY=first_below SKIP_BASELINE=0 +# Abort with a usage error when a value-taking option has no value ($1 is the +# option name, $2 the remaining positional-argument count of the caller); +# without this, `set -u` aborts on the unbound $2 instead. +require_value() { + [[ $2 -ge 2 ]] || { echo "error: $1 requires a value" >&2; usage >&2; exit 2; } +} + while [[ $# -gt 0 ]]; do case "$1" in - --model) MODEL=$2; shift 2 ;; + --model) require_value --model $#; MODEL=$2; shift 2 ;; --model=*) MODEL=${1#*=}; shift ;; - --drafter) DRAFTER=$2; shift 2 ;; + --drafter) require_value --drafter $#; DRAFTER=$2; shift 2 ;; --drafter=*) DRAFTER=${1#*=}; shift ;; - --image) CONTAINER_IMAGE=$2; shift 2 ;; + --image) require_value --image $#; CONTAINER_IMAGE=$2; shift 2 ;; --image=*) CONTAINER_IMAGE=${1#*=}; shift ;; - --outdir) OUTDIR=$2; shift 2 ;; + --outdir) require_value --outdir $#; OUTDIR=$2; shift 2 ;; --outdir=*) OUTDIR=${1#*=}; shift ;; - --num-prompts) NUM_PROMPTS=$2; shift 2 ;; - --max-tokens) MAX_TOKENS=$2; shift 2 ;; - --confidence-threshold) CONF_THRESHOLD=$2; shift 2 ;; - --confidence-policy) CONF_POLICY=$2; shift 2 ;; + --num-prompts) require_value --num-prompts $#; NUM_PROMPTS=$2; shift 2 ;; + --max-tokens) require_value --max-tokens $#; MAX_TOKENS=$2; shift 2 ;; + --confidence-threshold) require_value --confidence-threshold $#; CONF_THRESHOLD=$2; shift 2 ;; + --confidence-policy) require_value --confidence-policy $#; CONF_POLICY=$2; shift 2 ;; --skip-baseline) SKIP_BASELINE=1; shift ;; -h|--help) usage; exit 0 ;; *) echo "error: unknown argument: $1" >&2; usage >&2; exit 2 ;; diff --git a/examples/kimi_k3/run_eval_kimi_k3.sbatch b/examples/kimi_k3/run_eval_kimi_k3.sbatch index cb23693ecf2c..952b2a5183b0 100644 --- a/examples/kimi_k3/run_eval_kimi_k3.sbatch +++ b/examples/kimi_k3/run_eval_kimi_k3.sbatch @@ -204,7 +204,10 @@ case "$MODE" in dflash) [[ -e "$DFLASH_DRAFTER" ]] || { echo "error: drafter path does not exist: $DFLASH_DRAFTER" >&2; exit 2; } EVAL_CONFIG=$REPO/examples/kimi_k3/.eval_dflash_runtime.$SLURM_JOB_ID.yaml - sed "s|speculative_model: .*|speculative_model: $DFLASH_DRAFTER|" \ + # Escape sed replacement metacharacters (\, & and the | delimiter) so + # arbitrary drafter paths substitute literally. + DFLASH_ESC=$(printf '%s' "$DFLASH_DRAFTER" | sed 's/[&|\\]/\\&/g') + sed "s|speculative_model: .*|speculative_model: $DFLASH_ESC|" \ "$REPO/examples/kimi_k3/eval_extra_llm_options_dflash.yaml" > "$EVAL_CONFIG" CLEANUP_FILES+=("$EVAL_CONFIG") MAX_BATCH_SIZE=8 diff --git a/tensorrt_llm/_torch/configs/kimi_k3.py b/tensorrt_llm/_torch/configs/kimi_k3.py index d865e7d3ab54..7327bdaf1a63 100644 --- a/tensorrt_llm/_torch/configs/kimi_k3.py +++ b/tensorrt_llm/_torch/configs/kimi_k3.py @@ -77,12 +77,10 @@ def __init__( self.vt_intermediate_size = vt_intermediate_size self.merge_kernel_size = tuple(merge_kernel_size) self.merge_type = merge_type - self._attn_implementation = _attn_implementation # MM Projector config self.mm_projector_type = mm_projector_type - self.mm_hidden_size = (mm_hidden_size - if mm_hidden_size is not None else vt_hidden_size) + self.mm_hidden_size = mm_hidden_size if mm_hidden_size is not None else vt_hidden_size self.projector_hidden_act = projector_hidden_act self.projector_ln_eps = projector_ln_eps self.text_hidden_size = text_hidden_size @@ -100,6 +98,12 @@ def __init__( self.ignore_index = ignore_index self.media_placeholder_token_id = media_placeholder_token_id + # transformers v5 PretrainedConfig.__init__ assigns + # `attn_implementation` (default None) over any `_attn_implementation` + # set beforehand, so route the default through the kwarg instead of + # assigning the private attribute directly. An explicit + # `attn_implementation` passed by the caller still wins. + kwargs.setdefault("attn_implementation", _attn_implementation) super().__init__(pad_token_id=pad_token_id, **kwargs) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py index 7c72eb47a607..9b4ce83df05a 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py @@ -357,6 +357,26 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: self.merge_kernel_size = (2, 2) self.merge_type = vision_cfg.get("merge_type", "sd2_tpool") + # This tower hardcodes the released K3 vision architecture + # (K3PatchEmbed3d: bias-free conv; K3EncoderLayer: RMSNorm + + # bias-free attention; K3VisionMLP: bias-free, gelu_tanh). Reject a + # checkpoint that asks for a variant these modules do not build, so + # the load fails loudly instead of silently producing wrong outputs. + for field, supported in ( + ("norm_type", "rmsnorm"), + ("mlp_type", "mlp2"), + ("activation_func", "gelu_pytorch_tanh"), + ("pos_emb_type", "divided_fixed"), + ("attn_bias", False), + ("patch_embed_proj_bias", False), + ("linear_bias", False), + ): + value = vision_cfg.get(field, supported) + if value != supported: + raise ValueError( + f"Kimi K3 vision tower supports {field}={supported!r}, got {value!r}" + ) + text_config = getattr(pretrained_config, "text_config", pretrained_config) self.model_dtype = model_dtype self.text_hidden_size = ( From affa3e6745c4e1ed3b8c809e83b30ca55bbbe8af Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Tue, 11 Aug 2026 16:04:16 -0700 Subject: [PATCH 8/9] [None][fix] Address follow-up CodeRabbit review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - run_dspark_acceptance.sbatch: %q-quote the arguments run_leg forwards into the nested bash -c string, so values containing whitespace (e.g. a drafter path) keep their word boundaries. - run_eval_kimi_k3.sbatch: pass --output_path without the .json suffix — the evaluator treats it as a directory, so the old value created a directory literally named *.json. - Make the job virtual-environment path configurable via TRTLLM_VENV in both Slurm scripts (defaulting to the repository-root .venv-3.12) and document the export in the README, so the documented .venv-. substitution actually works for the jobs. - Add loader-level Kimi K3 routing tests: config.json -> load_pretrained_config for the composite VLM (KimiK3Config + KimiK3ForConditionalGeneration architectures), the language_model_only opt-out, and the text-only kimi_linear flatten. Signed-off-by: Michal Guzek --- examples/kimi_k3/README.md | 4 +- examples/kimi_k3/run_dspark_acceptance.sbatch | 12 ++++- examples/kimi_k3/run_eval_kimi_k3.sbatch | 9 ++-- .../modeling/test_kimi_k3_config_routing.py | 46 ++++++++++++++++++- 4 files changed, 64 insertions(+), 7 deletions(-) diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 67055748fc2a..fc8d75879bd8 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -24,7 +24,9 @@ other GPU architectures may be added in a future release. `build_wheel.py` creates the `.venv-3.12` virtual environment at the repository root (named after the container's Python version). If your container ships a different Python, substitute `.venv-.` - for `.venv-3.12` in every command below. Adjust + for `.venv-3.12` in every command below and export + `TRTLLM_VENV=/path/to/repo/.venv-.` when submitting the + Slurm jobs (they default to the repository-root `.venv-3.12`). Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300). - A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g. [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded diff --git a/examples/kimi_k3/run_dspark_acceptance.sbatch b/examples/kimi_k3/run_dspark_acceptance.sbatch index ba07b6ba377a..83510ab1c0de 100644 --- a/examples/kimi_k3/run_dspark_acceptance.sbatch +++ b/examples/kimi_k3/run_dspark_acceptance.sbatch @@ -83,6 +83,9 @@ done [[ -e "$CONTAINER_IMAGE" ]] || { echo "error: image path does not exist: $CONTAINER_IMAGE" >&2; exit 2; } REPO=${REPO:-$SLURM_SUBMIT_DIR} +# Virtual environment created by build_wheel.py; export TRTLLM_VENV to point +# at .venv-. when the container ships another Python version. +VENV=${TRTLLM_VENV:-$REPO/.venv-3.12} OUTDIR=${OUTDIR:-$REPO/dspark-accept-$SLURM_JOB_ID} mkdir -p "$OUTDIR" @@ -121,6 +124,11 @@ HF_CACHE_HOST="$CACHE_HOST/huggingface" run_leg() { local leg=$1; shift + # %q-quote every forwarded argument: they are spliced into the nested + # bash -c string, where a bare $* would re-split values containing + # whitespace (e.g. a drafter path with spaces). + local escaped_args + printf -v escaped_args ' %q' "$@" srun --mpi=pmix \ --container-image="$CONTAINER_IMAGE" \ --container-mount-home \ @@ -140,7 +148,7 @@ run_leg() { # workers where DFlashWorker lives (symptom: no accept-stats # files). Empty for the spec-off leg (recorder off = clean TPOT). ${LEG_ENV_EXPORT:-true} - export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + export PATH=\"$VENV/bin:\$PATH\" # Import tensorrt_llm from \$REPO, not from wherever the venv's # in-place install points (worktree submits differ). export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" @@ -148,7 +156,7 @@ run_leg() { '$REPO/examples/kimi_k3/measure_dspark_acceptance.py' \ --model '$MODEL' --tp-size 16 \ --num-prompts $NUM_PROMPTS --max-tokens $MAX_TOKENS \ - $* + $escaped_args " 2>&1 | tee "$OUTDIR/$leg.log" } diff --git a/examples/kimi_k3/run_eval_kimi_k3.sbatch b/examples/kimi_k3/run_eval_kimi_k3.sbatch index 952b2a5183b0..1c269fbf93a9 100644 --- a/examples/kimi_k3/run_eval_kimi_k3.sbatch +++ b/examples/kimi_k3/run_eval_kimi_k3.sbatch @@ -175,6 +175,9 @@ case "$PARALLEL" in esac REPO=${REPO:-$SLURM_SUBMIT_DIR} +# Virtual environment created by build_wheel.py; export TRTLLM_VENV to point +# at .venv-. when the container ships another Python version. +VENV=${TRTLLM_VENV:-$REPO/.venv-3.12} # Per-mode LLM options. # - sa: SA speculative decoding needs its own LLM options (eager, plain-EP; @@ -230,7 +233,7 @@ case "$TASK" in ;; mmmu) MAX_SEQ_LEN=24576 - TASK_CMD="mmmu --post_process_fn kimi_k3_mmmu --max_input_length 8192 --max_output_length 16384 --preserve_caller_max_tokens --output_path '$REPO/kimi-k3-mmmu-results-$SLURM_JOB_ID.json'" + TASK_CMD="mmmu --post_process_fn kimi_k3_mmmu --max_input_length 8192 --max_output_length 16384 --preserve_caller_max_tokens --output_path '$REPO/kimi-k3-mmmu-results-$SLURM_JOB_ID'" ;; esac @@ -314,7 +317,7 @@ srun --mpi=pmix \ # Run trtllm-eval from the in-place installation created in the # README Prerequisites section (build_wheel.py creates .venv-3.12). - export PATH=\"$REPO/.venv-3.12/bin:\$PATH\" + export PATH=\"$VENV/bin:\$PATH\" # Log a running partial score every N completed responses (0 = off). export TLLM_EVAL_PARTIAL_SCORES_EVERY=\"\${TLLM_EVAL_PARTIAL_SCORES_EVERY:-100}\" @@ -330,7 +333,7 @@ srun --mpi=pmix \ export PYTHONPATH=\"$REPO\${PYTHONPATH:+:\$PYTHONPATH}\" exec '$REPO/tensorrt_llm/llmapi/trtllm-llmapi-launch' python3 \ - '$REPO/.venv-3.12/bin/trtllm-eval' \ + '$VENV/bin/trtllm-eval' \ --model \"$MODEL\" \ --backend pytorch \ --tp_size 16 \ diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py index a7bc1c63eff6..9321d03d7591 100644 --- a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py +++ b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py @@ -10,11 +10,17 @@ the checkpoint. """ +import json +import tempfile import unittest +from pathlib import Path from types import SimpleNamespace from tensorrt_llm._torch.models.modeling_kimi_k25 import _vision_requires_replication -from tensorrt_llm._torch.pyexecutor.config_utils import is_kimi_k3_multimodal_config +from tensorrt_llm._torch.pyexecutor.config_utils import ( + is_kimi_k3_multimodal_config, + load_pretrained_config, +) def _composite_config(): @@ -100,5 +106,43 @@ def test_attention_dp_always_replicates(self): self.assertTrue(_vision_requires_replication(_vision_model_config(16, True), num_heads=16)) +def _load_config_from_dict(cfg): + """Round-trip a raw config dict through load_pretrained_config.""" + with tempfile.TemporaryDirectory() as model_dir: + (Path(model_dir) / "config.json").write_text(json.dumps(cfg)) + return load_pretrained_config(model_dir) + + +class TestKimiK3LoaderRouting(unittest.TestCase): + """End-to-end config.json -> load_pretrained_config routing. + + Complements the pure-predicate tests above: these exercise the loader + branch itself — composite KimiK3Config construction, the architecture + assignment, and the text-only flatten — so a change to the branch (not + just the predicate) fails in unit CI. + """ + + def test_composite_checkpoint_loads_as_vlm(self): + config = _load_config_from_dict(_composite_config()) + self.assertEqual(config.architectures, ["KimiK3ForConditionalGeneration"]) + self.assertEqual(config.model_type, "kimi_k3") + self.assertEqual(config.text_config.model_type, "kimi_linear") + self.assertIsNotNone(config.vision_config) + self.assertEqual(config.vision_config.vt_num_attention_heads, 12) + + def test_language_model_only_checkpoint_flattens_to_text(self): + cfg = _composite_config() + cfg["language_model_only"] = True + config = _load_config_from_dict(cfg) + self.assertEqual(config.architectures, ["KimiLinearForCausalLM"]) + self.assertEqual(config.model_type, "kimi_linear") + + def test_text_only_kimi_linear_checkpoint_flattens(self): + config = _load_config_from_dict({"model_type": "kimi_linear", "hidden_size": 7168}) + self.assertEqual(config.architectures, ["KimiLinearForCausalLM"]) + self.assertEqual(config.model_type, "kimi_linear") + self.assertEqual(config.hidden_size, 7168) + + if __name__ == "__main__": unittest.main() From d5660960169b67453c3122928f36b8d5ca5a37a7 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Tue, 11 Aug 2026 16:19:29 -0700 Subject: [PATCH 9/9] [None][chore] Annotate the new Kimi K3 loader-routing test helpers Add the precise input/return annotations CodeRabbit's trivial nitpick requested for _load_config_from_dict and -> None on the new loader test methods, per the coding guidelines. Signed-off-by: Michal Guzek --- .../_torch/modeling/test_kimi_k3_config_routing.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py index 9321d03d7591..c8f23083dbb0 100644 --- a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py +++ b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py @@ -15,6 +15,9 @@ import unittest from pathlib import Path from types import SimpleNamespace +from typing import Any + +from transformers import PretrainedConfig from tensorrt_llm._torch.models.modeling_kimi_k25 import _vision_requires_replication from tensorrt_llm._torch.pyexecutor.config_utils import ( @@ -106,7 +109,7 @@ def test_attention_dp_always_replicates(self): self.assertTrue(_vision_requires_replication(_vision_model_config(16, True), num_heads=16)) -def _load_config_from_dict(cfg): +def _load_config_from_dict(cfg: dict[str, Any]) -> PretrainedConfig: """Round-trip a raw config dict through load_pretrained_config.""" with tempfile.TemporaryDirectory() as model_dir: (Path(model_dir) / "config.json").write_text(json.dumps(cfg)) @@ -122,7 +125,7 @@ class TestKimiK3LoaderRouting(unittest.TestCase): just the predicate) fails in unit CI. """ - def test_composite_checkpoint_loads_as_vlm(self): + def test_composite_checkpoint_loads_as_vlm(self) -> None: config = _load_config_from_dict(_composite_config()) self.assertEqual(config.architectures, ["KimiK3ForConditionalGeneration"]) self.assertEqual(config.model_type, "kimi_k3") @@ -130,14 +133,14 @@ def test_composite_checkpoint_loads_as_vlm(self): self.assertIsNotNone(config.vision_config) self.assertEqual(config.vision_config.vt_num_attention_heads, 12) - def test_language_model_only_checkpoint_flattens_to_text(self): + def test_language_model_only_checkpoint_flattens_to_text(self) -> None: cfg = _composite_config() cfg["language_model_only"] = True config = _load_config_from_dict(cfg) self.assertEqual(config.architectures, ["KimiLinearForCausalLM"]) self.assertEqual(config.model_type, "kimi_linear") - def test_text_only_kimi_linear_checkpoint_flattens(self): + def test_text_only_kimi_linear_checkpoint_flattens(self) -> None: config = _load_config_from_dict({"model_type": "kimi_linear", "hidden_size": 7168}) self.assertEqual(config.architectures, ["KimiLinearForCausalLM"]) self.assertEqual(config.model_type, "kimi_linear")