From 47f71bb6763f9ba198c5371078203670afac5852 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Sun, 2 Aug 2026 15:06:25 -0400 Subject: [PATCH 1/6] feat: add Kimi-K3 model definition and quantization example Add vendored Kimi-K3 model implementation (KimiK3ForConditionalGeneration) with configuration, vision processing, tokenization, and encoding support. Add quantization example that loads a quantization config from pretrained, adds ignore patterns for residual projections and routed experts, and passes the config as a loading argument. Also adds _apply_attn_res to default tracing ignore list for Kimi-K3 compatibility with the sequential pipeline. Co-Authored-By: Claude Opus 4.6 --- examples/quantizing_moe/kimi_k3_example.py | 87 + src/llmcompressor/args/dataset_arguments.py | 1 + src/llmcompressor/entrypoints/oneshot.py | 1 + .../modeling/kimi_k3/__init__.py | 1 + .../modeling/kimi_k3/configuration_kimi_k3.py | 286 ++++ .../modeling/kimi_k3/encoding_k3.py | 651 ++++++++ .../modeling/kimi_k3/kimi_k3_processor.py | 189 +++ .../kimi_k3/kimi_k3_vision_processing.py | 199 +++ .../modeling/kimi_k3/media_utils.py | 376 +++++ .../modeling/kimi_k3/modeling_kimi_k3.py | 1355 +++++++++++++++ .../kimi_k3/modeling_kimi_k3_linear.py | 1479 +++++++++++++++++ .../modeling/kimi_k3/modeling_kimi_linear.py | 1 + .../modeling/kimi_k3/tokenization_kimi.py | 430 +++++ 13 files changed, 5056 insertions(+) create mode 100644 examples/quantizing_moe/kimi_k3_example.py create mode 100644 src/llmcompressor/modeling/kimi_k3/__init__.py create mode 100644 src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py create mode 100644 src/llmcompressor/modeling/kimi_k3/encoding_k3.py create mode 100644 src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py create mode 100644 src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py create mode 100644 src/llmcompressor/modeling/kimi_k3/media_utils.py create mode 100644 src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py create mode 100644 src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py create mode 100644 src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py create mode 100644 src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py diff --git a/examples/quantizing_moe/kimi_k3_example.py b/examples/quantizing_moe/kimi_k3_example.py new file mode 100644 index 0000000000..de7756111b --- /dev/null +++ b/examples/quantizing_moe/kimi_k3_example.py @@ -0,0 +1,87 @@ +from compressed_tensors.quantization import QuantizationConfig +from transformers import AutoTokenizer + +from datasets import load_dataset +from llmcompressor import oneshot +from llmcompressor.modeling.kimi_k3 import KimiK3ForConditionalGeneration +from llmcompressor.modifiers.quantization import QuantizationModifier +from llmcompressor.utils import load_context + +MODEL_ID = "moonshotai/Kimi-K3" + +# Load quantization config from pretrained and add ignore patterns +# for modules that should not be quantized +qconfig = QuantizationConfig.from_pretrained(MODEL_ID) +qconfig.ignore += [ + "re:.*mlp_res_proj.*", + "re:.*self_attention_res_proj.*", + "re:.*routed_expert.*", + "re:.*output_attn_res_proj.*", +] + +# Load model with the modified quantization config +with load_context(KimiK3ForConditionalGeneration): + model = KimiK3ForConditionalGeneration.from_pretrained( + MODEL_ID, + quantization_config=qconfig, + device_map="auto", + torch_dtype="auto", + trust_remote_code=True, + ) +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + +DATASET_ID = "HuggingFaceH4/ultrachat_200k" +DATASET_SPLIT = "train_sft" +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 2048 + +# Load dataset and preprocess +ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]") +ds = ds.shuffle(seed=42) + + +def preprocess(example): + return { + "text": tokenizer.apply_chat_template( + example["messages"], + tokenize=False, + ) + } + + +ds = ds.map(preprocess) + + +def tokenize(sample): + return tokenizer( + sample["text"], + padding=False, + max_length=MAX_SEQUENCE_LENGTH, + truncation=True, + add_special_tokens=False, + ) + + +ds = ds.map(tokenize, remove_columns=ds.column_names) + +recipe = QuantizationModifier( + targets="Linear", + scheme="NVFP4", + ignore=[ + "lm_head", + r"re:.*block_sparse_moe\.gate", + "re:.*vision_tower.*", + ], +) + +oneshot( + model=model, + dataset=ds, + recipe=recipe, + max_seq_length=MAX_SEQUENCE_LENGTH, + num_calibration_samples=NUM_CALIBRATION_SAMPLES, +) + +SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4" +model.save_pretrained(SAVE_DIR) +tokenizer.save_pretrained(SAVE_DIR) diff --git a/src/llmcompressor/args/dataset_arguments.py b/src/llmcompressor/args/dataset_arguments.py index a192f5fcdf..93d14d5829 100644 --- a/src/llmcompressor/args/dataset_arguments.py +++ b/src/llmcompressor/args/dataset_arguments.py @@ -225,6 +225,7 @@ class DatasetArguments(CustomDatasetArguments): "_prepare_4d_causal_attention_mask_with_cache_position", "_update_linear_attn_mask", "project_per_layer_inputs", + "_apply_attn_res", ], metadata={ "help": "List of functions to ignore during tracing, either " diff --git a/src/llmcompressor/entrypoints/oneshot.py b/src/llmcompressor/entrypoints/oneshot.py index 3f985977d8..43a2898746 100644 --- a/src/llmcompressor/entrypoints/oneshot.py +++ b/src/llmcompressor/entrypoints/oneshot.py @@ -362,6 +362,7 @@ def oneshot( "_prepare_4d_causal_attention_mask_with_cache_position", "_update_linear_attn_mask", "project_per_layer_inputs", + "_apply_attn_res", ], sequential_targets: list[str] | None = None, sequential_offload_device: str = "cpu", diff --git a/src/llmcompressor/modeling/kimi_k3/__init__.py b/src/llmcompressor/modeling/kimi_k3/__init__.py new file mode 100644 index 0000000000..fc914f1a5e --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/__init__.py @@ -0,0 +1 @@ +from .modeling_kimi_k3 import KimiK3ForConditionalGeneration diff --git a/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py b/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py new file mode 100644 index 0000000000..efca3a0c23 --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py @@ -0,0 +1,286 @@ +from typing import Optional + +from transformers.configuration_utils import PretrainedConfig + + +class KimiLinearConfig(PretrainedConfig): + model_type = "kimi_linear" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + model_type="kimi_linear", + vocab_size=163840, + hidden_size=4096, + head_dim=None, + intermediate_size=11008, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + hidden_act="silu", + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + rope_theta=10000.0, + rope_scaling=None, + tie_word_embeddings=False, + moe_intermediate_size: Optional[int] = None, + moe_renormalize: bool = True, + moe_router_activation_func: str = "sigmoid", + num_experts: Optional[int] = None, + num_experts_per_token: Optional[int] = None, + num_shared_experts: int = 0, + routed_scaling_factor: float = 1.0, + first_k_dense_replace: int = 0, + moe_layer_freq: int = 1, + use_grouped_topk: bool = True, + num_expert_group: int = 1, + topk_group: int = 1, + q_lora_rank: Optional[int] = None, + kv_lora_rank: Optional[int] = None, + qk_nope_head_dim: Optional[int] = None, + qk_rope_head_dim: Optional[int] = None, + v_head_dim: Optional[int] = None, + mla_use_nope: Optional[bool] = False, + mla_use_output_gate: Optional[bool] = False, + num_nextn_predict_layers: int = 0, + linear_attn_config: Optional[dict] = None, + attn_res_block_size: Optional[int] = None, + latent_moe_use_norm: bool = False, + activation_situ_beta: Optional[float] = None, + activation_situ_linear_beta: Optional[float] = None, + max_position_embeddings: int = 4096, + routed_expert_hidden_size: Optional[int] = None, + topk_method: str = "noaux_tc", + **kwargs, + ): + self.model_type = model_type + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = ( + head_dim if head_dim is not None else hidden_size // num_attention_heads + ) + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.mla_use_nope = mla_use_nope + self.mla_use_output_gate = mla_use_output_gate + # moe config + self.num_experts = num_experts + self.num_experts_per_token = num_experts_per_token + self.moe_renormalize = moe_renormalize + self.num_shared_experts = num_shared_experts + self.routed_scaling_factor = routed_scaling_factor + self.moe_router_activation_func = moe_router_activation_func + assert self.moe_router_activation_func in ("softmax", "sigmoid") + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.moe_layer_freq = moe_layer_freq + self.use_grouped_topk = use_grouped_topk + self.num_expert_group = num_expert_group + self.topk_group = topk_group + self.num_nextn_predict_layers = num_nextn_predict_layers + + self.attn_res_block_size = attn_res_block_size + self.latent_moe_use_norm = latent_moe_use_norm + self.activation_situ_beta = activation_situ_beta + self.activation_situ_linear_beta = activation_situ_linear_beta + self.max_position_embeddings = max_position_embeddings + self.routed_expert_hidden_size = routed_expert_hidden_size + self.topk_method = topk_method + + if linear_attn_config is not None: + assert linear_attn_config["kda_layers"] is not None + assert linear_attn_config["full_attn_layers"] is not None + self.linear_attn_config = linear_attn_config + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + @property + def is_mla(self): + return ( + self.q_lora_rank is not None + or self.kv_lora_rank is not None + or self.qk_nope_head_dim is not None + or self.qk_rope_head_dim is not None + or self.v_head_dim is not None + or self.mla_use_nope is True + ) + + @property + def is_moe(self): + return self.num_experts is not None + + @property + def is_linear_attn(self) -> bool: + return not ( + self.linear_attn_config is None + or ( + isinstance(self.linear_attn_config, dict) + and self.linear_attn_config["kda_layers"] is not None + and len(self.linear_attn_config["kda_layers"]) == 0 + ) + ) + + def is_kda_layer(self, layer_idx: int): + return ( + self.linear_attn_config is not None + and (layer_idx + 1) in self.linear_attn_config["kda_layers"] + ) + + +class KimiK3VisionConfig(PretrainedConfig): + 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: int | None = 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=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 = 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 + + super().__init__(**kwargs) + + +class KimiK3Config(PretrainedConfig): + """Kimi-K3 model configuration. + + Args: + text_config (dict | KimiLinearConfig): Configuration for the text model. + + Vision Tower Parameters (from MoonViT3dConfig): + patch_size (int): Patch size for vision tower. + init_pos_emb_height (int): Initial position embedding height. + init_pos_emb_width (int): Initial position embedding width. + init_pos_emb_time (int): Initial position embedding time dimension. + pos_emb_type (str): Type of position embedding. + vt_num_attention_heads (int): Number of attention heads in vision tower. + vt_num_hidden_layers (int): Number of hidden layers in vision tower. + vt_hidden_size (int): Hidden size of vision tower. + vt_intermediate_size (int): Intermediate size in vision tower FFN. + merge_kernel_size (tuple): Kernel size for patch merging. + merge_type (str): Type of merge operation. + _attn_implementation (str): Attention implementation type. + + MM Projector Parameters (from MultiModalProjectorConfig): + mm_projector_type (str): Type of multimodal projector. + mm_hidden_size (int): Hidden size from vision tower (should match vt_hidden_size). + projector_hidden_act (str): Activation function for projector. + projector_ln_eps (float): Layer norm epsilon for projector. + + Other Parameters: + ignore_index (int): The ignore index for the loss function. + media_placeholder_token_id (int): The token ID to use for media placeholders. + pad_token_id (int): The token ID to use for padding. + """ + + model_type = "kimi_k3" + + def __init__( + self, + text_config: dict | KimiLinearConfig = None, + vision_config: dict | KimiK3VisionConfig = None, + # Other parameters + 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 + # Other config + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + 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/src/llmcompressor/modeling/kimi_k3/encoding_k3.py b/src/llmcompressor/modeling/kimi_k3/encoding_k3.py new file mode 100644 index 0000000000..2a8f1271ed --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/encoding_k3.py @@ -0,0 +1,651 @@ +"""Kimi K3 XTML encoding helpers. + +This module keeps chat rendering in Python. +Callers that need token IDs should consume ``EncodeSegment`` objects directly: +structural markers may be encoded as tiktoken special tokens, while user/tool +text and attribute values are encoded as ordinary text. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Iterable, Optional + +OPEN_TOKEN = "<|open|>" +CLOSE_TOKEN = "<|close|>" +SEP_TOKEN = "<|sep|>" +END_OF_MSG_TOKEN = "<|end_of_msg|>" +IMAGE_PLACEHOLDER = "<|kimi_image_placeholder|>" + +_VALID_THINKING_EFFORTS = {"low", "high", "max"} + + +@dataclass(frozen=True) +class EncodeSegment: + text: str + allow_special: bool = False + + +class _ImagePromptState: + def __init__(self, image_prompts: Optional[list[str]] = None): + self.image_prompts = image_prompts + self.index = 0 + + def next_prompt(self) -> str: + if self.image_prompts is None: + return IMAGE_PLACEHOLDER + if self.index >= len(self.image_prompts): + raise ValueError("More image placeholders than image prompts.") + prompt = self.image_prompts[self.index] + self.index += 1 + return prompt + + def assert_consumed(self) -> None: + if self.image_prompts is None: + return + if self.index != len(self.image_prompts): + raise ValueError( + f"image prompt count {len(self.image_prompts)} != " + f"consumed placeholder count {self.index}" + ) + + +def _segment(text: Any, *, allow_special: bool = False) -> list[EncodeSegment]: + text = str(text) + if not text: + return [] + return [EncodeSegment(text, allow_special=allow_special)] + + +def _control(text: str) -> list[EncodeSegment]: + return _segment(text, allow_special=True) + + +def _text(text: Any) -> list[EncodeSegment]: + return _segment(text, allow_special=False) + + +def _append_text( + segments: list[EncodeSegment], + text: Any, + image_state: _ImagePromptState, +) -> None: + text = str(text) + if text == "": + return + if image_state.image_prompts is None or IMAGE_PLACEHOLDER not in text: + segments.extend(_text(text)) + return + + parts = text.split(IMAGE_PLACEHOLDER) + for i, part in enumerate(parts): + segments.extend(_text(part)) + if i < len(parts) - 1: + segments.extend(_segment(image_state.next_prompt(), allow_special=True)) + + +def _escape_attr_value(value: Any) -> str: + return str(value).replace("&", "&").replace('"', """) + + +def _attr(key: str, value: Any) -> list[EncodeSegment]: + return ( + _text(f" {key}") + _text('="') + _text(_escape_attr_value(value)) + _text('"') + ) + + +def _open_tag(tag: str, attrs: Iterable[tuple[str, Any]] = ()) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(OPEN_TOKEN)) + segments.extend(_text(tag)) + for key, value in attrs: + segments.extend(_attr(key, value)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _close_tag(tag: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_control(CLOSE_TOKEN)) + segments.extend(_text(tag)) + segments.extend(_control(SEP_TOKEN)) + return segments + + +def _end_of_msg() -> list[EncodeSegment]: + return _control(END_OF_MSG_TOKEN) + + +def _json_compact(value: Any) -> str: + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def _is_mapping(value: Any) -> bool: + return isinstance(value, dict) + + +def _xtml_type(value: Any) -> str: + if isinstance(value, bool): + return "boolean" + if value is None: + return "null" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return "number" + if isinstance(value, str): + return "string" + if _is_mapping(value): + return "object" + return "array" + + +def _xtml_value(value: Any) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def _get_value(obj: Any, key: str, default: Any = None) -> Any: + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + +def extract_response_schema(response_format: Any) -> Any: + if response_format is None: + return None + + json_schema = _get_value(response_format, "json_schema") + if json_schema is None: + return None + + if isinstance(json_schema, dict): + return json_schema.get( + "schema", + json_schema.get("json_schema", json_schema), + ) + + schema = _get_value(json_schema, "schema") + if schema is not None: + return schema + + schema = _get_value(json_schema, "json_schema") + if schema is not None: + return schema + + return json_schema + + +def deep_sort_dict(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: deep_sort_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [deep_sort_dict(item) for item in obj] + return obj + + +def normalize_tool_arguments(arguments: Any) -> tuple[dict[str, Any], Optional[str]]: + if arguments is None: + return {}, None + if isinstance(arguments, dict): + return arguments, None + if isinstance(arguments, str): + if not arguments.strip(): + return {}, None + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + return {}, arguments + if not isinstance(parsed, dict): + raise ValueError("Kimi K3 tool call arguments must be a JSON object.") + return parsed, None + raise TypeError( + "Kimi K3 tool call arguments must be a dict or a JSON object string." + ) + + +def normalize_message(message: Any) -> Any: + if not isinstance(message, dict): + return message + + normalized = dict(message) + + tools = normalized.get("tools") + if tools is not None: + normalized["tools"] = deep_sort_dict(tools) + + tool_calls = normalized.get("tool_calls") + if not tool_calls: + return normalized + + normalized_calls = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + + tc = dict(tool_call) + function = tc.get("function") + if isinstance(function, dict): + fn = dict(function) + arguments, json_block = normalize_tool_arguments(fn.get("arguments")) + fn["arguments"] = arguments + if json_block is None: + fn.pop("_xtml_json_block", None) + else: + fn["_xtml_json_block"] = json_block + tc["function"] = fn + else: + arguments, json_block = normalize_tool_arguments(tc.get("arguments")) + tc["arguments"] = arguments + if json_block is None: + tc.pop("_xtml_json_block", None) + else: + tc["_xtml_json_block"] = json_block + normalized_calls.append(tc) + + normalized["tool_calls"] = normalized_calls + return normalized + + +def normalize_conversation(conversation: Any) -> Any: + if not isinstance(conversation, list): + return conversation + + def normalize_messages(messages: list[Any]) -> list[Any]: + return [normalize_message(message) for message in messages] + + if conversation and isinstance(conversation[0], list): + return [normalize_messages(messages) for messages in conversation] + return normalize_messages(conversation) + + +def _tool_call_id_index(tool_calls: Any) -> dict: + """Map assistant ``tool_calls[].id`` to ``(1-based position, function name)``. + + The position mirrors the chat template's enumeration over ``tool_calls`` + (every entry advances the position, even an id-less one). Duplicate ids keep + their first occurrence. + """ + index: dict = {} + if not isinstance(tool_calls, list): + return index + for position, tool_call in enumerate(tool_calls, start=1): + if not isinstance(tool_call, dict): + continue + call_id = tool_call.get("id") + if call_id is None: + continue + key = str(call_id) + if key in index: + continue + function = tool_call.get("function") + name = ( + function.get("name") + if isinstance(function, dict) + else tool_call.get("name") + ) + index[key] = (position, name) + return index + + +def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]: + """Re-sort K3 XTML tool results into assistant ``tool_calls`` order. + + Serving frameworks generally deliver tool results already in call order. A + direct Transformers caller, however, may pass OpenAI-style tool messages in any + order, so each run of consecutive tool messages is matched against the most + recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` == + ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and + sorted by the matched 1-based position. The matched call is authoritative, + so each matched message's ``tool`` is set to that call's function name -- + this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting + out of sync with the reordered position. ``index`` is still derived from the + rendered position by the chat template. A run that cannot be fully matched is + left untouched. Re-running is idempotent. + + This function is side-effect free: matched tool messages are shallow-copied + before their ``tool``/``name`` is rewritten, and every other message is + appended to the output as-is. The input list and its message objects are + never mutated. + """ + if not isinstance(messages, list): + return messages + + output: list[Any] = [] + current_index: dict = {} + i = 0 + n = len(messages) + + while i < n: + message = messages[i] + + if isinstance(message, dict) and message.get("role") == "assistant": + tool_calls = message.get("tool_calls") + current_index = _tool_call_id_index(tool_calls) if tool_calls else {} + output.append(message) + i += 1 + continue + + if not isinstance(message, dict) or message.get("role") != "tool": + output.append(message) + i += 1 + continue + + run: list[tuple] = [] # (position, original_offset, message, name) + unresolved = False + offset = 0 + while ( + i < n + and isinstance(messages[i], dict) + and messages[i].get("role") == "tool" + ): + tool_message = messages[i] + call_id = tool_message.get("tool_call_id", tool_message.get("id")) + matched = current_index.get(str(call_id)) if call_id is not None else None + if matched is None: + unresolved = True + run.append((None, offset, tool_message, None)) + else: + position, name = matched + run.append((position, offset, tool_message, name)) + offset += 1 + i += 1 + + if unresolved: + output.extend(item[2] for item in run) + else: + run.sort(key=lambda item: (item[0], item[1])) + for _, _, tool_message, name in run: + if name is None: + output.append(tool_message) + continue + # The id-matched call is authoritative: align tool (and any + # explicit name) so the rendered XTML tool attribute cannot + # disagree with the reordered position. Copy first so the + # caller's message object is never mutated. + resolved = dict(tool_message) + resolved["tool"] = name + if "name" in resolved: + resolved["name"] = name + output.append(resolved) + + return output + + +def is_batched_conversation(conversation: Any) -> bool: + return ( + isinstance(conversation, list) + and bool(conversation) + and isinstance(conversation[0], list) + ) + + +def _render_content_segments( + content: Any, + image_state: _ImagePromptState, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + if isinstance(content, str): + _append_text(segments, content, image_state) + elif content is not None: + for part in content: + if part["type"] in ["image", "image_url"]: + segments.extend(_segment(image_state.next_prompt(), allow_special=True)) + else: + _append_text(segments, part["text"], image_state) + return segments + + +def _internal_system_message(message_type: str, body: str) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + segments.extend(_open_tag("message", [("role", "system"), ("type", message_type)])) + segments.extend(_text(body.strip())) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def _render_assistant_segments( + message: dict[str, Any], + image_state: _ImagePromptState, + thinking: bool = True, +) -> list[EncodeSegment]: + segments: list[EncodeSegment] = [] + # The channel is structural: in thinking mode every assistant + # message carries the open/close tags even when there is no reasoning + # content to fill in. In non-thinking mode the channel is dropped + # entirely. + if thinking: + reasoning_content = message.get("reasoning_content") or message.get("reasoning") + segments.extend(_open_tag("think")) + if reasoning_content is not None and str(reasoning_content).strip(): + _append_text(segments, reasoning_content, image_state) + segments.extend(_close_tag("think")) + + segments.extend(_open_tag("response")) + segments.extend(_render_content_segments(message.get("content"), image_state)) + segments.extend(_close_tag("response")) + + tool_calls = message.get("tool_calls") + if tool_calls: + segments.extend(_open_tag("tools")) + for index, tool_call in enumerate(tool_calls, start=1): + fn = tool_call.get("function", tool_call) + segments.extend(_open_tag("call", [("tool", fn["name"]), ("index", index)])) + args = fn.get("arguments", {}) + json_block = fn.get("_xtml_json_block") + if json_block is not None: + segments.extend(_open_tag("json", [("type", "object")])) + _append_text(segments, json_block, image_state) + segments.extend(_close_tag("json")) + elif _is_mapping(args): + for key, value in args.items(): + segments.extend( + _open_tag( + "argument", + [("key", key), ("type", _xtml_type(value))], + ) + ) + _append_text(segments, _xtml_value(value), image_state) + segments.extend(_close_tag("argument")) + segments.extend(_close_tag("call")) + segments.extend(_close_tag("tools")) + + return segments + + +def _render_tool_declare(tools: Any, *, dynamic: bool = False) -> list[EncodeSegment]: + if dynamic: + body = ( + "## New Tools Available\n" + "The system dynamically extends the toolset via lazy-loading.\n" + "You have access to all existing and extended tools.\n" + "Here are the specs for the extended tools.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + else: + body = ( + "# Tools\n" + "Here are the available tools, described in JSONSchema.\n\n" + "```json\n" + f"{_json_compact(tools)}\n" + "```" + ) + segments: list[EncodeSegment] = [] + segments.extend( + _open_tag("message", [("role", "system"), ("type", "tool-declare")]) + ) + segments.extend(_text(body)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + return segments + + +def build_chat_segments( + messages: list[Any], + tools: Optional[list[dict]] = None, + *, + add_generation_prompt: bool = True, + thinking: bool = True, + image_prompts: Optional[list[str]] = None, + **kwargs: Any, +) -> list[EncodeSegment]: + # Re-sort tool results by tool_call_id at the lowest layer so every caller + # (processor or direct tokenizer) gets correctly ordered XTML. The helper is + # side-effect free, so the caller's message objects are left untouched. + messages = normalize_xtml_tool_result_messages(messages) + messages = normalize_conversation(messages) + tools = deep_sort_dict(tools) + + kwargs = dict(kwargs) + response_format = kwargs.get("response_format") + if "response_schema" not in kwargs: + response_schema = extract_response_schema(response_format) + if response_schema is not None: + kwargs["response_schema"] = response_schema + if kwargs.get("response_schema") is not None: + kwargs["response_schema"] = deep_sort_dict(kwargs["response_schema"]) + + image_state = _ImagePromptState(image_prompts) + segments: list[EncodeSegment] = [] + + tool_calls = None + tool_index = 0 + + if tools: + segments.extend(_render_tool_declare(tools)) + + thinking_effort = kwargs.get("thinking_effort") + if thinking and thinking_effort is not None: + assert thinking_effort in _VALID_THINKING_EFFORTS, ( + f"Unsupported thinking_effort={thinking_effort!r}; " + f"supported values are {sorted(_VALID_THINKING_EFFORTS)}." + ) + if thinking and thinking_effort in _VALID_THINKING_EFFORTS: + segments.extend( + _internal_system_message( + "thinking-effort", + "`thinking_effort` guides on how much to think in your " + "thinking channel (not including the response channel), " + "supported values include `low`, `medium`, `high`, and `max`.\n" + f"Now the system is invoked with `thinking_effort={thinking_effort}`.", + ) + ) + + for message_index, message in enumerate(messages): + if not isinstance(message, dict): + continue + + role = message["role"] + if role == "user": + attrs = [("role", "user")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend( + _render_content_segments(message.get("content"), image_state) + ) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "system" and message.get("tools"): + segments.extend(_render_tool_declare(message["tools"], dynamic=True)) + elif role == "system": + attrs = [("role", "system")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend( + _render_content_segments(message.get("content"), image_state) + ) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "tool": + tool_index += 1 + tool_name = message.get("tool", message.get("name")) + if ( + tool_name is None + and tool_calls is not None + and tool_index <= len(tool_calls) + ): + tc = tool_calls[tool_index - 1] + fn = tc.get("function", tc) + tool_name = fn["name"] + if tool_name is None: + raise ValueError( + "Kimi K3 tool messages need a resolvable tool name: " + "carry `tool`/`name`, or match a preceding assistant " + "tool_call by order." + ) + segments.extend( + _open_tag( + "message", + [("role", "tool"), ("tool", tool_name), ("index", tool_index)], + ) + ) + segments.extend( + _render_content_segments(message.get("content"), image_state) + ) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + elif role == "assistant": + tool_calls = message.get("tool_calls") + tool_index = 0 + attrs = [("role", "assistant")] + if message.get("name"): + attrs.append(("name", message["name"])) + segments.extend(_open_tag("message", attrs)) + segments.extend(_render_assistant_segments(message, image_state, thinking)) + segments.extend(_close_tag("message")) + segments.extend(_end_of_msg()) + + tool_choice = kwargs.get("tool_choice") + if tool_choice == "required": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=required`.\n" + "You MUST call tools in the next message.", + ) + ) + elif tool_choice == "none": + segments.extend( + _internal_system_message( + "tool-choice", + "The system is invoked with `tool_choice=none`.\n" + "You MUST NOT call any tools in the next message.", + ) + ) + + rf = kwargs.get("response_format") + rf_type = _get_value(rf, "type", rf) if isinstance(rf, dict) else rf + if rf_type == "json_object": + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_object`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.", + ) + ) + elif rf_type == "json_schema": + schema = _json_compact(kwargs.get("response_schema")) + segments.extend( + _internal_system_message( + "response-format", + "The system is invoked with `response_format=json_schema`.\n" + "Your response must be raw JSON data without markdown code " + "blocks (```json) or any additional formatting.\n" + "The JSON data must match the following schema:\n" + f"```json\n{schema}\n```", + ) + ) + + if add_generation_prompt: + segments.extend(_open_tag("message", [("role", "assistant")])) + segments.extend(_open_tag("think" if thinking else "response")) + + image_state.assert_consumed() + return segments diff --git a/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py b/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py new file mode 100644 index 0000000000..3248c69aba --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py @@ -0,0 +1,189 @@ +"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface. + +Chat rendering (including XTML tool-result ordering) is handled by the +tokenizer's Python encoder; this processor adds multimodal media preprocessing. +""" + +from transformers.feature_extraction_utils import BatchFeature +from transformers.processing_utils import ProcessorMixin +from transformers.utils import logging + +from .media_utils import ensure_media_type + +logger = logging.get_logger(__name__) + +# ── KimiK3Processor ─────────────────────────────────────────────────── + + +class KimiK3Processor(ProcessorMixin): + r""" + Constructs a KimiK3 processor which wraps a KimiK3 image processor + and a tokenizer into a single processor. + + [`KimiK3Processor`] offers all the functionalities of + [`KimiK3VisionProcessor`] and [`TikTokenTokenizer`]. + + Args: + image_processor ([`KimiK3VisionProcessor`], *optional*): + The image processor is a required input. + tokenizer ([`TikTokenTokenizer`], *optional*): + The tokenizer is a required input. + chat_template (`str`, *optional*): Kept for ProcessorMixin + compatibility. Kimi K3 chat encoding is implemented in Python by + the tokenizer. + """ + + attributes = ["image_processor", "tokenizer"] + valid_kwargs = ["chat_template"] + image_processor_class = "AutoImageProcessor" + tokenizer_class = "AutoTokenizer" + + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + **kwargs, + ): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + self.media_processor = image_processor + self.image_placeholder = "<|kimi_image_placeholder|>" + + # ── Media preprocessing ──────────────────────────────────────────── + + def update_raw_text(self, text: str, image_prompts: list[str]) -> str: + # Replace image placeholders + image_count = text.count(self.image_placeholder) + if image_count > 0: + assert image_count == len(image_prompts), ( + f"image placeholder count {image_count} != " + f"image_prompts count {len(image_prompts)}" + ) + text_parts = text.split(self.image_placeholder) + assert len(text_parts) == len(image_prompts) + 1 + text = "".join( + [text_parts[i] + image_prompts[i] for i in range(len(image_prompts))] + ) + text += text_parts[-1] + + return text + + def preprocess_medias(self, medias: list[dict]) -> tuple[list[dict], list[str]]: + """Process media items and generate corresponding prompts. + + Returns: + A tuple of (updated_medias, image_prompts). + """ + updated_medias = [] + image_prompts = [] + for media in medias: + if media["type"] == "image": + updated_medias.append(media) + img = ensure_media_type( + media, + transparent_bg_config=self.media_processor._transparent_bg_config, + transparent_bg_fill_stage=self.media_processor._transparent_bg_fill_stage, + )["image"] + w, h = img.size + image_prompts.append(self.media_processor.make_image_prompt(w, h)) + else: + raise ValueError(f"unsupported media type: {media['type']}") + return updated_medias, image_prompts + + # ── Main entry points ────────────────────────────────────────────── + + def __call__( + self, + messages: list[dict] = None, + medias: list[dict] = None, + text: str = None, + return_tensors: str = "pt", + **kwargs, + ) -> BatchFeature: + """ + Process multimodal inputs for Kimi-K3 model. + + Args: + messages: List of message dicts with 'role' and 'content' fields. + If provided, medias and text will be extracted automatically. + medias: Pre-extracted list of media dicts. + text: Pre-formatted text string. + return_tensors: Format of returned tensors. Default: 'pt'. + **kwargs: Additional arguments passed to apply_chat_template. + + Returns: + BatchFeature with fields: input_ids, attention_mask, + pixel_values, grid_thws. + """ + if messages is None and (medias is None or text is None): + raise ValueError("Provide either 'messages' or both 'medias' and 'text'") + + if medias is not None and text is not None: + updated_medias, image_prompts = self.preprocess_medias(medias) + preprocessed = self.media_processor.preprocess( + updated_medias, return_tensors=return_tensors + ) + text = self.update_raw_text(text, image_prompts) + text_inputs = self.tokenizer(text, return_tensors=return_tensors) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + if medias is None: + medias = self._extract_medias_from_messages(messages) + updated_medias, image_prompts = self.preprocess_medias(medias) + preprocessed = self.media_processor.preprocess( + updated_medias, return_tensors=return_tensors + ) + + if text is None: + text_inputs = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + return_tensors=return_tensors, + return_dict=True, + image_prompts=image_prompts, + **kwargs, + ) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + text = self.update_raw_text(text, image_prompts) + text_inputs = self.tokenizer(text, return_tensors=return_tensors) + return BatchFeature(data={**text_inputs, **preprocessed.data}) + + @staticmethod + def _extract_medias_from_messages(messages: list[dict]) -> list[dict]: + """Extract media items from messages in a single pass.""" + medias = [] + for msg in messages: + if msg["role"] != "user" or not msg.get("content"): + continue + + for content_part in msg["content"]: + if not isinstance(content_part, dict): + continue + + content_type = content_part.get("type") + if content_type in ["image_url", "image"]: + image_data = content_part.get(content_type) + assert ( + image_data is not None + ), f"image data is missing for content part: {content_part}" + medias.append( + { + "type": "image", + "image": image_data, + } + ) + return medias + + def apply_chat_template(self, messages, **kwargs): + return self.tokenizer.apply_chat_template(messages, **kwargs) + + def batch_decode(self, *args, **kwargs): + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + return ["input_ids", "attention_mask", "pixel_values", "grid_thws"] diff --git a/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py b/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py new file mode 100644 index 0000000000..705c3f167a --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py @@ -0,0 +1,199 @@ +"""Image processor class for Kimi-K3.""" + +import json +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from PIL import Image +from transformers.image_processing_utils import BaseImageProcessor, BatchFeature +from transformers.utils import TensorType + +from .media_utils import ( + MediaInput, + TransparentBgConfig, + _to_tensor, + ensure_media_type, + image_to_np, + navit_patchify, + navit_resize_image, + normalize, +) + + +class KimiK3VisionProcessor(BaseImageProcessor): + model_type = "kimi_k3" + + def __init__( + self, + media_proc_cfg: dict, + **kwargs, + ): + super().__init__(**kwargs) + self.media_proc_cfg = media_proc_cfg + + @property + def _transparent_bg_config(self) -> Optional[TransparentBgConfig]: + cfg = self.media_proc_cfg.get("transparent_bg_config") + if cfg is None: + return None + if isinstance(cfg, TransparentBgConfig): + return cfg + return TransparentBgConfig(**cfg) + + @property + def _transparent_bg_fill_stage(self) -> str: + return self.media_proc_cfg.get("transparent_bg_fill_stage", "before_resize") + + def media_tokens_calculator(self, media: MediaInput): + media = ensure_media_type( + media, + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + ret = self.get_resize_config(media) + return ret["num_tokens"] + + @classmethod + def make_image_prompt(cls, width: int, height: int) -> str: + """Build the K3 image placeholder with resolution info.""" + return ( + f"<|media_begin|>image {width}x{height}" + f"<|media_content|><|media_pad|><|media_end|>" + ) + + def get_resize_config(self, media_input: MediaInput) -> dict: + if media_input["type"] == "image": + w, h = media_input["image"].size + ret = navit_resize_image( + w, + h, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + self.media_proc_cfg["in_patch_limit"], + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["fixed_output_tokens"], + ) + return ret + else: + raise ValueError("Unsupported type: {}".format(media_input["type"])) + + def resize_image( + self, + image: Image.Image, + new_width: int, + new_height: int, + pad_width: int, + pad_height: int, + ) -> np.ndarray: + image_np = image_to_np( + image, + (new_width, new_height), + "resize", + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + image_np = np.pad( + image_np, + ((0, pad_height), (0, pad_width), (0, 0)), + mode="constant", + constant_values=0, + ) + return image_np + + def preprocess( + self, + medias: list[MediaInput], + return_tensors: Optional[Union[str, TensorType]] = None, + ) -> BatchFeature: + """ + Preprocess a atom vision input (images) into model-ready tensors. + + Args: + medias: List of MediaInput. + return_tensors: Desired output format ('pt', 'np', 'tf', or None). + + Returns: + BatchFeature containing 'pixel_values' and 'grid_thws' tensors. + """ + if not isinstance(medias, list): + medias = [medias] + if medias: + pixel_values = [] + for item in medias: + item = ensure_media_type( + item, + transparent_bg_config=self._transparent_bg_config, + transparent_bg_fill_stage=self._transparent_bg_fill_stage, + ) + resize_config = self.get_resize_config(item) + new_width, new_height, pad_width, pad_height = ( + resize_config["new_width"], + resize_config["new_height"], + resize_config["pad_width"], + resize_config["pad_height"], + ) + if item["type"] == "image": + image = item["image"] + image_np = self.resize_image( + image, new_width, new_height, pad_width, pad_height + ) + pixel_values.append(np.expand_dims(image_np, axis=0)) + else: + raise ValueError("Unsupported type: {}".format(item["type"])) + normalized_pixel_values = [] + image_std_inv = 1.0 / np.array(self.media_proc_cfg["image_std"]) + image_mean = np.array(self.media_proc_cfg["image_mean"]) + for pixels in pixel_values: + pixels = normalize(pixels, image_mean, image_std_inv) + pixels_and_thw = navit_patchify( + pixels, + self.media_proc_cfg["patch_size"], + ) + normalized_pixel_values.append(pixels_and_thw) + + pixel_values = torch.cat( + [ + _to_tensor(pixel_value["pixel_values"]) + for pixel_value in normalized_pixel_values + ] + ) + grid_thws = torch.cat( + [ + _to_tensor(pixel_value["grid_thw"], dtype=torch.int64).unsqueeze(0) + for pixel_value in normalized_pixel_values + ] + ) + + data = { + "pixel_values": pixel_values, + "grid_thws": grid_thws, + } + + else: + data = {} + + return BatchFeature(data=data, tensor_type=return_tensors) + + def __repr__(self): + return f"KimiK3VisionProcessor(media_proc_cfg={self.media_proc_cfg})" + + def to_dict(self) -> Dict[str, Any]: + output = super().to_dict() + output["media_proc_cfg"] = self.media_proc_cfg + if "media_processor" in output: + del output["media_processor"] + return output + + @classmethod + def from_dict(cls, config_dict: Dict[str, Any], **kwargs): + config = config_dict.copy() + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + def to_json_string(self): + dictionary = self.to_dict() + for key, value in dictionary.items(): + if hasattr(value, "tolist"): + dictionary[key] = value.tolist() + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" diff --git a/src/llmcompressor/modeling/kimi_k3/media_utils.py b/src/llmcompressor/modeling/kimi_k3/media_utils.py new file mode 100644 index 0000000000..560397b0c8 --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/media_utils.py @@ -0,0 +1,376 @@ +import base64 +import functools +import io +import math +from dataclasses import dataclass +from typing import Literal, TypedDict + +import numpy as np +from PIL import Image + + +class ImageInput(TypedDict): + type: Literal["image"] + image: Image.Image + + +MediaInput = ImageInput + + +@dataclass +class TransparentBgConfig: + """The config of the transparent background.""" + + pattern: Literal["white", "black", "gray", "chessboard"] = "black" + """The pattern of the transparent background.""" + + chessboard_square_size: int = 16 + """The size of the squares in the chessboard background.""" + + chessboard_square_on_top_left: bool = True + """Whether to start the chessboard with a white square on the top left.""" + + chessboard_white_value: int = 255 + """The value of the white pixels in the background.""" + + chessboard_gray_value: int = 200 + """The value of the gray pixels in the background.""" + + +@functools.lru_cache(maxsize=256) +def _create_chessboard_background( + height: int, + width: int, + square_size: int, + square_on_top_left: bool, + white_value: int, + gray_value: int, +) -> np.ndarray: + """Create a chessboard background.""" + bg = np.ones((height, width, 3), dtype=np.uint8) * white_value + for y in range(0, height, square_size): + for x in range(0, width, square_size): + if (y // square_size + x // square_size) % 2 == ( + 1 if square_on_top_left else 0 + ): + bg[y : y + square_size, x : x + square_size] = gray_value + return bg + + +def fill_transparent_bg_with( + image: Image.Image, + transparent_bg_config: TransparentBgConfig | None = None, +) -> Image.Image: + """Composite a (possibly) transparent image onto a configured background. + + When ``transparent_bg_config`` is ``None``, the image is simply converted + to RGB (preserving the historical behavior). Otherwise the alpha channel + is alpha-composited over a background generated according to the config. + """ + if transparent_bg_config is None: + return image.convert("RGB") + + if image.mode == "RGB": + return image + + has_alpha = "A" in image.getbands() or "transparency" in image.info + if not has_alpha: + return image.convert("RGB") + + img = np.array(image.convert("RGBA")) + height, width = img.shape[:2] + bg_pattern = transparent_bg_config.pattern + if bg_pattern == "white": + bg = np.full((height, width, 3), 255, dtype=np.uint8) + elif bg_pattern == "black": + bg = np.zeros((height, width, 3), dtype=np.uint8) + elif bg_pattern == "gray": + bg = np.full((height, width, 3), 128, dtype=np.uint8) + elif bg_pattern == "chessboard": + bg = _create_chessboard_background( + height, + width, + transparent_bg_config.chessboard_square_size, + transparent_bg_config.chessboard_square_on_top_left, + transparent_bg_config.chessboard_white_value, + transparent_bg_config.chessboard_gray_value, + ) + else: + raise ValueError(f"Invalid background pattern: {bg_pattern}") + + alpha = img[:, :, 3] + img_rgb = img[:, :, :3] + alpha_normalized = alpha.astype(np.float32) / 255.0 + alpha_3d = np.stack([alpha_normalized] * 3, axis=2) + result = alpha_3d * img_rgb + (1 - alpha_3d) * bg + result = result.astype(np.uint8) + return Image.fromarray(result) + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +): + # Apply the patch limits. + s1 = math.sqrt( + in_patch_limit + / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) + ) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) + new_w = min(new_w, patch_limit_on_one_side * patch_size) + new_h = min(new_h, patch_limit_on_one_side * patch_size) + + # Calculate the padding to make the height and width divisible by the merge kernel size and patch size. + factor = merge_kernel_size * patch_size + + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + # Calculate new dimensions after padding and patching + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + + assert ( + token_height * merge_kernel_size <= patch_limit_on_one_side + ), f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" + assert ( + token_width * merge_kernel_size <= patch_limit_on_one_side + ), f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > patch_limit_on_one_side {patch_limit_on_one_side}" + + num_tokens = token_height * token_width + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +def _to_pil( + data: str | bytes | Image.Image, + transparent_bg_config: TransparentBgConfig | None = None, + to_rgb: bool = True, +) -> Image.Image: + """Load an image and (optionally) composite its transparent background. + + Args: + data: A PIL Image, a base64 ``data:`` URL, a file path, or raw bytes. + transparent_bg_config: The config used to fill the transparent + background. ``None`` keeps the historical behavior of converting + to RGB without compositing. + to_rgb: If ``False`` the image is returned as-is (the + ``transparent_bg_config`` is ignored). The caller is then + expected to call :func:`fill_transparent_bg_with` later — e.g. + after a resize. + """ + if isinstance(data, Image.Image): + image = data + elif isinstance(data, str): + if data.startswith("data:"): + raw_base64 = data.split(",")[1] + image = Image.open(io.BytesIO(base64.b64decode(raw_base64))) + else: + image = Image.open(data) + elif isinstance(data, bytes): + image = Image.open(io.BytesIO(data)) + else: + raise ValueError(f"Unsupported data type: {type(data)}") + + if not to_rgb: + return image + + return fill_transparent_bg_with(image, transparent_bg_config) + + +def ensure_media_type( + media: MediaInput, + transparent_bg_config: TransparentBgConfig | None = None, + transparent_bg_fill_stage: Literal[ + "before_resize", "after_resize" + ] = "before_resize", +) -> MediaInput: + if media["type"] == "image": + media["image"] = _to_pil( + media["image"], + transparent_bg_config=transparent_bg_config, + to_rgb=transparent_bg_fill_stage == "before_resize", + ) + return media + else: + raise ValueError(f"Unsupported media type: {media['type']}") + + +def image_to_np( + image: Image.Image, + resize_to: tuple[int, int] | None = None, + mode: str = "resize", + raise_error_for_ill_resize: bool = True, + transparent_bg_config: TransparentBgConfig | None = None, + transparent_bg_fill_stage: Literal[ + "before_resize", "after_resize" + ] = "before_resize", +) -> np.ndarray: + """Convert an image to a numpy array. + + Args: + content: The image to convert. + resize_to: The size to resize the image to. + mode: The mode to resize the image to. + raise_error_for_ill_resize: Whether to raise an error for ill-sized resize. + transparent_bg_config: The config of the transparent background. Only + used when ``transparent_bg_fill_stage == "after_resize"`` (the + caller is responsible for filling before resize otherwise). + transparent_bg_fill_stage: When to composite the transparent + background — before or after the resize step. + + Returns: + A numpy array. + """ + assert isinstance(image, Image.Image), "image must be a PIL Image" + if resize_to is not None: + if mode == "resize": + image = image.resize(resize_to, resample=Image.Resampling.BICUBIC) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + + elif mode == "rescale_and_pad_to_center": + scale = min(resize_to[0] / image.width, resize_to[1] / image.height, 1.0) + new_width = round(image.width * scale) + new_height = round(image.height * scale) + if new_width == 0 or new_height == 0: + if raise_error_for_ill_resize: + raise ValueError( + f"Invalid resize to: {resize_to}, from image size: {image.size}" + ) + else: + return np.zeros((resize_to[1], resize_to[0], 3), dtype=np.uint8) + + image = image.resize( + (new_width, new_height), resample=Image.Resampling.BICUBIC + ) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + padding_left = (resize_to[0] - new_width) // 2 + padding_right = resize_to[0] - new_width - padding_left + padding_top = (resize_to[1] - new_height) // 2 + padding_bottom = resize_to[1] - new_height - padding_top + image = np.asarray(image) + image = np.pad( + image, + ((padding_top, padding_bottom), (padding_left, padding_right), (0, 0)), + mode="constant", + constant_values=0, + ) + assert image.shape == (resize_to[1], resize_to[0], 3) + + elif mode == "rescale_and_pad_to_rightbottom": + scale = min(resize_to[0] / image.width, resize_to[1] / image.height, 1.0) + new_width = round(image.width * scale) + new_height = round(image.height * scale) + if new_width == 0 or new_height == 0: + if raise_error_for_ill_resize: + raise ValueError( + f"Invalid resize to: {resize_to}, from image size: {image.size}" + ) + else: + return np.zeros((resize_to[1], resize_to[0], 3), dtype=np.uint8) + + image = image.resize( + (new_width, new_height), resample=Image.Resampling.BICUBIC + ) + if transparent_bg_fill_stage == "after_resize": + image = fill_transparent_bg_with(image, transparent_bg_config) + padding_right = resize_to[0] - new_width + padding_bottom = resize_to[1] - new_height + image = np.asarray(image) + image = np.pad( + image, + ((0, padding_bottom), (0, padding_right), (0, 0)), + mode="constant", + constant_values=0, + ) + assert image.shape == (resize_to[1], resize_to[0], 3) + + else: + raise ValueError(f"Invalid mode: {mode}") + + if isinstance(image, Image.Image): + return np.asarray(image) + else: + return image + + +def navit_patchify(pixel_values: np.ndarray, patch_size: int) -> dict[str, np.ndarray]: + """Reshape the pixel values to a navit shape. + + Args: + pixel_values: np.ndarray, shape (t, h, w, c) + patch_size: int + + Returns: + dict[str, np.ndarray] + - patches: np.ndarray, shape (t * h//patch_size * w//patch_size, c, patch_size, patch_size) + - grid_thw: np.ndarray, (t, h//patch_size, w//patch_size) + """ + T, H, W, C = pixel_values.shape + assert C == 3, "pixel_values must have 3 channels" + + patches = pixel_values.reshape( + T, H // patch_size, patch_size, W // patch_size, patch_size, C + ) + # (T, H//patch_size, W//patch_size, C, patch_size, patch_size) + patches = patches.transpose(0, 1, 3, 5, 2, 4) + patches = patches.reshape(-1, C, patch_size, patch_size) + grid_thw = np.array([T, H // patch_size, W // patch_size]) + return {"pixel_values": patches, "grid_thw": grid_thw} + + +def normalize( + x: np.ndarray, mean, std_inv, pixels_dtype: np.dtype = np.float32 +) -> np.ndarray: + """Normalize the image. + + Args: + x: The image to normalize. The shape is (..., 3). The dtype is uint8. The range is [0, 255]. + mean: The mean of the image. + std_inv: The inverse of the std of the image. + pixels_dtype: The dtype of the image. + Returns: + The normalized image. The shape is (..., 3). The dtype is determined by the pixels_dtype. + """ + x = (x / 255.0).astype(pixels_dtype) + x -= mean + x *= std_inv + return x + + +def _to_tensor(data, **kwargs): + import torch + + if isinstance(data, np.ndarray): + return torch.from_numpy(data).to(**kwargs) + elif isinstance(data, torch.Tensor): + return data.to(**kwargs) + elif isinstance(data, list): + return [_to_tensor(item, **kwargs) for item in data] + elif isinstance(data, tuple): + return tuple(_to_tensor(item, **kwargs) for item in data) + elif isinstance(data, dict): + return {k: _to_tensor(v, **kwargs) for k, v in data.items()} + elif data is None: + return None + else: + raise ValueError(f"Unsupported data type: {type(data)}") diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py new file mode 100644 index 0000000000..00dc8e042c --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py @@ -0,0 +1,1355 @@ +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team and HuggingFace Inc. team. All rights reserved. +# +# The code is based on llava (llava/modeling_llava.py), but modified for Kimi-K3. +# +# Licensing Information: +# - Code derived from llava (llava/modeling_llava.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 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. + + +# NOTE: Reference implementation for model architecture; see the model card for production deployment. +import math +from collections.abc import Sequence +from copy import deepcopy +from typing import Optional + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import activations + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh +from transformers.activations import PytorchGELUTanh +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_utils import PreTrainedModel +from transformers.models.llava.modeling_llava import LlavaCausalLMOutputWithPast +from transformers.utils import is_flash_attn_2_available + +from .configuration_kimi_k3 import KimiK3Config +from .modeling_kimi_linear import KimiLinearForCausalLM + +# Flash attention imports +if is_flash_attn_2_available(): + from flash_attn import flash_attn_varlen_func +else: + flash_attn_varlen_func = None + + +def multihead_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_cu_seqlens: torch.Tensor | None = None, + k_cu_seqlens: torch.Tensor | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + deterministic: bool = False, +): + """Multi-head attention using flash attention 2. + + Args: + q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim), + or (tot_seqlens, num_heads, head_dim) if packing. + q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q. + The first element should be 0 and the last element should be q.shape[0]. + k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k. + The first element should be 0 and the last element should be k.shape[0]. + + Returns: + output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing, + where dim = num_heads * head_dim + """ + attn_out = flash_attn_varlen_func( + q, + k, + v, + q_cu_seqlens, + k_cu_seqlens, + max_seqlen_q, + max_seqlen_k, + causal=False, + deterministic=deterministic, + ) + if isinstance(attn_out, tuple): + attn_out = attn_out[0] + + attn_out = attn_out.flatten(start_dim=-2) + + return attn_out + + +def eager_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_cu_seqlens: Optional[torch.Tensor] = None, + k_cu_seqlens: Optional[torch.Tensor] = None, + **kwargs, +) -> torch.Tensor: + seq_length = q.shape[0] + attention_mask = torch.zeros( + [1, seq_length, seq_length], device=q.device, dtype=torch.bool + ) + for i in range(1, len(q_cu_seqlens)): + attention_mask[ + ..., + q_cu_seqlens[i - 1] : q_cu_seqlens[i], + q_cu_seqlens[i - 1] : q_cu_seqlens[i], + ] = True + q = q.transpose(0, 1) + k = k.transpose(0, 1) + v = v.transpose(0, 1) + + attn_weight = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) + attn_weight = attn_weight.masked_fill( + ~attention_mask, torch.finfo(attn_weight.dtype).min + ) + attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32).to(q.dtype) + + attn_output = attn_weight @ v + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(seq_length, -1) + return attn_output + + +VL_VISION_ATTENTION_FUNCTIONS = { + "flash_attention_2": multihead_attention, + "eager": eager_attention, +} + + +def _apply_rope_input_validation(x, freqs_cis): + assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape) + assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape) + assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape) + assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +@get_rope_shape_decorate +@torch.compile(dynamic=True) +def get_rope_shape(org, interpolation_mode, shape): + return ( + F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ) + .squeeze(0) + .permute((1, 2, 0)) + .flatten(end_dim=1) + ) + + +def apply_rope( + xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: (The leading dimensions of all inputs should be the same) + xq: query, tensor of shape (..., num_heads, head_dim) + xk: key, tensor of shape (..., num_heads, head_dim) + freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. + Returns: + xq_out, xk_out: tensors of shape (..., num_heads, head_dim) + """ + _apply_rope_input_validation(xq, freqs_cis) + _apply_rope_input_validation(xk, freqs_cis) + + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 + # ..., num_heads, head_dim/2 + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + From: + https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): + """ + t_size: int of the temporal size + return: + pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +class Learnable2DInterpPosEmbDivided_fixed(nn.Module): + def __init__( + self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = "bicubic", + ) -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer( + "time_weight", + torch.from_numpy(get_1d_sincos_pos_embed(self.dim, self.num_frames)) + .float() + .unsqueeze(1), + persistent=False, + ) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.normal_(self.weight) + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + pos_embs = [] + for t, h, w in grid_thws.tolist(): + assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}" + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = get_rope_shape( + self.weight, + interpolation_mode=self.interpolation_mode, + shape=(h, w), + ) + + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = ( + pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t] + ) + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + out = x + torch.cat(pos_embs) + return out + + +class MoonVision3dPatchEmbed(nn.Module): + def __init__( + self, + out_dim: int, + in_dim: int = 3, + patch_size: int | tuple[int, int] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + patch_embed_proj_bias: bool = True, + pos_emb_interpolation_mode: str = "bicubic", + ): + super().__init__() + assert isinstance( + patch_size, int | Sequence + ), f"Invalid patch_size type: {type(patch_size)}" + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + assert ( + len(patch_size) == 2 + ), f"Expected patch_size to be a tuple of 2, got {patch_size}" + self.patch_size = patch_size + + self.proj = nn.Conv2d( + in_dim, + out_dim, + kernel_size=patch_size, + stride=patch_size, + bias=patch_embed_proj_bias, + ) + + if pos_emb_type == "divided_fixed": + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + height=pos_emb_height, + width=pos_emb_width, + num_frames=pos_emb_time, + dim=out_dim, + interpolation_mode=pos_emb_interpolation_mode, + ) + else: + raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + x (L, Channels): input tensor + grid_hws (N, 3): temporal, height and width + + Returns: + (L, Cout) tensor + """ + x = self.proj(x).view(x.size(0), -1) + # apply positional embedding + x = self.pos_emb(x, grid_thws) + return x + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + + This class is intended to be used in the following way: + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis. + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration. + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation. + The rope is shared across all attention layers and all heads. + + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + + Args: + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed) + max_height (int): the maximum height of the 2D grid + max_width (int): the maximum width of the 2D grid + theta_base (float): the base of the theta + device (str): the device to store the precomputed cis + """ + + def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000): + super().__init__() + self.dim = dim + assert self.dim % 4 == 0, "dim must be divisible by 4" + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = ( + torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) + ) # C/4 + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat( + [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 + ) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis( + self, grid_thws: torch.Tensor, device: torch.device + ) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + if not hasattr(self, "freqs_cis"): + self.register_buffer( + "freqs_cis", self._precompute_freqs_cis(device), persistent=False + ) + + shapes = grid_thws.tolist() + assert all( + 1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes + ), ( + shapes, + self.max_height, + self.max_width, + ) + freqs_cis = torch.cat( + [ + self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) + for t, h, w in shapes + ], + dim=0, + ) + return freqs_cis + + +class MLP2(nn.Module): + """ + Args: + dims: [in_dim, hidden_dim, out_dim] + bias: whether to use bias in linear layer. + """ + + def __init__(self, dims: list[int], activation, bias=True): + super().__init__() + assert len(dims) == 3 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + self.activation = activation + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.fc0(x) + x = self.activation(x) + return self.fc1(x) + + +class MoonViTEncoderLayer(nn.Module): + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + qkv_hidden_size: int | None = None, + norm_type: str = "layernorm", + mlp_type: str = "mlp2", + *, + attn_implementation: str = "flash_attention_2", + activation=F.gelu, + attn_bias: bool = False, + linear_bias: bool = True, + use_deterministic_attn: bool = False, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + self.qkv_hidden_size = ( + hidden_dim if qkv_hidden_size is None else qkv_hidden_size + ) + self.hidden_size_per_attention_head = self.qkv_hidden_size // self.num_heads + self.attn_implementation = attn_implementation + self.use_deterministic_attn = use_deterministic_attn + + if norm_type == "layernorm": + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + elif norm_type == "rmsnorm": + self.norm0 = nn.RMSNorm(hidden_dim) + self.norm1 = nn.RMSNorm(hidden_dim) + else: + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + if mlp_type == "mlp2": + self.mlp = MLP2( + [hidden_dim, mlp_dim, hidden_dim], activation, bias=linear_bias + ) + else: + raise NotImplementedError(f"Not support mlp_type: {mlp_type}") + + self.wqkv = nn.Linear(hidden_dim, self.qkv_hidden_size * 3, bias=attn_bias) + self.wo = nn.Linear(self.qkv_hidden_size, hidden_dim, bias=attn_bias) + + def attention_qkvpacked( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + rope_freqs_cis: torch.Tensor | None = None, + ): + """ + Args: + x (torch.Tensor): (batch_size, seqlen, hidden_dim) + cu_seqlens (torch.Tensor): + """ + xqkv = self.wqkv(x) + + qkv_shape = xqkv.size()[:-1] + ( + 3, + self.num_heads, + self.hidden_size_per_attention_head, + ) + # xqkv: (batch_size, seqlen, 3, nheads, headdim) + xqkv = xqkv.view(*qkv_shape) + xq, xk, xv = torch.unbind(xqkv, dim=-3) + + xq, xk = apply_rope(xq, xk, rope_freqs_cis) + + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation] + attn_out = attn_func( + xq, + xk, + xv, + q_cu_seqlens=cu_seqlens, + k_cu_seqlens=cu_seqlens, + max_seqlen_k=max_seqlen, + max_seqlen_q=max_seqlen, + deterministic=self.use_deterministic_attn, + ) + + attn_out = self.wo(attn_out) + return attn_out + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attention_qkvpacked( + hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class MoonViT3dEncoder(nn.Module): + def __init__( + self, + hidden_dim: int, + num_layers: int, + block_cfg: dict, + use_deterministic_attn: bool = False, + ) -> None: + super().__init__() + self.use_deterministic_attn = use_deterministic_attn + + qkv_hidden_size = ( + block_cfg["hidden_dim"] + if block_cfg.get("qkv_hidden_size") is None + else block_cfg["qkv_hidden_size"] + ) + self.rope_2d = Rope2DPosEmbRepeated( + qkv_hidden_size // block_cfg["num_heads"], 512, 512 + ) + self.blocks = nn.ModuleList( + [ + MoonViTEncoderLayer( + **block_cfg, use_deterministic_attn=self.use_deterministic_attn + ) + for _ in range(num_layers) + ] + ) + norm_type = block_cfg.get("norm_type", "layernorm") + if norm_type == "layernorm": + self.final_layernorm = nn.LayerNorm(hidden_dim) + elif norm_type == "rmsnorm": + self.final_layernorm = nn.RMSNorm(hidden_dim) + else: + raise NotImplementedError(f"Not support norm_type: {norm_type}") + + def forward( + self, + hidden_states: torch.Tensor, + grid_thws: torch.Tensor, + ) -> torch.Tensor: + rope_freqs_cis = self.rope_2d.get_freqs_cis( + grid_thws=grid_thws, device=hidden_states.device + ) + + lengths = torch.cat( + ( + torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device), + grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2], + ) + ) + + max_seqlen = lengths.max() + cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32) + for block in self.blocks: + hidden_states = block( + hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis + ) + + hidden_states = self.final_layernorm(hidden_states) + return hidden_states + + +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> list[torch.Tensor]: + d_model = x.size(-1) + + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + # Get the current sequence + seq = x[pre_sum : pre_sum + t * h * w] + # Reshape along self.merge_kernel_size and concat to the last dimension + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view( + t, new_height, kernel_height, new_width, kernel_width, d_model + ) + reshaped_seq = ( + reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) + ) # temporal pooling + padded_seq = reshaped_seq.view( + new_height * new_width, kernel_height * kernel_width, -1 + ) + outputs.append(padded_seq) + pre_sum += t * h * w + + return outputs + + +class MoonViT3dPretrainedModel(PreTrainedModel): + config_class = None + model_type = "moonvit3d" + _no_split_modules = ["MoonViTEncoderLayer"] + _supports_flash_attn_2 = True + _supports_sdpa = True + + def __init__(self, config, *inputs, **kwargs): + super().__init__(config, *inputs, **kwargs) + config = deepcopy(config) + self.merge_kernel_size = config.merge_kernel_size + self.patch_size = config.patch_size + self.merge_type = config.merge_type + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=config.hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + patch_embed_proj_bias=getattr(config, "patch_embed_proj_bias", True), + pos_emb_interpolation_mode=getattr( + config, "pos_emb_interpolation_mode", "bicubic" + ), + ) + + self.encoder = MoonViT3dEncoder( + hidden_dim=config.hidden_size, + num_layers=config.num_hidden_layers, + block_cfg={ + "num_heads": config.num_attention_heads, + "hidden_dim": config.hidden_size, + "qkv_hidden_size": getattr(config, "qkv_hidden_size", None), + "mlp_dim": config.intermediate_size, + "norm_type": getattr(config, "norm_type", "layernorm"), + "mlp_type": getattr(config, "mlp_type", "mlp2"), + "activation": PytorchGELUTanh(), + "attn_bias": getattr(config, "attn_bias", True), + "linear_bias": getattr(config, "linear_bias", True), + "attn_implementation": config._attn_implementation, + }, + use_deterministic_attn=getattr(self, "use_deterministic_attn", False), + ) + + def forward( + self, pixel_values: torch.Tensor, grid_thws: torch.Tensor + ) -> torch.Tensor: + """ + Args: + pixel_values (torch.Tensor): The input pixel values. + grid_thws (torch.Tensor): Temporal, height and width. + + Returns: + torch.Tensor: The output tokens. + """ + # grid_thws = grid_thws.to('cpu') + assert grid_thws.ndim == 2, f"grid_thws should be 2D, got {grid_thws.ndim}" + assert grid_thws.size(1) == 3, f"No support for thw: {grid_thws}" + hidden_states = self.patch_embed(pixel_values, grid_thws) + hidden_states = self.encoder(hidden_states, grid_thws) + if ( + self.merge_type == "sd2_tpool" + ): # spatial downsampling 2x with temporal pooling all + hidden_states = tpool_patch_merger( + hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size + ) + else: + raise NotImplementedError(f"Not support {self.merge_type}") + + return hidden_states + + +# ============================================================================ +# MM Projector Helper Classes (from mm_projector/modeling_mm_projectors.py) +# ============================================================================ + + +class IdentityMap(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, *args, **kwargs): + return x + + +class MLP(nn.Module): + def __init__(self, config): + super().__init__() + # TODO, use faster LayerNorm + self.pre_norm = nn.LayerNorm(config.mm_hidden_size) + self.proj = nn.Sequential( + nn.Linear(config.mm_hidden_size, config.hidden_size), + nn.GELU(), + nn.Linear(config.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + assert isinstance(x, list | tuple), f"x is not a list or tuple: {type(x)}" + lengths = [item.shape[0] for item in x] + x = torch.cat(x, dim=0) + x = self.pre_norm(x) + x = self.proj(x) + x = torch.split(x, lengths, dim=0) + + return x + + +class PatchMergerMLP(nn.Module): + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * ( + config.merge_kernel_size[0] * config.merge_kernel_size[1] + ) + self.pre_norm = nn.LayerNorm(config.mm_hidden_size, eps=eps) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size), + ) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + x = [self.proj(self.pre_norm(item).view(item.shape[0], -1)) for item in x] + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(self.pre_norm(x).view(B, -1, self.hidden_size)) + return x + + +class PatchMergerMLPV2(nn.Module): + def __init__(self, config): + super().__init__() + eps = config.projector_ln_eps + self.hidden_size = config.mm_hidden_size * ( + config.merge_kernel_size[0] * config.merge_kernel_size[1] + ) + self.proj = nn.Sequential( + nn.Linear(self.hidden_size, self.hidden_size, bias=False), + nn.GELU(), + nn.Linear(self.hidden_size, config.hidden_size, bias=False), + ) + self.post_norm = nn.RMSNorm(config.hidden_size, eps=eps) + for m in self.proj.modules(): + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x, *args, **kwargs): + if isinstance(x, list) or isinstance(x, tuple): + lengths = [item.shape[0] for item in x] + x = torch.concat([item.view(item.shape[0], -1) for item in x], dim=0) + x = self.post_norm(self.proj(x)) + x = torch.split(x, lengths, dim=0) + else: + # B, N, N_k, C = x.shape + B = x.shape[0] + x = self.proj(x.view(B, -1, self.hidden_size)) + x = self.post_norm(x) + return x + + +class KimiK3PreTrainedModel(PreTrainedModel): + config_class = KimiK3Config + base_model_prefix = "model" + _no_split_modules = [ + "MoonViT3dPretrainedModel", + "MoonViTEncoderLayer", + "KimiDecoderLayer", + "PatchMergerMLP", + "PatchMergerMLPV2", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = False + + def _init_weights(self, module): + # HOTFIX: disk offloading attempts to initialize the meta tensors + # but this is bad programming: we shouldn't be initializing these + # params in the first place + # the init attempt attempts to get `module.weight`, which DNE for qmodels + + return + + # important: this ported version of Llava isn't meant for training from scratch - only + # inference and fine-tuning - so the proper init weights code has been removed - the original codebase + # https://github.com/haotian-liu/LLaVA/tree/main/llava should serve for that purpose + std = ( + self.config.initializer_range + if hasattr(self.config, "initializer_range") + else self.config.text_config.initializer_range + ) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class VisionTowerConfig(PretrainedConfig): + model_type = "moonvit3d" + + def __init__(self, config: KimiK3Config, **kwargs): + super().__init__(**kwargs) + self.patch_size = config.patch_size + self.init_pos_emb_height = config.init_pos_emb_height + self.init_pos_emb_width = config.init_pos_emb_width + self.init_pos_emb_time = config.init_pos_emb_time + self.pos_emb_type = config.pos_emb_type + self.num_attention_heads = config.vt_num_attention_heads + self.num_hidden_layers = config.vt_num_hidden_layers + self.hidden_size = config.vt_hidden_size + self.intermediate_size = config.vt_intermediate_size + self.merge_kernel_size = config.merge_kernel_size + self.merge_type = config.merge_type + self._attn_implementation = config._attn_implementation + self.qkv_hidden_size = getattr(config, "qkv_hidden_size", None) + self.norm_type = getattr(config, "norm_type", "layernorm") + self.attn_bias = getattr(config, "attn_bias", True) + self.patch_embed_proj_bias = getattr(config, "patch_embed_proj_bias", True) + self.mlp_type = getattr(config, "mlp_type", "mlp2") + self.linear_bias = getattr(config, "linear_bias", True) + self.pos_emb_interpolation_mode = getattr( + config, "pos_emb_interpolation_mode", "bilinear" + ) + + +class ProjectorConfig: + def __init__(self, config: KimiK3Config): + self.mm_projector_type = config.mm_projector_type + self.mm_hidden_size = config.mm_hidden_size + self.hidden_size = config.text_hidden_size + self.merge_kernel_size = config.merge_kernel_size + self.projector_hidden_act = config.projector_hidden_act + self.projector_ln_eps = config.projector_ln_eps + + +# ref https://github.com/huggingface/transformers/blob/78b2929c0554b79e0489b451ce4ece14d265ead2/src/transformers/models/llava/modeling_llava.py#L240 +class KimiK3ForConditionalGeneration(KimiK3PreTrainedModel): + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + def __init__(self, config: KimiK3Config): + super().__init__(config) + + vt_config = VisionTowerConfig(config.vision_config) + self.vision_tower = MoonViT3dPretrainedModel(vt_config) + + proj_config = ProjectorConfig(config.vision_config) + if proj_config.mm_projector_type == "identity": + self.mm_projector = IdentityMap() + elif proj_config.mm_projector_type == "mlp": + self.mm_projector = MLP(proj_config) + elif proj_config.mm_projector_type == "patchmerger": + self.mm_projector = PatchMergerMLP(proj_config) + elif proj_config.mm_projector_type == "patchmergerv2": + self.mm_projector = PatchMergerMLPV2(proj_config) + else: + raise ValueError( + f"Unsupported mm_projector_type: {proj_config.mm_projector_type}" + ) + + self.language_model = KimiLinearForCausalLM(config.text_config) + self.post_init() + + if hasattr(self.language_model, "dtype"): + target_dtype = self.language_model.dtype + self.vision_tower = self.vision_tower.to(dtype=target_dtype) + self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + def get_input_embeddings(self): + return self.language_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.language_model.set_input_embeddings(value) + + def get_output_embeddings(self): + return self.language_model.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + self.language_model.set_output_embeddings(new_embeddings) + + def set_decoder(self, decoder): + self.language_model.set_decoder(decoder) + + def get_decoder(self): + return self.language_model.get_decoder() + + def tie_weights(self, **kwargs): + return self.language_model.tie_weights(**kwargs) + + def resize_token_embeddings( + self, new_num_tokens: int | None = None, pad_to_multiple_of=None + ) -> nn.Embedding: + model_embeds = self.language_model.resize_token_embeddings( + new_num_tokens, pad_to_multiple_of + ) + # update vocab size + self.config.text_config.vocab_size = model_embeds.num_embeddings + self.vocab_size = model_embeds.num_embeddings + return model_embeds + + def _merge_input_ids_with_image_features( + self, + image_features: list[torch.Tensor], + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + labels: torch.Tensor | None = None, + ): + """ + Args: + image_features (:obj:`torch.Tensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The image features to merge with the input embeddings. + inputs_embeds (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length, embed_dim)`): + The input embeddings. + input_ids (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The input ids. + attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`): + The attention mask. + labels (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, *optional*): + The labels. + """ + _, embed_dim = image_features[0].shape + feature_lengths = [x.shape[0] for x in image_features] + image_features = torch.cat(image_features, dim=0) + + image_token_index: int = self.config.media_placeholder_token_id + pad_token_id: int = self.config.pad_token_id + ignore_index: int = self.config.ignore_index + + batch_size, sequence_length = input_ids.shape + left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(pad_token_id)) + + # 1. Create a mask to know where special image tokens are + _token_occupation_table = torch.ones_like(input_ids.flatten()) + _token_occupation_table[input_ids.flatten() == image_token_index] = ( + torch.tensor(feature_lengths, dtype=torch.long, device=input_ids.device) + ) + _token_occupation_table = _token_occupation_table.reshape(input_ids.shape) + + max_embed_dim = _token_occupation_table.sum(-1).max().item() + assert ( + max_embed_dim >= sequence_length + ), f"The maximum embedding dimension ({max_embed_dim}) is less than the sequence length ({sequence_length})" + batch_indices, non_image_indices = torch.where(input_ids != image_token_index) + + # 2. Compute the positions where text should be written + # Calculate new positions for text tokens in merged image-text sequence. + new_token_positions = torch.cumsum(_token_occupation_table, -1) - 1 + nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] + if left_padding: + new_token_positions += nb_image_pad[:, None] # offset for left padding + text_to_overwrite = new_token_positions[batch_indices, non_image_indices] + + # 3. Create the full embedding, already padded to the maximum position + final_embedding = torch.zeros( + batch_size, + max_embed_dim, + embed_dim, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + final_attention_mask = torch.zeros( + batch_size, + max_embed_dim, + dtype=attention_mask.dtype, + device=inputs_embeds.device, + ) + if labels is not None: + final_labels = torch.full( + (batch_size, max_embed_dim), + ignore_index, + dtype=input_ids.dtype, + device=input_ids.device, + ) + # In case the Vision model or the Language model has been offloaded to CPU, we need to manually + # set the corresponding tensors into their correct target device. + target_device = inputs_embeds.device + batch_indices, non_image_indices, text_to_overwrite = ( + batch_indices.to(target_device), + non_image_indices.to(target_device), + text_to_overwrite.to(target_device), + ) + attention_mask = attention_mask.to(target_device) + + # 4. Fill the embeddings based on the mask. + final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[ + batch_indices, non_image_indices + ] + final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[ + batch_indices, non_image_indices + ] + if labels is not None: + final_labels[batch_indices, text_to_overwrite] = labels[ + batch_indices, non_image_indices + ] + + # 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835) + image_to_overwrite = torch.full( + (batch_size, max_embed_dim), + True, + dtype=torch.bool, + device=inputs_embeds.device, + ) + image_to_overwrite[batch_indices, text_to_overwrite] = False + image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[ + :, None + ].to(target_device) + + if image_to_overwrite.sum() != image_features.shape[:-1].numel(): + raise ValueError( + f"The input provided to the model are wrong. The number of image tokens is {image_to_overwrite.sum()} while" + f" the number of image features given to the model is {image_features.shape[:-1].numel()}. " + "This prevents correct indexing and breaks batch generation." + ) + + final_embedding[image_to_overwrite] = ( + image_features.contiguous().reshape(-1, embed_dim).to(target_device) + ) + final_attention_mask |= image_to_overwrite + position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_( + (final_attention_mask == 0), 1 + ) + + # 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens. + batch_indices, pad_indices = torch.where(input_ids == pad_token_id) + indices_to_mask = new_token_positions[batch_indices, pad_indices] + + final_embedding[batch_indices, indices_to_mask] = 0 + + if labels is None: + final_labels = None + + return final_embedding, final_attention_mask, final_labels, position_ids + + def _extract_image_features( + self, pixel_values: torch.Tensor, grid_thws: torch.Tensor + ) -> list[torch.Tensor]: + """ + Args: + pixel_values (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_channels, height, width)`): + The pixel values of the images processed by image processor. + grid_thws (:obj:`torch.Tensor` of shape :obj:`(batch_size, 3)`): + The grid, height, width of the images. + + Returns: + selected_image_feature (:obj:`torch.FloatTensor` of shape :obj:`(num_image_tokens, embed_dim)`): + The selected image features to use as input to the projector head. + + """ + + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = pixel_values.to(target_dtype) + + image_features = self.vision_tower(pixel_values, grid_thws) + return image_features + + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | list[torch.FloatTensor] | None = None, + grid_thws: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> tuple | LlavaCausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + ```""" + assert self.vision_tower is not None, "vision_tower is not loaded" + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if inputs_embeds is None: + # 1. Extra the input embeddings + inputs_embeds = self.get_input_embeddings()(input_ids) + + # 2. Merge text and images + if ( + pixel_values is not None + and len(pixel_values) > 0 + and input_ids.shape[1] != 1 + ): + image_features = self._extract_image_features(pixel_values, grid_thws) + if self.mm_projector: + image_features = self.mm_projector(image_features) + + inputs_embeds = inputs_embeds.to( + image_features[0].dtype + ) # num_tokens, embed_dim + inputs_embeds, attention_mask, labels, position_ids = ( + self._merge_input_ids_with_image_features( + image_features, + inputs_embeds, + input_ids, + attention_mask, + labels, + ) + ) + + # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of + # generation with cache + elif ( + past_key_values is not None + and pixel_values is not None + and input_ids.shape[1] == 1 + ): + # Retrieve the first layer to inspect the logits and mask out the hidden states + # that are set to 0 + first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] + + # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941 + batch_index, non_attended_tokens = torch.where( + first_layer_past_key_value.float().sum(-2) == 0 + ) + + # Get the target length + target_length = input_ids.shape[1] + past_length = first_layer_past_key_value.shape[-1] + + extended_attention_mask = torch.ones( + (attention_mask.shape[0], past_length), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Filter out only the tokens that can be un-attended, this can happen + # if one uses Llava + Fused modules where the cache on the + # first iteration is already big enough, or if one passes custom cache + valid_indices = non_attended_tokens < extended_attention_mask.size(-1) + new_batch_index = batch_index[valid_indices] + new_non_attended_tokens = non_attended_tokens[valid_indices] + + # Zero-out the places where we don't need to attend + extended_attention_mask[new_batch_index, new_non_attended_tokens] = 0 + + attention_mask = torch.cat( + (extended_attention_mask, attention_mask[:, -target_length:]), dim=1 + ) + position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1 + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + logits = outputs[0] + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + if attention_mask is not None: + shift_attention_mask = attention_mask[..., 1:] + shift_logits = logits[..., :-1, :][ + shift_attention_mask.to(logits.device) != 0 + ].contiguous() + shift_labels = labels[..., 1:][ + shift_attention_mask.to(labels.device) != 0 + ].contiguous() + else: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1).to(shift_logits.device), + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return LlavaCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + grid_thws=None, + attention_mask=None, + **kwargs, + ): + if past_key_values is not None: + if hasattr(past_key_values, "get_seq_length"): + cache_length = past_key_values.get_seq_length() + past_length = getattr(past_key_values, "seen_tokens", cache_length) + else: + cache_length = past_length = past_key_values[0][0].shape[2] + + # Keep only the unprocessed tokens: + # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where + # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as + # input) + if ( + attention_mask is not None + and attention_mask.shape[1] > input_ids.shape[1] + ): + input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] + # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard + # input_ids based on the past_length. + elif past_length < input_ids.shape[1]: + input_ids = input_ids[:, past_length:] + # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens. + elif self.config.media_placeholder_token_id in input_ids: + input_ids = input_ids[:, input_ids.shape[1] - 1 :] + # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the + # older attention values, as their corresponding values are not part of the input. + if cache_length < past_length and attention_mask is not None: + attention_mask = attention_mask[ + :, -(cache_length + input_ids.shape[1]) : + ] + + position_ids = kwargs.get("position_ids", None) + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "position_ids": position_ids, + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "grid_thws": grid_thws, + } + ) + return model_inputs + + def _reorder_cache(self, *args, **kwargs): + return self.language_model._reorder_cache(*args, **kwargs) diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py new file mode 100644 index 0000000000..1873391e71 --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py @@ -0,0 +1,1479 @@ +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. +# +# The multi-head latent attention, MoE gating and sparse MoE block in this file are +# adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py). They have been +# extensively modified and extended for the Kimi-Linear architecture. +# +# Licensing Information: +# - Code adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 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. +import math +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn.functional as F +import transformers +from einops import rearrange +from packaging import version +from torch import nn +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.masking_utils import create_causal_mask +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS +from transformers.utils import ( + TransformersKwargs, + can_return_tuple, + logging, +) +from transformers.utils.generic import check_model_inputs +from transformers.utils.output_capturing import OutputRecorder + +try: + from fla.modules import FusedRMSNormGated, ShortConvolution + from fla.ops.kda import chunk_kda, fused_recurrent_kda + + # from fla.ops.kda.gate import fused_kda_gate # deprecated, gate is now computed inside chunk_kda/fused_recurrent_kda + from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask + from fla.utils import tensor_cache +except ImportError: + raise ImportError("Plese run `pip install -U fla-core`") + +from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag + +from .configuration_kimi_k3 import KimiLinearConfig + +assert version.parse(transformers.__version__) >= version.parse( + "4.56.0" +), "Please upgrade transformers to >= 4.56.0" + +logger = logging.get_logger(__name__) + + +# Register Moonshot-specific activation functions +class SituAndMul(nn.Module): + """ + SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up + When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta). + """ + + def __init__(self, beta: float = 1.0, linear_beta: float | None = None): + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].to(torch.float32) + up = x[..., d:].to(torch.float32) + situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (situ_a * up).to(x.dtype) + + +ACT2FN["situ"] = SituAndMul + + +def _get_situ_activation_params(config: KimiLinearConfig): + beta = getattr(config, "activation_situ_beta", None) + linear_beta = getattr(config, "activation_situ_linear_beta", None) + return beta or 1.0, linear_beta + + +def index_first_axis(x, indices): + return x[indices] + + +@tensor_cache +def get_unpad_data( + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, int]: + lens = prepare_lens_from_mask(attention_mask) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = lens.max().item() + cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) + return indices, cu_seqlens, max_seqlen_in_batch + + +def pad_input( + hidden_states: torch.Tensor, + indices: torch.LongTensor, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + out = hidden_states.new_zeros((batch_size * seq_len, *hidden_states.shape[1:])) + out[indices] = hidden_states + return out.view(batch_size, seq_len, *hidden_states.shape[1:]) + + +class KimiDynamicCache: + """ + Dynamic cache for Kimi model. + Inspired by Qwen3-Next + """ + + is_compileable = False + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + + if config.linear_attn_config is not None: + self.layer_types = [] + for i in range(config.num_hidden_layers): + if config.is_kda_layer(i): + self.layer_types.append("linear_attention") + else: + self.layer_types.append("full_attention") + else: + self.layer_types = ["full_attention"] * config.num_hidden_layers + + self.transformer_layers = [ + i + for i in range(config.num_hidden_layers) + if self.layer_types[i] == "full_attention" + ] + + linear_layers = [ + i + for i in range(config.num_hidden_layers) + if self.layer_types[i] == "linear_attention" + ] + self.last_linear_layer = linear_layers[-1] if linear_layers else -1 + + self.conv_states = [None for _ in range(config.num_hidden_layers)] + self.recurrent_states = [None for _ in range(config.num_hidden_layers)] + self.key_cache = [None for _ in range(config.num_hidden_layers)] + self.value_cache = [None for _ in range(config.num_hidden_layers)] + + def __len__(self): + return len(self.layer_types) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.key_cache[layer_idx] is None: + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat( + [self.key_cache[layer_idx], key_states], dim=2 + ) + self.value_cache[layer_idx] = torch.cat( + [self.value_cache[layer_idx], value_states], dim=2 + ) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx] is not None: + device = self.key_cache[layer_idx].device + beam_idx = beam_idx.to(device) + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select( + 0, beam_idx + ) + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select( + 0, beam_idx + ) + + if self.conv_states[layer_idx] is not None: + device = self.conv_states[layer_idx][0].device + beam_idx = beam_idx.to(device) + q_conv, k_conv, v_conv = self.conv_states[layer_idx] + self.conv_states[layer_idx] = ( + q_conv.index_select(0, beam_idx), + k_conv.index_select(0, beam_idx), + v_conv.index_select(0, beam_idx), + ) + self.recurrent_states[layer_idx] = self.recurrent_states[ + layer_idx + ].index_select(0, beam_idx) + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # take any layer that contains cache and not empty tensor + layer_idx = ( + self.transformer_layers[0] + if layer_idx not in self.transformer_layers + else layer_idx + ) + if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: + return 0 + return self.key_cache[layer_idx].shape[-2] + + def get_query_offset(self, layer_idx: int = 0) -> int: + return self.get_seq_length(layer_idx=layer_idx) + + def get_mask_sizes(self, cache_position, layer_idx: int) -> tuple[int, int]: + """ + Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for + the given layer at `layer_idx`. + The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. + """ + kv_offset = 0 + # cache_position may be an int (new API) or a 1-D tensor (old API) + query_length = ( + cache_position + if isinstance(cache_position, int) + else cache_position.shape[0] + ) + past_seen_tokens = self.get_seq_length(layer_idx) + kv_length = query_length + past_seen_tokens + return kv_length, kv_offset + + @property + def has_previous_state(self): + """We have a previous state if the last linear (conv) layer was already updated.""" + if self.last_linear_layer == -1: + return False + return self.conv_states[self.last_linear_layer] is not None + + +class KimiRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + dtype = hidden_states.dtype + x = hidden_states.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + return self.weight * x.to(dtype) + + +ALL_LAYERNORM_LAYERS.append(KimiRMSNorm) + + +class KimiBlockSparseMLP(nn.Module): + def __init__( + self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None + ): + super().__init__() + self.config = config + self.ffn_dim = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size + + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down + self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up + + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + if self.config.hidden_act == "situ": + gate_up = torch.cat( + [self.w1(hidden_states), self.w3(hidden_states)], dim=-1 + ) + current_hidden_states = self.act_fn(gate_up) + else: + current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3( + hidden_states + ) + current_hidden_states = self.w2(current_hidden_states) + return current_hidden_states + + +class KimiMLP(nn.Module): + def __init__( + self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + if self.config.hidden_act == "situ": + gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1) + down_proj = self.down_proj(self.act_fn(gate_up)) + else: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand the key/value heads from `num_key_value_heads` to `num_attention_heads`.""" + if n_rep == 1: + return hidden_states + return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling + if attention_mask is not None: + scores = scores + attention_mask[:, :, :, : key.shape[-2]] + + probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) + probs = F.dropout(probs, p=dropout, training=module.training) + out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous() + + return out, probs + + +class KimiMLAAttention(nn.Module): + """ + Multi-Latent Attention adapted from deepseek-v3 + """ + + def __init__(self, config: KimiLinearConfig, layer_idx: int): + nn.Module.__init__(self) + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + + self.attention_dropout = getattr(config, "attention_dropout", 0.0) + + try: + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.use_nope = config.mla_use_nope + self.scaling = self.q_head_dim ** (-0.5) + except Exception as e: + raise ValueError( + f"Kimi MLA config is not found or not properly formatted: {e}" + ) + + if self.q_lora_rank is not None: + self.q_a_proj = nn.Linear( + self.hidden_size, + self.q_lora_rank, + bias=False, + ) + self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank) + self.q_b_proj = nn.Linear( + self.q_lora_rank, + self.num_heads * self.q_head_dim, + bias=False, + ) + else: + self.q_proj = nn.Linear( + self.hidden_size, + self.num_heads * self.q_head_dim, + bias=False, + ) + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + ) + self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + ) + self.is_causal = True + assert self.use_nope + + self.use_output_gate = getattr(config, "mla_use_output_gate", False) + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.rotary_emb = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.q_head_dim) + key_shape = ( + batch_size, + seq_length, + -1, + self.qk_nope_head_dim + self.v_head_dim, + ) + + if self.q_lora_rank is not None: + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + else: + q_states = self.q_proj(hidden_states) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split( + q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + k_pass = ( + self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) + ) + k_pass, value_states = torch.split( + k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx + ) + + if ( + self.config._attn_implementation == "flash_attention_2" + and self.q_head_dim != self.v_head_dim + ): + value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[ + self.config._attn_implementation + ] + + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + if ( + self.config._attn_implementation == "flash_attention_2" + and self.q_head_dim != self.v_head_dim + ): + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + if self.use_output_gate: + g = self.g_proj(hidden_states).sigmoid() + attn_output = attn_output * g + attn_output = self.o_proj(attn_output) + return attn_output + + +class KimiDeltaAttention(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.config = config + self.mode = "chunk" + + self.hidden_size = config.hidden_size + self.conv_size = config.linear_attn_config["short_conv_kernel_size"] + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + self.head_k_dim = self.head_dim + self.num_k_heads = self.num_heads + + self.layer_idx = layer_idx + + assert self.mode in [ + "chunk", + "fused_recurrent", + ], f"Not supported mode `{self.mode}`." + + projection_k_size = self.head_k_dim * self.num_k_heads + projection_size = self.head_dim * self.num_heads + + self.q_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False) + self.k_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False) + self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.q_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=projection_size, + kernel_size=self.conv_size, + activation="silu", + ) + + self.A_log = torch.nn.Parameter( + torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)) + ) + + self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32)) + + self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False) + + self.use_full_rank_gate = config.linear_attn_config.get( + "use_full_rank_gate", False + ) + self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None) + if self.use_full_rank_gate: + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + else: + self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.o_norm = FusedRMSNormGated( + self.head_dim, eps=config.rms_norm_eps, activation="sigmoid" + ) + self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + cache_params: KimiDynamicCache | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + if attention_mask.dim() != 2: + attention_mask = kwargs.get("padding_mask") + + if attention_mask is not None and attention_mask.dim() != 2: + raise ValueError( + "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] " + "(0 = padding). 3D masks are not supported here.", + ) + use_cache = cache_params is not None + batch_size, q_len, _ = hidden_states.shape + mode = "fused_recurrent" if use_cache and q_len == 1 else self.mode + if self.training: + assert mode == "chunk", "Only chunk mode is supported in training." + + cu_seqlens = kwargs.get("cu_seqlens") + indices = None + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis( + rearrange(hidden_states, "b s ... -> (b s) ..."), indices + ).unsqueeze(0) + + conv_state_q, conv_state_k, conv_state_v = None, None, None + recurrent_state = None + if cache_params is not None: + if cache_params.conv_states[self.layer_idx] is not None: + conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[ + self.layer_idx + ] + recurrent_state = cache_params.recurrent_states[self.layer_idx] + + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + q, conv_state_q = self.q_conv1d( + x=q_proj_states, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=k_proj_states, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=v_proj_states, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + g = self.f_b_proj(self.f_a_proj(hidden_states)) + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) + beta = self.b_proj(hidden_states).float() + + q, k = map( + lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim), (q, k) + ) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + if mode == "chunk": + o, recurrent_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self.gate_lower_bound is not None, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + else: + o, recurrent_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + if cache_params is not None: + cache_params.recurrent_states[self.layer_idx] = recurrent_state + cache_params.conv_states[self.layer_idx] = ( + conv_state_q, + conv_state_k, + conv_state_v, + ) + + if self.use_full_rank_gate: + g = self.g_proj(hidden_states) + else: + g = self.g_b_proj(self.g_a_proj(hidden_states)) + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, g) + + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o + + +class KimiMoEGate(nn.Module): + """ + MoEGate adapted from Deepseek-V3. + Parameter correspondences: + num_experts -> n_routed_experts + num_experts_per_token -> num_experts_per_tok + num_expert_group -> n_group + moe_router_activation_func -> scoring_func + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + + # topk selection algorithm + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim)), + ) + + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + # compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear( + hidden_states.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + if self.moe_router_activation_func == "sigmoid": + scores = logits.sigmoid() + elif self.moe_router_activation_func == "softmax": + scores = logits.softmax(dim=1) + else: + raise NotImplementedError( + f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}", + ) + + # select top-k experts + scores = scores.view(bsz * seq_len, -1) + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: + group_scores = ( + scores_for_choice.view(bsz * seq_len, self.num_expert_group, -1) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) # [n, num_expert_group] + group_idx = torch.topk( + group_scores, + k=self.topk_group, + dim=-1, + sorted=False, + )[1] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, num_expert_group] + group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + bsz * seq_len, + self.num_expert_group, + self.num_experts // self.num_expert_group, + ) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill( + ~score_mask.bool(), float("-inf") + ) # [n, e] + else: + tmp_scores = scores_for_choice + _, topk_idx = torch.topk( + tmp_scores, + k=self.top_k, + dim=-1, + sorted=False, + ) + topk_weight = scores.gather(1, topk_idx) + + # norm gate to sum 1 + if self.top_k > 1 and self.moe_renormalize: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + # must multiply the scaling factor + topk_weight = topk_weight * self.routed_scaling_factor + + return topk_idx, topk_weight + + +class KimiSparseMoeBlock(nn.Module): + """ + Adapted from Deepseek-V3's MOE implementation + The namings are consistent with Kimi's version. + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.hidden_dim = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_token + self.moe_renormalize = config.moe_renormalize + + self.use_latent_moe = ( + getattr(config, "routed_expert_hidden_size", None) is not None + ) + self.moe_hidden_size = ( + config.routed_expert_hidden_size + if self.use_latent_moe + else config.hidden_size + ) + self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False) + + self.ep_size = 1 + self.experts_per_rank = config.num_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + KimiBlockSparseMLP( + config, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + ) + for _ in range(config.num_experts) + ], + ) + self.gate = KimiMoEGate(config) + if config.num_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = KimiMLP( + config=config, + intermediate_size=intermediate_size, + ) + + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear( + config.hidden_size, + self.moe_hidden_size, + bias=False, + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, + config.hidden_size, + bias=False, + ) + if self.latent_moe_use_norm: + self.routed_expert_norm = KimiRMSNorm( + self.moe_hidden_size, + eps=config.rms_norm_eps, + ) + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + if self.use_latent_moe: + hidden_states = self.routed_expert_down_proj(hidden_states) + + if False:#if not self.training: + y = self.moe_infer(hidden_states, topk_idx, topk_weight) + else: + y = self.moe_train(hidden_states, topk_idx, topk_weight) + + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + + y = y.view(*orig_shape) + + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + def moe_train(self, x, topk_ids, topk_weight): + """Training-compatible MoE dispatch with gradient flow.""" + y = torch.zeros_like(x) + + with torch.no_grad(): + expert_mask = F.one_hot(topk_ids, self.num_experts).permute(2, 1, 0) + + for expert_idx, expert in enumerate(self.experts): + top_k_pos, token_indices = torch.where(expert_mask[expert_idx]) + + if get_calibrate_all_experts_flag(): + expert_out = expert(x)[token_indices] + else: + expert_out = expert(x[token_indices]) + + expert_weights = topk_weight[token_indices, top_k_pos, None] + y.index_add_(0, token_indices, (expert_out * expert_weights).to(y.dtype)) + + return y + + @torch.no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +class KimiDecoderLayer(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + self.layer_idx = layer_idx + if config.is_kda_layer(layer_idx): + self.is_linear_attn = True + self.self_attn = KimiDeltaAttention(config=config, layer_idx=layer_idx) + elif config.is_mla: + self.is_linear_attn = False + self.self_attn = KimiMLAAttention(config=config, layer_idx=layer_idx) + else: + raise NotImplementedError + if ( + config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % getattr(config, "moe_layer_freq", 1) == 0 + ): + self.block_sparse_moe = KimiSparseMoeBlock(config) + else: + self.mlp = KimiMLP(config) + self.input_layernorm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + # Attention residual + self.use_attn_residuals = ( + getattr(config, "attn_res_block_size", None) is not None + ) + if self.use_attn_residuals: + self.attn_res_block_size = config.attn_res_block_size + self.self_attention_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = nn.Linear(config.hidden_size, 1, bias=False) + self.mlp_res_proj = nn.Linear(config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + if self.use_attn_residuals: + return self._forward_attn_residual( + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + block_residual, + **kwargs, + ) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + def _forward_attn_residual( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + batch_size, seq_len, hidden_size = hidden_states.shape + prefix_sum = hidden_states + is_block_start = self.layer_idx % self.attn_res_block_size == 0 + + hidden_states, block_residual = _attn_res_pre_update( + block_residual, + prefix_sum, + hidden_size, + batch_size, + seq_len, + self.self_attention_res_proj, + self.self_attention_res_norm, + is_block_start, + ) + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + if is_block_start: + prefix_sum = hidden_states + else: + prefix_sum = prefix_sum + hidden_states + + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + ).view(batch_size, seq_len, hidden_size) + + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + + prefix_sum = prefix_sum + hidden_states + + return prefix_sum, block_residual + + +class KimiPreTrainedModel(PreTrainedModel): + config_class = KimiLinearConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["KimiDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _can_record_outputs = { + "router_logits": OutputRecorder(KimiBlockSparseMLP, index=1), + "hidden_states": KimiDecoderLayer, + "attentions": KimiMLAAttention, + } + _is_stateful = True + + def _init_weights(self, module): + # HOTFIX: disk offloading attempts to initialize the meta tensors + # but this is bad programming: we shouldn't be initializing these + # params in the first place + # the init attempt attempts to get `module.weight`, which DNE for qmodels + + return + + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +def _apply_attn_res(prefix_sum, block_residual, proj, norm): + """ + prefix_sum: (num_tokens, hidden_size) + block_residual: (num_tokens, num_blocks, hidden_size) + """ + v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.variance_epsilon) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +@torch.fx.wrap +def _attn_res_pre_update( + block_residual, + prefix_sum, + hidden_size, + batch_size, + seq_len, + self_attention_res_proj, + self_attention_res_norm, + is_block_start, +): + hidden_states = prefix_sum + if block_residual is not None and block_residual.shape[1] > 0: + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self_attention_res_proj, + self_attention_res_norm, + ).view(batch_size, seq_len, hidden_size) + if is_block_start: + block_residual = torch.cat( + [block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1 + ) + return hidden_states, block_residual + + +class KimiLinearModel(KimiPreTrainedModel): + def __init__(self, config: KimiLinearConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + KimiDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.use_attn_residuals = ( + getattr(config, "attn_res_block_size", None) is not None + ) + if self.use_attn_residuals: + self.output_attn_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False) + + from transformers.utils import is_flash_attn_2_available as _fa2_avail + + _requested = getattr(config, "_attn_implementation", None) + if _requested not in (None, "flash_attention_2") or not _fa2_avail(): + # Fall back gracefully when flash-attn2 is unavailable or a different impl is requested + if _requested == "flash_attention_2" and not _fa2_avail(): + logger.warning_once( + "flash_attention_2 requested but not available; falling back to sdpa." + ) + config._attn_implementation = ( + _requested if _requested not in (None, "flash_attention_2") else "eager" + ) + else: + config._attn_implementation = "flash_attention_2" + + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def _update_linear_attn_mask(self, attention_mask, cache_position): + """ + NOTE: Left-padding is used for linear attention mask. + No need for zeroing states when + 1. Cached forward + 2. Attending to all inputs + """ + linear_attn_mask = attention_mask + if cache_position[0] > 0 or ( + attention_mask is not None and torch.all(attention_mask == 1) + ): + linear_attn_mask = None + return linear_attn_mask + + @check_model_inputs + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) and (inputs_embeds is None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + # Get inputs_embeds + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = KimiDynamicCache(config=self.config) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position: torch.Tensor = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position) + + hidden_states = inputs_embeds + if past_key_values is not None: + assert isinstance(past_key_values, KimiDynamicCache) + + block_residual = None + if self.use_attn_residuals: + block_residual = hidden_states.new_zeros( + hidden_states.shape[0] * hidden_states.shape[1], + 0, + hidden_states.shape[2], + ) + + for decoder_layer in self.layers: + layer_mask = ( + linear_attn_mask if decoder_layer.is_linear_attn else causal_mask + ) + + if self.use_attn_residuals: + hidden_states, block_residual = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + block_residual=block_residual, + **kwargs, + ) + else: + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + if self.use_attn_residuals: + hidden_states = self._apply_output_attn_res(hidden_states, block_residual) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + def _apply_output_attn_res(self, hidden_states, block_residual): + batch_size, seq_len, hidden_size = hidden_states.shape + return _apply_attn_res( + hidden_states.view(-1, hidden_size), + block_residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + ).view(batch_size, seq_len, hidden_size) + + +class KimiLinearForCausalLM(KimiPreTrainedModel, GenerationMixin): + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = KimiLinearModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + generation_mode: bool | None = None, + return_dict: bool | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + """ + + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + logits = outputs[0] + if generation_mode: + logits = logits[:, -1:] + logits = self.lm_head(logits) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py new file mode 100644 index 0000000000..a6be474412 --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py @@ -0,0 +1 @@ +from .modeling_kimi_k3_linear import * diff --git a/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py b/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py new file mode 100644 index 0000000000..a9d00030cf --- /dev/null +++ b/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py @@ -0,0 +1,430 @@ +import os +from logging import getLogger +from pathlib import Path +from shutil import copyfile +from typing import Dict, Iterator, List, Optional, Tuple, Union, cast + +import tiktoken +from tiktoken.load import load_tiktoken_bpe +from tokenizers import AddedToken +from transformers.convert_slow_tokenizer import bytes_to_unicode +from transformers.tokenization_utils import PreTrainedTokenizer + +try: + from .encoding_k3 import build_chat_segments, is_batched_conversation +except ImportError: # pragma: no cover - supports direct file execution/import. + from encoding_k3 import build_chat_segments, is_batched_conversation + +logger = getLogger(__name__) +VOCAB_FILES_NAMES = {"vocab_file": "tiktoken.model"} + + +class TikTokenTokenizer(PreTrainedTokenizer): + """ + Tokenizing and encoding/decoding text using the Tiktoken tokenizer. See megatron/tokenizer/tiktoken_tokenizer.py. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + The path to the Tiktoken model file. + bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|begin_of_text|>",`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|end_of_text|>"`): + The end of sequence token. + unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_249|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. The second to last item in special_tokens. + pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<|reserved_special_token_250|>"`): + The token used for padding, for example when batching sequences of different lengths. + additional_special_tokens (list of `str`, *optional*): + A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be + skipped when decoding if `skip_special_tokens` is set to `True`. + """ + + vocab_files_names = VOCAB_FILES_NAMES + + model_input_names = ["input_ids", "attention_mask"] + + special_tokens: Dict[str, int] + + num_reserved_special_tokens = 256 + + pat_str = "|".join( + [ + r"""[\p{Han}]+""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""", + r"""\p{N}{1,3}""", + r""" ?[^\s\p{L}\p{N}]+[\r\n]*""", + r"""\s*[\r\n]+""", + r"""\s+(?!\S)""", + r"""\s+""", + ] + ) + + def __init__( + self, + vocab_file, + bos_token: Union[str, AddedToken] = "[BOS]", + eos_token: Union[str, AddedToken] = "[EOS]", + unk_token: Union[str, AddedToken, None] = None, + pad_token: Union[str, AddedToken, None] = None, + additional_special_tokens: List[str] = None, + added_tokens_decoder: Optional[dict] = None, + **kwargs, + ): + assert os.path.isfile(vocab_file), vocab_file + + if additional_special_tokens is None: + additional_special_tokens = [ + "<|im_end|>", + "<|im_user|>", + "<|im_assistant|>", + "<|start_header_id|>", + "<|end_header_id|>", + "[EOT]", + "<|im_system|>", + "<|im_middle|>", + ] + + if added_tokens_decoder: + special_tokens_mapping = { + i: added_tokens_decoder[i].content for i in added_tokens_decoder + } + else: + special_tokens_mapping = {} + + self.vocab_file = vocab_file + mergeable_ranks = load_tiktoken_bpe(vocab_file) + num_base_tokens = len(mergeable_ranks) + self.special_tokens = { + special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i + for i in range( + num_base_tokens, num_base_tokens + self.num_reserved_special_tokens + ) + } + + self.model = tiktoken.Encoding( + name=Path(vocab_file).name, + pat_str=self.pat_str, + mergeable_ranks=mergeable_ranks, + special_tokens=self.special_tokens, + ) + logger.info(f"Reloaded tiktoken model from {vocab_file}") + + self.n_words: int = self.model.n_vocab + # BOS / EOS token IDs + self.bos_id: int = self.special_tokens[str(bos_token)] + self.eos_id: int = self.special_tokens[str(eos_token)] + logger.info( + f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}" + ) + + self.pad_id: int = self.special_tokens[str(pad_token)] + self.unk_id: int = self.special_tokens[str(unk_token)] + + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + + self.decoder = {} + for i in range(self.n_words): + # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee + decoding = "".join( + [ + self.byte_encoder[ord(char)] + for char in self.model.decode_single_token_bytes(i).decode( + "latin-1" + ) + ] + ) + self.decoder[i] = decoding + + self.encoder = {} + for i in range(self.n_words): + if i in self.decoder: + self.encoder[self.decoder[i]] = i + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + additional_special_tokens=additional_special_tokens, + added_tokens_decoder=added_tokens_decoder, + **kwargs, + ) + self.all_special_ids_set = set(self.all_special_ids) + + def _encode_text_piece( + self, text: str, allow_special_tokens: bool = True + ) -> List[int]: + # The tiktoken tokenizer can handle <=400k chars without + # pyo3_runtime.PanicException. + TIKTOKEN_MAX_ENCODE_CHARS = 400_000 + + # https://github.com/openai/tiktoken/issues/195 + # Here we iterate over subsequences and split if we exceed the limit + # of max consecutive non-whitespace or whitespace characters. + MAX_NO_WHITESPACES_CHARS = 25_000 + + t: List[int] = [] + for i in range(0, len(text), TIKTOKEN_MAX_ENCODE_CHARS): + for substr in self._split_whitespaces_or_nonwhitespaces( + text[i : i + TIKTOKEN_MAX_ENCODE_CHARS], + MAX_NO_WHITESPACES_CHARS, + ): + if allow_special_tokens: + t.extend( + # structural markers: encode <|...|> as their special token IDs + self.model.encode( + substr, + allowed_special="all", + ) + ) + else: + t.extend( + # user/tool text: encode any <|...|> as ordinary BPE tokens (never as control tokens) + self.model.encode( + substr, + disallowed_special=(), + ) + ) + + return t + + def encode( + self, text: str, allow_special_tokens: bool = True, **kwargs + ) -> List[int]: + """ + Encodes a string into a list of token IDs. + + Args: + text (str): The input string to be encoded. + + Returns: + list[int]: A list of token IDs. + """ + # If there are other args, we should call super().encode because there are a lot of code + # to handle those args. supper().encode finally will call _tokenize and _convert_token_to_id. + # NOTE: our encode method is not compatible with the super().encode method, + # e.g. split_special_tokens' default is True in our encode method. + if len(kwargs) > 0: + logger.warning(f"Calling super().encode with {kwargs}") + return super().encode(text, **kwargs) + + assert type(text) is str + return self._encode_text_piece(text, allow_special_tokens=allow_special_tokens) + + def decode(self, token_ids: Union[int, List[int]], **kwargs) -> str: + """ + Decodes a list of token IDs into a string. + + Args: + token_ids (List[int]): The list of token IDs to be decoded. + + Returns: + str: The decoded string. + """ + # If there are other args, we should call super().decode because there are a lot of code + # to handle those args. supper().encode finally will call convert_tokens_to_string and _convert_id_to_token. + if len(kwargs) > 0: + return super().decode(token_ids, **kwargs) + + if type(token_ids) is int: + token_ids = [token_ids] + + return self.model.decode(cast(List[int], token_ids)) + + @staticmethod + def _split_whitespaces_or_nonwhitespaces( + s: str, max_consecutive_slice_len: int + ) -> Iterator[str]: + """ + Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len` + consecutive whitespaces or consecutive non-whitespaces. + """ + current_slice_len = 0 + current_slice_is_space = s[0].isspace() if len(s) > 0 else False + slice_start = 0 + + for i in range(len(s)): + is_now_space = s[i].isspace() + + if current_slice_is_space ^ is_now_space: + current_slice_len = 1 + current_slice_is_space = is_now_space + else: + current_slice_len += 1 + if current_slice_len > max_consecutive_slice_len: + yield s[slice_start:i] + slice_start = i + current_slice_len = 1 + yield s[slice_start:] + + def _encode_chat_segments(self, segments) -> List[int]: + token_ids: List[int] = [] + for segment in segments: + token_ids.extend( + self._encode_text_piece( + segment.text, + allow_special_tokens=segment.allow_special, + ) + ) + return token_ids + + @staticmethod + def _truncate( + ids: List[int], truncation: bool = False, max_length: Optional[int] = None + ) -> List[int]: + if truncation and max_length is not None: + return ids[:max_length] + return ids + + def _format_chat_token_output( + self, + encoded_inputs: List[List[int]], + *, + is_batched: bool, + padding=False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors=None, + return_dict: bool = False, + ): + encoded_inputs = [ + self._truncate(ids, truncation=truncation, max_length=max_length) + for ids in encoded_inputs + ] + + needs_batch_encoding = ( + is_batched or padding or return_tensors is not None or return_dict + ) + if not needs_batch_encoding: + return encoded_inputs[0] + + features = [ + {"input_ids": ids, "attention_mask": [1] * len(ids)} + for ids in encoded_inputs + ] + batch = self.pad( + features, + padding=padding, + max_length=max_length if padding else None, + return_attention_mask=True, + return_tensors=return_tensors, + ) + + if return_dict: + return batch + if is_batched: + return batch["input_ids"] + return batch["input_ids"][0] if return_tensors is None else batch["input_ids"] + + """ ----- Below are the abstract methods required by PreTrainedTokenizer ----- """ + + @property + def vocab_size(self) -> int: + return self.n_words + + def get_vocab(self) -> Dict[str, int]: + return self.encoder + + def _tokenize(self, text: str, **kwargs) -> List[str]: + return [self.decoder[t] for t in self.encode(text)] + + def _convert_token_to_id(self, token: str) -> int: + return self.encoder.get(token, self.unk_id) + + def _convert_id_to_token(self, index: int) -> str: + return self.decoder.get(index) + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + return out_string + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + text = "".join(tokens) + text = bytearray([self.byte_decoder[c] for c in text]).decode( + "utf-8", "replace" + ) + return text + + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + if not os.path.isdir(save_directory): + raise ValueError( + f"vocabulary path ({save_directory}) should be a directory" + ) + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath( + out_vocab_file + ) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + + return (out_vocab_file,) + + def apply_chat_template( + self, + conversation, + tools: Optional[list[dict]] = None, + tokenize: bool = False, + add_generation_prompt: bool = True, + thinking: bool = True, + padding=False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors=None, + return_dict: bool = False, + **kwargs, + ): + # Tokenizer-level rendering reorders tool result messages to match + # assistant tool_calls, normalizes per-call arguments and response + # schema, then encodes the resulting XTML structure segment-by-segment. + is_batched = is_batched_conversation(conversation) + conversations = conversation if is_batched else [conversation] + image_prompts = kwargs.pop("image_prompts", None) + if is_batched and image_prompts is not None: + raise ValueError("image_prompts is only supported for one chat.") + + # by default set thinking effort to max + kwargs.setdefault("thinking_effort", "max") + + segment_batches = [ + build_chat_segments( + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + thinking=thinking, + image_prompts=image_prompts, + **kwargs, + ) + for messages in conversations + ] + + if not tokenize: + rendered = [ + "".join(segment.text for segment in segments) + for segments in segment_batches + ] + return rendered if is_batched else rendered[0] + + encoded_inputs = [ + self._encode_chat_segments(segments) for segments in segment_batches + ] + return self._format_chat_token_output( + encoded_inputs, + is_batched=is_batched, + padding=padding, + truncation=truncation, + max_length=max_length, + return_tensors=return_tensors, + return_dict=return_dict, + ) From 696791a1f973628b79c49d1ea6ef4d0dfcecc95c Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Mon, 10 Aug 2026 15:45:10 -0400 Subject: [PATCH 2/6] Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Kyle Sayers --- src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py index 1873391e71..d64f42ecc9 100644 --- a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py @@ -908,7 +908,7 @@ def forward(self, hidden_states): if self.use_latent_moe: hidden_states = self.routed_expert_down_proj(hidden_states) - + if not self.training: if False:#if not self.training: y = self.moe_infer(hidden_states, topk_idx, topk_weight) else: From d7f930f668fc246427f5d70c661de24fcef11877 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Mon, 10 Aug 2026 15:50:14 -0400 Subject: [PATCH 3/6] docs: add comments noting differences from upstream Kimi-K3 modeling files Each kimi_k3 file now has a comment at the top documenting implementation differences from the original files at moonshotai/Kimi-K3 on Hugging Face. Co-Authored-By: Claude Opus 4.6 --- .../modeling/kimi_k3/configuration_kimi_k3.py | 3 +++ .../modeling/kimi_k3/encoding_k3.py | 3 +++ .../modeling/kimi_k3/kimi_k3_processor.py | 3 +++ .../kimi_k3/kimi_k3_vision_processing.py | 3 +++ .../modeling/kimi_k3/media_utils.py | 3 +++ .../modeling/kimi_k3/modeling_kimi_k3.py | 5 ++++ .../kimi_k3/modeling_kimi_k3_linear.py | 24 +++++++++++++++++++ .../modeling/kimi_k3/tokenization_kimi.py | 3 +++ 8 files changed, 47 insertions(+) diff --git a/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py b/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py index efca3a0c23..40703fe090 100644 --- a/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py +++ b/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py @@ -1,3 +1,6 @@ +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - Formatting only (no functional changes) + from typing import Optional from transformers.configuration_utils import PretrainedConfig diff --git a/src/llmcompressor/modeling/kimi_k3/encoding_k3.py b/src/llmcompressor/modeling/kimi_k3/encoding_k3.py index 2a8f1271ed..e3306fc242 100644 --- a/src/llmcompressor/modeling/kimi_k3/encoding_k3.py +++ b/src/llmcompressor/modeling/kimi_k3/encoding_k3.py @@ -1,3 +1,6 @@ +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - Formatting only (no functional changes) + """Kimi K3 XTML encoding helpers. This module keeps chat rendering in Python. diff --git a/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py b/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py index 3248c69aba..14aaee41f4 100644 --- a/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py +++ b/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py @@ -1,3 +1,6 @@ +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - Formatting only (no functional changes) + """Kimi-K3 processor: wraps vision processor + tokenizer into a single interface. Chat rendering (including XTML tool-result ordering) is handled by the diff --git a/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py b/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py index 705c3f167a..0cafb35045 100644 --- a/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py +++ b/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py @@ -1,3 +1,6 @@ +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - Formatting only (no functional changes) + """Image processor class for Kimi-K3.""" import json diff --git a/src/llmcompressor/modeling/kimi_k3/media_utils.py b/src/llmcompressor/modeling/kimi_k3/media_utils.py index 560397b0c8..c6f36a1c9c 100644 --- a/src/llmcompressor/modeling/kimi_k3/media_utils.py +++ b/src/llmcompressor/modeling/kimi_k3/media_utils.py @@ -1,3 +1,6 @@ +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - Formatting only (no functional changes) + import base64 import functools import io diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py index 00dc8e042c..45556d0021 100644 --- a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3.py @@ -20,6 +20,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - KimiK3PreTrainedModel._init_weights: early return added to prevent +# initialization of meta tensors during disk offloading (quantized models) +# - KimiK3ForConditionalGeneration.tie_weights: accepts **kwargs for +# compatibility with transformers API # NOTE: Reference implementation for model architecture; see the model card for production deployment. import math diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py index d64f42ecc9..162774f036 100644 --- a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py @@ -21,6 +21,30 @@ # 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. + +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main +# (originally modeling_kimi_linear.py): +# - Renamed file from modeling_kimi_linear.py to modeling_kimi_k3_linear.py +# - KimiSparseMoeBlock: added moe_train() method with gradient flow for +# calibration; uses get_calibrate_all_experts_flag() to run all experts +# during calibration. Inference path disabled (if False) to force training +# path during quantization +# - KimiMoEGate.forward: removed `assert not self.training` to allow +# calibration in training mode +# - KimiPreTrainedModel._init_weights: early return to prevent initialization +# of meta tensors during disk offloading (quantized models) +# - KimiDecoderLayer._forward_attn_residual: refactored attention residual +# logic into _attn_res_pre_update() decorated with @torch.fx.wrap for +# torch.compile compatibility +# - KimiDynamicCache: added get_query_offset() method; get_mask_sizes() +# handles both int and tensor cache_position (new transformers API) +# - KimiLinearModel.__init__: attention implementation fallback changed from +# forcing flash_attention_2 to graceful fallback to sdpa/eager when FA2 +# is unavailable +# - create_causal_mask call: fixed input_embeds -> inputs_embeds typo, +# removed cache_position kwarg +# - KimiLinearForCausalLM._tied_weights_keys: changed from list to dict +# mapping format import math from collections.abc import Callable from typing import Any diff --git a/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py b/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py index a9d00030cf..5bca540870 100644 --- a/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py +++ b/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py @@ -1,3 +1,6 @@ +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - Formatting only (no functional changes) + import os from logging import getLogger from pathlib import Path From b6dd6de836d311942386af95d2d2c1c9b613ac04 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Mon, 10 Aug 2026 15:59:17 -0400 Subject: [PATCH 4/6] refactor: rename modeling_kimi_k3_linear back to modeling_kimi_linear, remove formatting-only comments - Rename modeling_kimi_k3_linear.py back to modeling_kimi_linear.py to match the upstream HuggingFace filename - Remove the re-export shim that was in modeling_kimi_linear.py - Remove "Differences: formatting only" comments from files with no substantive changes Co-Authored-By: Claude Opus 4.6 --- .../modeling/kimi_k3/configuration_kimi_k3.py | 3 - .../modeling/kimi_k3/encoding_k3.py | 3 - .../modeling/kimi_k3/kimi_k3_processor.py | 3 - .../kimi_k3/kimi_k3_vision_processing.py | 3 - .../modeling/kimi_k3/media_utils.py | 3 - .../kimi_k3/modeling_kimi_k3_linear.py | 1503 ----------------- .../modeling/kimi_k3/modeling_kimi_linear.py | 1502 +++++++++++++++- .../modeling/kimi_k3/tokenization_kimi.py | 3 - 8 files changed, 1501 insertions(+), 1522 deletions(-) delete mode 100644 src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py diff --git a/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py b/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py index 40703fe090..efca3a0c23 100644 --- a/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py +++ b/src/llmcompressor/modeling/kimi_k3/configuration_kimi_k3.py @@ -1,6 +1,3 @@ -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: -# - Formatting only (no functional changes) - from typing import Optional from transformers.configuration_utils import PretrainedConfig diff --git a/src/llmcompressor/modeling/kimi_k3/encoding_k3.py b/src/llmcompressor/modeling/kimi_k3/encoding_k3.py index e3306fc242..2a8f1271ed 100644 --- a/src/llmcompressor/modeling/kimi_k3/encoding_k3.py +++ b/src/llmcompressor/modeling/kimi_k3/encoding_k3.py @@ -1,6 +1,3 @@ -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: -# - Formatting only (no functional changes) - """Kimi K3 XTML encoding helpers. This module keeps chat rendering in Python. diff --git a/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py b/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py index 14aaee41f4..3248c69aba 100644 --- a/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py +++ b/src/llmcompressor/modeling/kimi_k3/kimi_k3_processor.py @@ -1,6 +1,3 @@ -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: -# - Formatting only (no functional changes) - """Kimi-K3 processor: wraps vision processor + tokenizer into a single interface. Chat rendering (including XTML tool-result ordering) is handled by the diff --git a/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py b/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py index 0cafb35045..705c3f167a 100644 --- a/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py +++ b/src/llmcompressor/modeling/kimi_k3/kimi_k3_vision_processing.py @@ -1,6 +1,3 @@ -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: -# - Formatting only (no functional changes) - """Image processor class for Kimi-K3.""" import json diff --git a/src/llmcompressor/modeling/kimi_k3/media_utils.py b/src/llmcompressor/modeling/kimi_k3/media_utils.py index c6f36a1c9c..560397b0c8 100644 --- a/src/llmcompressor/modeling/kimi_k3/media_utils.py +++ b/src/llmcompressor/modeling/kimi_k3/media_utils.py @@ -1,6 +1,3 @@ -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: -# - Formatting only (no functional changes) - import base64 import functools import io diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py deleted file mode 100644 index 162774f036..0000000000 --- a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_k3_linear.py +++ /dev/null @@ -1,1503 +0,0 @@ -# coding=utf-8 -# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. -# -# The multi-head latent attention, MoE gating and sparse MoE block in this file are -# adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py). They have been -# extensively modified and extended for the Kimi-Linear architecture. -# -# Licensing Information: -# - Code adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0. -# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). -# -# Apache License, Version 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. - -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main -# (originally modeling_kimi_linear.py): -# - Renamed file from modeling_kimi_linear.py to modeling_kimi_k3_linear.py -# - KimiSparseMoeBlock: added moe_train() method with gradient flow for -# calibration; uses get_calibrate_all_experts_flag() to run all experts -# during calibration. Inference path disabled (if False) to force training -# path during quantization -# - KimiMoEGate.forward: removed `assert not self.training` to allow -# calibration in training mode -# - KimiPreTrainedModel._init_weights: early return to prevent initialization -# of meta tensors during disk offloading (quantized models) -# - KimiDecoderLayer._forward_attn_residual: refactored attention residual -# logic into _attn_res_pre_update() decorated with @torch.fx.wrap for -# torch.compile compatibility -# - KimiDynamicCache: added get_query_offset() method; get_mask_sizes() -# handles both int and tensor cache_position (new transformers API) -# - KimiLinearModel.__init__: attention implementation fallback changed from -# forcing flash_attention_2 to graceful fallback to sdpa/eager when FA2 -# is unavailable -# - create_causal_mask call: fixed input_embeds -> inputs_embeds typo, -# removed cache_position kwarg -# - KimiLinearForCausalLM._tied_weights_keys: changed from list to dict -# mapping format -import math -from collections.abc import Callable -from typing import Any - -import torch -import torch.nn.functional as F -import transformers -from einops import rearrange -from packaging import version -from torch import nn -from transformers.activations import ACT2FN -from transformers.cache_utils import Cache -from transformers.generation import GenerationMixin -from transformers.masking_utils import create_causal_mask -from transformers.modeling_flash_attention_utils import FlashAttentionKwargs -from transformers.modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, -) -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel -from transformers.processing_utils import Unpack -from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS -from transformers.utils import ( - TransformersKwargs, - can_return_tuple, - logging, -) -from transformers.utils.generic import check_model_inputs -from transformers.utils.output_capturing import OutputRecorder - -try: - from fla.modules import FusedRMSNormGated, ShortConvolution - from fla.ops.kda import chunk_kda, fused_recurrent_kda - - # from fla.ops.kda.gate import fused_kda_gate # deprecated, gate is now computed inside chunk_kda/fused_recurrent_kda - from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask - from fla.utils import tensor_cache -except ImportError: - raise ImportError("Plese run `pip install -U fla-core`") - -from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag - -from .configuration_kimi_k3 import KimiLinearConfig - -assert version.parse(transformers.__version__) >= version.parse( - "4.56.0" -), "Please upgrade transformers to >= 4.56.0" - -logger = logging.get_logger(__name__) - - -# Register Moonshot-specific activation functions -class SituAndMul(nn.Module): - """ - SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up - When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta). - """ - - def __init__(self, beta: float = 1.0, linear_beta: float | None = None): - super().__init__() - self.beta = beta - self.linear_beta = linear_beta - - def forward(self, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - gate = x[..., :d].to(torch.float32) - up = x[..., d:].to(torch.float32) - situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) - if self.linear_beta is not None: - up = self.linear_beta * torch.tanh(up / self.linear_beta) - return (situ_a * up).to(x.dtype) - - -ACT2FN["situ"] = SituAndMul - - -def _get_situ_activation_params(config: KimiLinearConfig): - beta = getattr(config, "activation_situ_beta", None) - linear_beta = getattr(config, "activation_situ_linear_beta", None) - return beta or 1.0, linear_beta - - -def index_first_axis(x, indices): - return x[indices] - - -@tensor_cache -def get_unpad_data( - attention_mask: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, int]: - lens = prepare_lens_from_mask(attention_mask) - indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = lens.max().item() - cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) - return indices, cu_seqlens, max_seqlen_in_batch - - -def pad_input( - hidden_states: torch.Tensor, - indices: torch.LongTensor, - batch_size: int, - seq_len: int, -) -> torch.Tensor: - out = hidden_states.new_zeros((batch_size * seq_len, *hidden_states.shape[1:])) - out[indices] = hidden_states - return out.view(batch_size, seq_len, *hidden_states.shape[1:]) - - -class KimiDynamicCache: - """ - Dynamic cache for Kimi model. - Inspired by Qwen3-Next - """ - - is_compileable = False - - def __init__(self, config: KimiLinearConfig): - super().__init__() - self.config = config - - if config.linear_attn_config is not None: - self.layer_types = [] - for i in range(config.num_hidden_layers): - if config.is_kda_layer(i): - self.layer_types.append("linear_attention") - else: - self.layer_types.append("full_attention") - else: - self.layer_types = ["full_attention"] * config.num_hidden_layers - - self.transformer_layers = [ - i - for i in range(config.num_hidden_layers) - if self.layer_types[i] == "full_attention" - ] - - linear_layers = [ - i - for i in range(config.num_hidden_layers) - if self.layer_types[i] == "linear_attention" - ] - self.last_linear_layer = linear_layers[-1] if linear_layers else -1 - - self.conv_states = [None for _ in range(config.num_hidden_layers)] - self.recurrent_states = [None for _ in range(config.num_hidden_layers)] - self.key_cache = [None for _ in range(config.num_hidden_layers)] - self.value_cache = [None for _ in range(config.num_hidden_layers)] - - def __len__(self): - return len(self.layer_types) - - def update( - self, - key_states: torch.Tensor, - value_states: torch.Tensor, - layer_idx: int, - cache_kwargs: dict[str, Any] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.key_cache[layer_idx] is None: - self.key_cache[layer_idx] = key_states - self.value_cache[layer_idx] = value_states - else: - self.key_cache[layer_idx] = torch.cat( - [self.key_cache[layer_idx], key_states], dim=2 - ) - self.value_cache[layer_idx] = torch.cat( - [self.value_cache[layer_idx], value_states], dim=2 - ) - - return self.key_cache[layer_idx], self.value_cache[layer_idx] - - def reorder_cache(self, beam_idx: torch.LongTensor): - """Reorders the cache for beam search, given the selected beam indices.""" - for layer_idx in range(len(self.key_cache)): - if self.key_cache[layer_idx] is not None: - device = self.key_cache[layer_idx].device - beam_idx = beam_idx.to(device) - self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select( - 0, beam_idx - ) - self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select( - 0, beam_idx - ) - - if self.conv_states[layer_idx] is not None: - device = self.conv_states[layer_idx][0].device - beam_idx = beam_idx.to(device) - q_conv, k_conv, v_conv = self.conv_states[layer_idx] - self.conv_states[layer_idx] = ( - q_conv.index_select(0, beam_idx), - k_conv.index_select(0, beam_idx), - v_conv.index_select(0, beam_idx), - ) - self.recurrent_states[layer_idx] = self.recurrent_states[ - layer_idx - ].index_select(0, beam_idx) - - def get_seq_length(self, layer_idx: int | None = 0) -> int: - """Returns the sequence length of the cached states. A layer index can be optionally passed.""" - # take any layer that contains cache and not empty tensor - layer_idx = ( - self.transformer_layers[0] - if layer_idx not in self.transformer_layers - else layer_idx - ) - if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: - return 0 - return self.key_cache[layer_idx].shape[-2] - - def get_query_offset(self, layer_idx: int = 0) -> int: - return self.get_seq_length(layer_idx=layer_idx) - - def get_mask_sizes(self, cache_position, layer_idx: int) -> tuple[int, int]: - """ - Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for - the given layer at `layer_idx`. - The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. - """ - kv_offset = 0 - # cache_position may be an int (new API) or a 1-D tensor (old API) - query_length = ( - cache_position - if isinstance(cache_position, int) - else cache_position.shape[0] - ) - past_seen_tokens = self.get_seq_length(layer_idx) - kv_length = query_length + past_seen_tokens - return kv_length, kv_offset - - @property - def has_previous_state(self): - """We have a previous state if the last linear (conv) layer was already updated.""" - if self.last_linear_layer == -1: - return False - return self.conv_states[self.last_linear_layer] is not None - - -class KimiRMSNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - dtype = hidden_states.dtype - x = hidden_states.float() - x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) - return self.weight * x.to(dtype) - - -ALL_LAYERNORM_LAYERS.append(KimiRMSNorm) - - -class KimiBlockSparseMLP(nn.Module): - def __init__( - self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None - ): - super().__init__() - self.config = config - self.ffn_dim = ( - config.intermediate_size if intermediate_size is None else intermediate_size - ) - self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size - - self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate - self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down - self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up - - if config.hidden_act == "situ": - beta, linear_beta = _get_situ_activation_params(config) - self.act_fn = SituAndMul( - beta=beta, - linear_beta=linear_beta, - ) - else: - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, hidden_states): - if self.config.hidden_act == "situ": - gate_up = torch.cat( - [self.w1(hidden_states), self.w3(hidden_states)], dim=-1 - ) - current_hidden_states = self.act_fn(gate_up) - else: - current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3( - hidden_states - ) - current_hidden_states = self.w2(current_hidden_states) - return current_hidden_states - - -class KimiMLP(nn.Module): - def __init__( - self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None - ): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size if hidden_size is None else hidden_size - self.intermediate_size = ( - config.intermediate_size if intermediate_size is None else intermediate_size - ) - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - if config.hidden_act == "situ": - beta, linear_beta = _get_situ_activation_params(config) - self.act_fn = SituAndMul( - beta=beta, - linear_beta=linear_beta, - ) - else: - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, x): - if self.config.hidden_act == "situ": - gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1) - down_proj = self.down_proj(self.act_fn(gate_up)) - else: - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """Expand the key/value heads from `num_key_value_heads` to `num_attention_heads`.""" - if n_rep == 1: - return hidden_states - return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep) - - -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key = repeat_kv(key, module.num_key_value_groups) - value = repeat_kv(value, module.num_key_value_groups) - - scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling - if attention_mask is not None: - scores = scores + attention_mask[:, :, :, : key.shape[-2]] - - probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) - probs = F.dropout(probs, p=dropout, training=module.training) - out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous() - - return out, probs - - -class KimiMLAAttention(nn.Module): - """ - Multi-Latent Attention adapted from deepseek-v3 - """ - - def __init__(self, config: KimiLinearConfig, layer_idx: int): - nn.Module.__init__(self) - self.config = config - self.layer_idx = layer_idx - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.num_key_value_heads = config.num_key_value_heads - self.num_key_value_groups = self.num_heads // self.num_key_value_heads - - self.attention_dropout = getattr(config, "attention_dropout", 0.0) - - try: - self.q_lora_rank = config.q_lora_rank - self.qk_rope_head_dim = config.qk_rope_head_dim - self.kv_lora_rank = config.kv_lora_rank - self.v_head_dim = config.v_head_dim - self.qk_nope_head_dim = config.qk_nope_head_dim - self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - self.use_nope = config.mla_use_nope - self.scaling = self.q_head_dim ** (-0.5) - except Exception as e: - raise ValueError( - f"Kimi MLA config is not found or not properly formatted: {e}" - ) - - if self.q_lora_rank is not None: - self.q_a_proj = nn.Linear( - self.hidden_size, - self.q_lora_rank, - bias=False, - ) - self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank) - self.q_b_proj = nn.Linear( - self.q_lora_rank, - self.num_heads * self.q_head_dim, - bias=False, - ) - else: - self.q_proj = nn.Linear( - self.hidden_size, - self.num_heads * self.q_head_dim, - bias=False, - ) - self.kv_a_proj_with_mqa = nn.Linear( - self.hidden_size, - self.kv_lora_rank + self.qk_rope_head_dim, - bias=False, - ) - self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank) - self.kv_b_proj = nn.Linear( - self.kv_lora_rank, - self.num_heads - * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), - bias=False, - ) - self.o_proj = nn.Linear( - self.num_heads * self.v_head_dim, - self.hidden_size, - bias=False, - ) - self.is_causal = True - assert self.use_nope - - self.use_output_gate = getattr(config, "mla_use_output_gate", False) - if self.use_output_gate: - projection_size = self.num_heads * self.v_head_dim - self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) - - self.rotary_emb = None - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - batch_size, seq_length = hidden_states.shape[:-1] - query_shape = (batch_size, seq_length, -1, self.q_head_dim) - key_shape = ( - batch_size, - seq_length, - -1, - self.qk_nope_head_dim + self.v_head_dim, - ) - - if self.q_lora_rank is not None: - q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - else: - q_states = self.q_proj(hidden_states) - q_states = q_states.view(query_shape).transpose(1, 2) - q_pass, q_rot = torch.split( - q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) - - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - k_pass, k_rot = torch.split( - compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - - k_pass = ( - self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) - ) - k_pass, value_states = torch.split( - k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 - ) - - k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) - - k_rot = k_rot.expand(*k_pass.shape[:-1], -1) - - query_states = torch.cat((q_pass, q_rot), dim=-1) - key_states = torch.cat((k_pass, k_rot), dim=-1) - - if past_key_values is not None: - key_states, value_states = past_key_values.update( - key_states, value_states, self.layer_idx - ) - - if ( - self.config._attn_implementation == "flash_attention_2" - and self.q_head_dim != self.v_head_dim - ): - value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[ - self.config._attn_implementation - ] - - attn_output, _ = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) - - if ( - self.config._attn_implementation == "flash_attention_2" - and self.q_head_dim != self.v_head_dim - ): - attn_output = attn_output[:, :, :, : self.v_head_dim] - - attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() - if self.use_output_gate: - g = self.g_proj(hidden_states).sigmoid() - attn_output = attn_output * g - attn_output = self.o_proj(attn_output) - return attn_output - - -class KimiDeltaAttention(nn.Module): - def __init__(self, config: KimiLinearConfig, layer_idx: int): - super().__init__() - self.config = config - self.mode = "chunk" - - self.hidden_size = config.hidden_size - self.conv_size = config.linear_attn_config["short_conv_kernel_size"] - self.head_dim = config.linear_attn_config["head_dim"] - self.num_heads = config.linear_attn_config["num_heads"] - self.head_k_dim = self.head_dim - self.num_k_heads = self.num_heads - - self.layer_idx = layer_idx - - assert self.mode in [ - "chunk", - "fused_recurrent", - ], f"Not supported mode `{self.mode}`." - - projection_k_size = self.head_k_dim * self.num_k_heads - projection_size = self.head_dim * self.num_heads - - self.q_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False) - self.k_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False) - self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False) - - self.q_conv1d = ShortConvolution( - hidden_size=projection_k_size, - kernel_size=self.conv_size, - activation="silu", - ) - self.k_conv1d = ShortConvolution( - hidden_size=projection_k_size, - kernel_size=self.conv_size, - activation="silu", - ) - self.v_conv1d = ShortConvolution( - hidden_size=projection_size, - kernel_size=self.conv_size, - activation="silu", - ) - - self.A_log = torch.nn.Parameter( - torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)) - ) - - self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) - self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) - - self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32)) - - self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False) - - self.use_full_rank_gate = config.linear_attn_config.get( - "use_full_rank_gate", False - ) - self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None) - if self.use_full_rank_gate: - self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) - else: - self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) - self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) - - self.o_norm = FusedRMSNormGated( - self.head_dim, eps=config.rms_norm_eps, activation="sigmoid" - ) - self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - cache_params: KimiDynamicCache | None = None, - **kwargs: Unpack[dict], - ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: - if attention_mask is not None: - if attention_mask.dim() != 2: - attention_mask = kwargs.get("padding_mask") - - if attention_mask is not None and attention_mask.dim() != 2: - raise ValueError( - "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] " - "(0 = padding). 3D masks are not supported here.", - ) - use_cache = cache_params is not None - batch_size, q_len, _ = hidden_states.shape - mode = "fused_recurrent" if use_cache and q_len == 1 else self.mode - if self.training: - assert mode == "chunk", "Only chunk mode is supported in training." - - cu_seqlens = kwargs.get("cu_seqlens") - indices = None - if attention_mask is not None: - indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) - hidden_states = index_first_axis( - rearrange(hidden_states, "b s ... -> (b s) ..."), indices - ).unsqueeze(0) - - conv_state_q, conv_state_k, conv_state_v = None, None, None - recurrent_state = None - if cache_params is not None: - if cache_params.conv_states[self.layer_idx] is not None: - conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[ - self.layer_idx - ] - recurrent_state = cache_params.recurrent_states[self.layer_idx] - - q_proj_states = self.q_proj(hidden_states) - k_proj_states = self.k_proj(hidden_states) - v_proj_states = self.v_proj(hidden_states) - q, conv_state_q = self.q_conv1d( - x=q_proj_states, - cache=conv_state_q, - output_final_state=use_cache, - cu_seqlens=cu_seqlens, - ) - k, conv_state_k = self.k_conv1d( - x=k_proj_states, - cache=conv_state_k, - output_final_state=use_cache, - cu_seqlens=cu_seqlens, - ) - v, conv_state_v = self.v_conv1d( - x=v_proj_states, - cache=conv_state_v, - output_final_state=use_cache, - cu_seqlens=cu_seqlens, - ) - g = self.f_b_proj(self.f_a_proj(hidden_states)) - g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) - beta = self.b_proj(hidden_states).float() - - q, k = map( - lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim), (q, k) - ) - v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) - - if mode == "chunk": - o, recurrent_state = chunk_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=recurrent_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - safe_gate=self.gate_lower_bound is not None, - lower_bound=self.gate_lower_bound, - transpose_state_layout=True, - cu_seqlens=cu_seqlens, - ) - else: - o, recurrent_state = fused_recurrent_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=recurrent_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - lower_bound=self.gate_lower_bound, - transpose_state_layout=True, - cu_seqlens=cu_seqlens, - ) - if cache_params is not None: - cache_params.recurrent_states[self.layer_idx] = recurrent_state - cache_params.conv_states[self.layer_idx] = ( - conv_state_q, - conv_state_k, - conv_state_v, - ) - - if self.use_full_rank_gate: - g = self.g_proj(hidden_states) - else: - g = self.g_b_proj(self.g_a_proj(hidden_states)) - g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) - o = self.o_norm(o, g) - - o = rearrange(o, "b t h d -> b t (h d)") - o = self.o_proj(o) - if attention_mask is not None: - o = pad_input(o.squeeze(0), indices, batch_size, q_len) - - return o - - -class KimiMoEGate(nn.Module): - """ - MoEGate adapted from Deepseek-V3. - Parameter correspondences: - num_experts -> n_routed_experts - num_experts_per_token -> num_experts_per_tok - num_expert_group -> n_group - moe_router_activation_func -> scoring_func - """ - - def __init__(self, config: KimiLinearConfig): - super().__init__() - self.config = config - self.top_k = config.num_experts_per_token - self.num_experts = config.num_experts - self.routed_scaling_factor = config.routed_scaling_factor - self.moe_router_activation_func = config.moe_router_activation_func - self.num_expert_group = getattr(config, "num_expert_group", 1) - self.topk_group = getattr(config, "topk_group", 1) - - # topk selection algorithm - self.moe_renormalize = config.moe_renormalize - self.gating_dim = config.hidden_size - self.weight = nn.Parameter( - torch.empty((self.num_experts, self.gating_dim)), - ) - - self.e_score_correction_bias = nn.Parameter( - torch.empty(self.num_experts), - ) - self.reset_parameters() - - def reset_parameters(self) -> None: - import torch.nn.init as init - - init.kaiming_uniform_(self.weight, a=math.sqrt(5)) - - def forward(self, hidden_states): - bsz, seq_len, h = hidden_states.shape - # compute gating score - hidden_states = hidden_states.view(-1, h) - logits = F.linear( - hidden_states.type(torch.float32), - self.weight.type(torch.float32), - None, - ) - if self.moe_router_activation_func == "sigmoid": - scores = logits.sigmoid() - elif self.moe_router_activation_func == "softmax": - scores = logits.softmax(dim=1) - else: - raise NotImplementedError( - f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}", - ) - - # select top-k experts - scores = scores.view(bsz * seq_len, -1) - scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) - if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: - group_scores = ( - scores_for_choice.view(bsz * seq_len, self.num_expert_group, -1) - .topk(2, dim=-1)[0] - .sum(dim=-1) - ) # [n, num_expert_group] - group_idx = torch.topk( - group_scores, - k=self.topk_group, - dim=-1, - sorted=False, - )[1] # [n, top_k_group] - group_mask = torch.zeros_like(group_scores) # [n, num_expert_group] - group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group] - score_mask = ( - group_mask.unsqueeze(-1) - .expand( - bsz * seq_len, - self.num_expert_group, - self.num_experts // self.num_expert_group, - ) - .reshape(bsz * seq_len, -1) - ) # [n, e] - tmp_scores = scores_for_choice.masked_fill( - ~score_mask.bool(), float("-inf") - ) # [n, e] - else: - tmp_scores = scores_for_choice - _, topk_idx = torch.topk( - tmp_scores, - k=self.top_k, - dim=-1, - sorted=False, - ) - topk_weight = scores.gather(1, topk_idx) - - # norm gate to sum 1 - if self.top_k > 1 and self.moe_renormalize: - denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 - topk_weight = topk_weight / denominator - # must multiply the scaling factor - topk_weight = topk_weight * self.routed_scaling_factor - - return topk_idx, topk_weight - - -class KimiSparseMoeBlock(nn.Module): - """ - Adapted from Deepseek-V3's MOE implementation - The namings are consistent with Kimi's version. - """ - - def __init__(self, config: KimiLinearConfig): - super().__init__() - self.config = config - self.hidden_dim = config.hidden_size - self.num_experts = config.num_experts - self.top_k = config.num_experts_per_token - self.moe_renormalize = config.moe_renormalize - - self.use_latent_moe = ( - getattr(config, "routed_expert_hidden_size", None) is not None - ) - self.moe_hidden_size = ( - config.routed_expert_hidden_size - if self.use_latent_moe - else config.hidden_size - ) - self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False) - - self.ep_size = 1 - self.experts_per_rank = config.num_experts - self.ep_rank = 0 - self.experts = nn.ModuleList( - [ - KimiBlockSparseMLP( - config, - hidden_size=self.moe_hidden_size, - intermediate_size=config.moe_intermediate_size, - ) - for _ in range(config.num_experts) - ], - ) - self.gate = KimiMoEGate(config) - if config.num_shared_experts is not None: - intermediate_size = config.moe_intermediate_size * config.num_shared_experts - self.shared_experts = KimiMLP( - config=config, - intermediate_size=intermediate_size, - ) - - if self.use_latent_moe: - self.routed_expert_down_proj = nn.Linear( - config.hidden_size, - self.moe_hidden_size, - bias=False, - ) - self.routed_expert_up_proj = nn.Linear( - self.moe_hidden_size, - config.hidden_size, - bias=False, - ) - if self.latent_moe_use_norm: - self.routed_expert_norm = KimiRMSNorm( - self.moe_hidden_size, - eps=config.rms_norm_eps, - ) - - def forward(self, hidden_states): - identity = hidden_states - orig_shape = hidden_states.shape - topk_idx, topk_weight = self.gate(hidden_states) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - - if self.use_latent_moe: - hidden_states = self.routed_expert_down_proj(hidden_states) - if not self.training: - if False:#if not self.training: - y = self.moe_infer(hidden_states, topk_idx, topk_weight) - else: - y = self.moe_train(hidden_states, topk_idx, topk_weight) - - if self.use_latent_moe: - if self.latent_moe_use_norm: - y = self.routed_expert_norm(y) - y = self.routed_expert_up_proj(y) - - y = y.view(*orig_shape) - - if self.config.num_shared_experts is not None: - y = y + self.shared_experts(identity) - return y - - def moe_train(self, x, topk_ids, topk_weight): - """Training-compatible MoE dispatch with gradient flow.""" - y = torch.zeros_like(x) - - with torch.no_grad(): - expert_mask = F.one_hot(topk_ids, self.num_experts).permute(2, 1, 0) - - for expert_idx, expert in enumerate(self.experts): - top_k_pos, token_indices = torch.where(expert_mask[expert_idx]) - - if get_calibrate_all_experts_flag(): - expert_out = expert(x)[token_indices] - else: - expert_out = expert(x[token_indices]) - - expert_weights = topk_weight[token_indices, top_k_pos, None] - y.index_add_(0, token_indices, (expert_out * expert_weights).to(y.dtype)) - - return y - - @torch.no_grad() - def moe_infer(self, x, topk_ids, topk_weight): - cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) - cnts.scatter_(1, topk_ids, 1) - tokens_per_expert = cnts.sum(dim=0) - idxs = topk_ids.view(-1).argsort() - sorted_tokens = x[idxs // topk_ids.shape[1]] - - tokens_per_expert = tokens_per_expert.cpu().numpy() - - outputs = [] - start_idx = 0 - for i, num_tokens in enumerate(tokens_per_expert): - end_idx = start_idx + num_tokens - if num_tokens == 0: - continue - expert = self.experts[i + self.ep_rank * self.experts_per_rank] - tokens_for_this_expert = sorted_tokens[start_idx:end_idx] - expert_out = expert(tokens_for_this_expert) - outputs.append(expert_out) - start_idx = end_idx - - outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) - - new_x = torch.empty_like(outs) - new_x[idxs] = outs - final_out = ( - new_x.view(*topk_ids.shape, -1) - .type(topk_weight.dtype) - .mul_(topk_weight.unsqueeze(dim=-1)) - .sum(dim=1) - .type(new_x.dtype) - ) - return final_out - - -class KimiDecoderLayer(nn.Module): - def __init__(self, config: KimiLinearConfig, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - self.config = config - self.layer_idx = layer_idx - if config.is_kda_layer(layer_idx): - self.is_linear_attn = True - self.self_attn = KimiDeltaAttention(config=config, layer_idx=layer_idx) - elif config.is_mla: - self.is_linear_attn = False - self.self_attn = KimiMLAAttention(config=config, layer_idx=layer_idx) - else: - raise NotImplementedError - if ( - config.num_experts is not None - and layer_idx >= config.first_k_dense_replace - and layer_idx % getattr(config, "moe_layer_freq", 1) == 0 - ): - self.block_sparse_moe = KimiSparseMoeBlock(config) - else: - self.mlp = KimiMLP(config) - self.input_layernorm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = KimiRMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - # Attention residual - self.use_attn_residuals = ( - getattr(config, "attn_res_block_size", None) is not None - ) - if self.use_attn_residuals: - self.attn_res_block_size = config.attn_res_block_size - self.self_attention_res_norm = KimiRMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - self.mlp_res_norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.self_attention_res_proj = nn.Linear(config.hidden_size, 1, bias=False) - self.mlp_res_proj = nn.Linear(config.hidden_size, 1, bias=False) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: tuple[torch.Tensor] | None = None, - output_attentions: bool | None = False, - use_cache: bool | None = False, - block_residual: torch.Tensor | None = None, - **kwargs: Unpack[FlashAttentionKwargs], - ): - if self.use_attn_residuals: - return self._forward_attn_residual( - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - block_residual, - **kwargs, - ) - - residual = hidden_states - - hidden_states = self.input_layernorm(hidden_states) - - # Self Attention - if self.is_linear_attn is False: - hidden_states = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - **kwargs, - ) - else: - hidden_states = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - cache_params=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - **kwargs, - ) - hidden_states = residual + hidden_states - - # Fully Connected - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - if hasattr(self, "block_sparse_moe"): - hidden_states = self.block_sparse_moe(hidden_states) - else: - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states - - def _forward_attn_residual( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: tuple[torch.Tensor] | None = None, - output_attentions: bool | None = False, - use_cache: bool | None = False, - block_residual: torch.Tensor | None = None, - **kwargs: Unpack[FlashAttentionKwargs], - ): - batch_size, seq_len, hidden_size = hidden_states.shape - prefix_sum = hidden_states - is_block_start = self.layer_idx % self.attn_res_block_size == 0 - - hidden_states, block_residual = _attn_res_pre_update( - block_residual, - prefix_sum, - hidden_size, - batch_size, - seq_len, - self.self_attention_res_proj, - self.self_attention_res_norm, - is_block_start, - ) - - hidden_states = self.input_layernorm(hidden_states) - - # Self Attention - if self.is_linear_attn is False: - hidden_states = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - **kwargs, - ) - else: - hidden_states = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - cache_params=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - **kwargs, - ) - - if is_block_start: - prefix_sum = hidden_states - else: - prefix_sum = prefix_sum + hidden_states - - hidden_states = _apply_attn_res( - prefix_sum.view(-1, hidden_size), - block_residual, - self.mlp_res_proj, - self.mlp_res_norm, - ).view(batch_size, seq_len, hidden_size) - - hidden_states = self.post_attention_layernorm(hidden_states) - if hasattr(self, "block_sparse_moe"): - hidden_states = self.block_sparse_moe(hidden_states) - else: - hidden_states = self.mlp(hidden_states) - - prefix_sum = prefix_sum + hidden_states - - return prefix_sum, block_residual - - -class KimiPreTrainedModel(PreTrainedModel): - config_class = KimiLinearConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["KimiDecoderLayer"] - _skip_keys_device_placement = "past_key_values" - _supports_flash_attn_2 = True - _can_record_outputs = { - "router_logits": OutputRecorder(KimiBlockSparseMLP, index=1), - "hidden_states": KimiDecoderLayer, - "attentions": KimiMLAAttention, - } - _is_stateful = True - - def _init_weights(self, module): - # HOTFIX: disk offloading attempts to initialize the meta tensors - # but this is bad programming: we shouldn't be initializing these - # params in the first place - # the init attempt attempts to get `module.weight`, which DNE for qmodels - - return - - std = self.config.initializer_range - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=std) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.Embedding): - module.weight.data.normal_(mean=0.0, std=std) - if module.padding_idx is not None: - module.weight.data[module.padding_idx].zero_() - - -def _apply_attn_res(prefix_sum, block_residual, proj, norm): - """ - prefix_sum: (num_tokens, hidden_size) - block_residual: (num_tokens, num_blocks, hidden_size) - """ - v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1) - v_float = v.float() - variance = v_float.pow(2).mean(-1, keepdim=True) - k = v_float * torch.rsqrt(variance + norm.variance_epsilon) - score_weight = norm.weight.float() * proj.weight.squeeze(0).float() - scores = (k * score_weight).sum(-1) - probs = scores.softmax(-1).unsqueeze(1) - hidden_states = torch.matmul(probs, v_float).squeeze(1) - return hidden_states.to(v.dtype) - - -@torch.fx.wrap -def _attn_res_pre_update( - block_residual, - prefix_sum, - hidden_size, - batch_size, - seq_len, - self_attention_res_proj, - self_attention_res_norm, - is_block_start, -): - hidden_states = prefix_sum - if block_residual is not None and block_residual.shape[1] > 0: - hidden_states = _apply_attn_res( - prefix_sum.view(-1, hidden_size), - block_residual, - self_attention_res_proj, - self_attention_res_norm, - ).view(batch_size, seq_len, hidden_size) - if is_block_start: - block_residual = torch.cat( - [block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1 - ) - return hidden_states, block_residual - - -class KimiLinearModel(KimiPreTrainedModel): - def __init__(self, config: KimiLinearConfig): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - - self.embed_tokens = nn.Embedding( - config.vocab_size, config.hidden_size, self.padding_idx - ) - self.layers = nn.ModuleList( - [ - KimiDecoderLayer(config, layer_idx) - for layer_idx in range(config.num_hidden_layers) - ] - ) - self.norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - self.use_attn_residuals = ( - getattr(config, "attn_res_block_size", None) is not None - ) - if self.use_attn_residuals: - self.output_attn_res_norm = KimiRMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - self.output_attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False) - - from transformers.utils import is_flash_attn_2_available as _fa2_avail - - _requested = getattr(config, "_attn_implementation", None) - if _requested not in (None, "flash_attention_2") or not _fa2_avail(): - # Fall back gracefully when flash-attn2 is unavailable or a different impl is requested - if _requested == "flash_attention_2" and not _fa2_avail(): - logger.warning_once( - "flash_attention_2 requested but not available; falling back to sdpa." - ) - config._attn_implementation = ( - _requested if _requested not in (None, "flash_attention_2") else "eager" - ) - else: - config._attn_implementation = "flash_attention_2" - - self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" - self.gradient_checkpointing = False - # Initialize weights and apply final processing - self.post_init() - - def _update_linear_attn_mask(self, attention_mask, cache_position): - """ - NOTE: Left-padding is used for linear attention mask. - No need for zeroing states when - 1. Cached forward - 2. Attending to all inputs - """ - linear_attn_mask = attention_mask - if cache_position[0] > 0 or ( - attention_mask is not None and torch.all(attention_mask == 1) - ): - linear_attn_mask = None - return linear_attn_mask - - @check_model_inputs - # @auto_docstring - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - cache_position: torch.LongTensor | None = None, - use_cache: bool | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple | BaseModelOutputWithPast: - use_cache = use_cache if use_cache is not None else self.config.use_cache - - if (input_ids is None) and (inputs_embeds is None): - raise ValueError( - "You must specify exactly one of input_ids or inputs_embeds" - ) - - # Get inputs_embeds - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - if use_cache and past_key_values is None: - past_key_values = KimiDynamicCache(config=self.config) - - if cache_position is None: - past_seen_tokens = ( - past_key_values.get_seq_length() if past_key_values is not None else 0 - ) - cache_position: torch.Tensor = torch.arange( - past_seen_tokens, - past_seen_tokens + inputs_embeds.shape[1], - device=inputs_embeds.device, - ) - - if position_ids is None: - position_ids = cache_position.unsqueeze(0) - - causal_mask = create_causal_mask( - config=self.config, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - past_key_values=past_key_values, - position_ids=position_ids, - ) - linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position) - - hidden_states = inputs_embeds - if past_key_values is not None: - assert isinstance(past_key_values, KimiDynamicCache) - - block_residual = None - if self.use_attn_residuals: - block_residual = hidden_states.new_zeros( - hidden_states.shape[0] * hidden_states.shape[1], - 0, - hidden_states.shape[2], - ) - - for decoder_layer in self.layers: - layer_mask = ( - linear_attn_mask if decoder_layer.is_linear_attn else causal_mask - ) - - if self.use_attn_residuals: - hidden_states, block_residual = decoder_layer( - hidden_states, - attention_mask=layer_mask, - past_key_values=past_key_values, - cache_position=cache_position, - block_residual=block_residual, - **kwargs, - ) - else: - hidden_states = decoder_layer( - hidden_states, - attention_mask=layer_mask, - past_key_values=past_key_values, - cache_position=cache_position, - **kwargs, - ) - - if self.use_attn_residuals: - hidden_states = self._apply_output_attn_res(hidden_states, block_residual) - - hidden_states = self.norm(hidden_states) - - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=past_key_values, - ) - - def _apply_output_attn_res(self, hidden_states, block_residual): - batch_size, seq_len, hidden_size = hidden_states.shape - return _apply_attn_res( - hidden_states.view(-1, hidden_size), - block_residual, - self.output_attn_res_proj, - self.output_attn_res_norm, - ).view(batch_size, seq_len, hidden_size) - - -class KimiLinearForCausalLM(KimiPreTrainedModel, GenerationMixin): - @classmethod - def _supports_default_dynamic_cache(cls) -> bool: - return False - - _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} - - def __init__(self, config): - super().__init__(config) - self.model = KimiLinearModel(config) - self.vocab_size = config.vocab_size - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - - # Initialize weights and apply final processing - self.post_init() - - @can_return_tuple - # @auto_docstring - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: list[torch.FloatTensor] | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - generation_mode: bool | None = None, - return_dict: bool | None = None, - cache_position: torch.LongTensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple | CausalLMOutputWithPast: - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - """ - - output_attentions = ( - output_attentions - if output_attentions is not None - else self.config.output_attentions - ) - output_hidden_states = ( - output_hidden_states - if output_hidden_states is not None - else self.config.output_hidden_states - ) - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - outputs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - cache_position=cache_position, - ) - - logits = outputs[0] - if generation_mode: - logits = logits[:, -1:] - logits = self.lm_head(logits) - - loss = None - if labels is not None: - loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) diff --git a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py index a6be474412..371dc61557 100644 --- a/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py +++ b/src/llmcompressor/modeling/kimi_k3/modeling_kimi_linear.py @@ -1 +1,1501 @@ -from .modeling_kimi_k3_linear import * +# coding=utf-8 +# Copyright 2025-2026 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. +# +# The multi-head latent attention, MoE gating and sparse MoE block in this file are +# adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py). They have been +# extensively modified and extended for the Kimi-Linear architecture. +# +# Licensing Information: +# - Code adapted from DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0. +# - Other parts of the code are licensed under the Kimi K3 License (see the LICENSE file in this repository). +# +# Apache License, Version 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. + +# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: +# - KimiPreTrainedModel._init_weights: early return to prevent initialization +# of meta tensors during disk offloading (quantized models) +# - KimiSparseMoeBlock: added moe_train() method with gradient flow for +# calibration; uses get_calibrate_all_experts_flag() to run all experts +# during calibration. Inference path disabled (if False) to force training +# path during quantization +# - KimiMoEGate.forward: removed `assert not self.training` to allow +# calibration in training mode +# - KimiDecoderLayer._forward_attn_residual: refactored attention residual +# logic into _attn_res_pre_update() decorated with @torch.fx.wrap for +# torch.compile compatibility +# - KimiDynamicCache: added get_query_offset() method; get_mask_sizes() +# handles both int and tensor cache_position (new transformers API) +# - KimiLinearModel.__init__: attention implementation fallback changed from +# forcing flash_attention_2 to graceful fallback to sdpa/eager when FA2 +# is unavailable +# - create_causal_mask call: fixed input_embeds -> inputs_embeds typo, +# removed cache_position kwarg +# - KimiLinearForCausalLM._tied_weights_keys: changed from list to dict +# mapping format +import math +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn.functional as F +import transformers +from einops import rearrange +from packaging import version +from torch import nn +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache +from transformers.generation import GenerationMixin +from transformers.masking_utils import create_causal_mask +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS +from transformers.utils import ( + TransformersKwargs, + can_return_tuple, + logging, +) +from transformers.utils.generic import check_model_inputs +from transformers.utils.output_capturing import OutputRecorder + +try: + from fla.modules import FusedRMSNormGated, ShortConvolution + from fla.ops.kda import chunk_kda, fused_recurrent_kda + + # from fla.ops.kda.gate import fused_kda_gate # deprecated, gate is now computed inside chunk_kda/fused_recurrent_kda + from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask + from fla.utils import tensor_cache +except ImportError: + raise ImportError("Plese run `pip install -U fla-core`") + +from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag + +from .configuration_kimi_k3 import KimiLinearConfig + +assert version.parse(transformers.__version__) >= version.parse( + "4.56.0" +), "Please upgrade transformers to >= 4.56.0" + +logger = logging.get_logger(__name__) + + +# Register Moonshot-specific activation functions +class SituAndMul(nn.Module): + """ + SituAndMul activation: beta * tanh(gate / beta) * sigmoid(gate) * up + When linear_beta is set, up is also transformed by linear_beta * tanh(up / linear_beta). + """ + + def __init__(self, beta: float = 1.0, linear_beta: float | None = None): + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d].to(torch.float32) + up = x[..., d:].to(torch.float32) + situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (situ_a * up).to(x.dtype) + + +ACT2FN["situ"] = SituAndMul + + +def _get_situ_activation_params(config: KimiLinearConfig): + beta = getattr(config, "activation_situ_beta", None) + linear_beta = getattr(config, "activation_situ_linear_beta", None) + return beta or 1.0, linear_beta + + +def index_first_axis(x, indices): + return x[indices] + + +@tensor_cache +def get_unpad_data( + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, int]: + lens = prepare_lens_from_mask(attention_mask) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = lens.max().item() + cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) + return indices, cu_seqlens, max_seqlen_in_batch + + +def pad_input( + hidden_states: torch.Tensor, + indices: torch.LongTensor, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + out = hidden_states.new_zeros((batch_size * seq_len, *hidden_states.shape[1:])) + out[indices] = hidden_states + return out.view(batch_size, seq_len, *hidden_states.shape[1:]) + + +class KimiDynamicCache: + """ + Dynamic cache for Kimi model. + Inspired by Qwen3-Next + """ + + is_compileable = False + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + + if config.linear_attn_config is not None: + self.layer_types = [] + for i in range(config.num_hidden_layers): + if config.is_kda_layer(i): + self.layer_types.append("linear_attention") + else: + self.layer_types.append("full_attention") + else: + self.layer_types = ["full_attention"] * config.num_hidden_layers + + self.transformer_layers = [ + i + for i in range(config.num_hidden_layers) + if self.layer_types[i] == "full_attention" + ] + + linear_layers = [ + i + for i in range(config.num_hidden_layers) + if self.layer_types[i] == "linear_attention" + ] + self.last_linear_layer = linear_layers[-1] if linear_layers else -1 + + self.conv_states = [None for _ in range(config.num_hidden_layers)] + self.recurrent_states = [None for _ in range(config.num_hidden_layers)] + self.key_cache = [None for _ in range(config.num_hidden_layers)] + self.value_cache = [None for _ in range(config.num_hidden_layers)] + + def __len__(self): + return len(self.layer_types) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.key_cache[layer_idx] is None: + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat( + [self.key_cache[layer_idx], key_states], dim=2 + ) + self.value_cache[layer_idx] = torch.cat( + [self.value_cache[layer_idx], value_states], dim=2 + ) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx] is not None: + device = self.key_cache[layer_idx].device + beam_idx = beam_idx.to(device) + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select( + 0, beam_idx + ) + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select( + 0, beam_idx + ) + + if self.conv_states[layer_idx] is not None: + device = self.conv_states[layer_idx][0].device + beam_idx = beam_idx.to(device) + q_conv, k_conv, v_conv = self.conv_states[layer_idx] + self.conv_states[layer_idx] = ( + q_conv.index_select(0, beam_idx), + k_conv.index_select(0, beam_idx), + v_conv.index_select(0, beam_idx), + ) + self.recurrent_states[layer_idx] = self.recurrent_states[ + layer_idx + ].index_select(0, beam_idx) + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # take any layer that contains cache and not empty tensor + layer_idx = ( + self.transformer_layers[0] + if layer_idx not in self.transformer_layers + else layer_idx + ) + if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx] is None: + return 0 + return self.key_cache[layer_idx].shape[-2] + + def get_query_offset(self, layer_idx: int = 0) -> int: + return self.get_seq_length(layer_idx=layer_idx) + + def get_mask_sizes(self, cache_position, layer_idx: int) -> tuple[int, int]: + """ + Return a tuple (kv_length, kv_offset) corresponding to the length and offset that will be returned for + the given layer at `layer_idx`. + The masks are then prepared according to the given lengths (kv_length, kv_offset) and patterns for each layer. + """ + kv_offset = 0 + # cache_position may be an int (new API) or a 1-D tensor (old API) + query_length = ( + cache_position + if isinstance(cache_position, int) + else cache_position.shape[0] + ) + past_seen_tokens = self.get_seq_length(layer_idx) + kv_length = query_length + past_seen_tokens + return kv_length, kv_offset + + @property + def has_previous_state(self): + """We have a previous state if the last linear (conv) layer was already updated.""" + if self.last_linear_layer == -1: + return False + return self.conv_states[self.last_linear_layer] is not None + + +class KimiRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + dtype = hidden_states.dtype + x = hidden_states.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + return self.weight * x.to(dtype) + + +ALL_LAYERNORM_LAYERS.append(KimiRMSNorm) + + +class KimiBlockSparseMLP(nn.Module): + def __init__( + self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None + ): + super().__init__() + self.config = config + self.ffn_dim = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + self.hidden_dim = config.hidden_size if hidden_size is None else hidden_size + + self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # gate + self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) # down + self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) # up + + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states): + if self.config.hidden_act == "situ": + gate_up = torch.cat( + [self.w1(hidden_states), self.w3(hidden_states)], dim=-1 + ) + current_hidden_states = self.act_fn(gate_up) + else: + current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3( + hidden_states + ) + current_hidden_states = self.w2(current_hidden_states) + return current_hidden_states + + +class KimiMLP(nn.Module): + def __init__( + self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + self.intermediate_size = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + if config.hidden_act == "situ": + beta, linear_beta = _get_situ_activation_params(config) + self.act_fn = SituAndMul( + beta=beta, + linear_beta=linear_beta, + ) + else: + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + if self.config.hidden_act == "situ": + gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1) + down_proj = self.down_proj(self.act_fn(gate_up)) + else: + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand the key/value heads from `num_key_value_heads` to `num_attention_heads`.""" + if n_rep == 1: + return hidden_states + return torch.repeat_interleave(hidden_states, dim=1, repeats=n_rep) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling + if attention_mask is not None: + scores = scores + attention_mask[:, :, :, : key.shape[-2]] + + probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) + probs = F.dropout(probs, p=dropout, training=module.training) + out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous() + + return out, probs + + +class KimiMLAAttention(nn.Module): + """ + Multi-Latent Attention adapted from deepseek-v3 + """ + + def __init__(self, config: KimiLinearConfig, layer_idx: int): + nn.Module.__init__(self) + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + + self.attention_dropout = getattr(config, "attention_dropout", 0.0) + + try: + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.use_nope = config.mla_use_nope + self.scaling = self.q_head_dim ** (-0.5) + except Exception as e: + raise ValueError( + f"Kimi MLA config is not found or not properly formatted: {e}" + ) + + if self.q_lora_rank is not None: + self.q_a_proj = nn.Linear( + self.hidden_size, + self.q_lora_rank, + bias=False, + ) + self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank) + self.q_b_proj = nn.Linear( + self.q_lora_rank, + self.num_heads * self.q_head_dim, + bias=False, + ) + else: + self.q_proj = nn.Linear( + self.hidden_size, + self.num_heads * self.q_head_dim, + bias=False, + ) + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + ) + self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads + * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + ) + self.is_causal = True + assert self.use_nope + + self.use_output_gate = getattr(config, "mla_use_output_gate", False) + if self.use_output_gate: + projection_size = self.num_heads * self.v_head_dim + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.rotary_emb = None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + batch_size, seq_length = hidden_states.shape[:-1] + query_shape = (batch_size, seq_length, -1, self.q_head_dim) + key_shape = ( + batch_size, + seq_length, + -1, + self.qk_nope_head_dim + self.v_head_dim, + ) + + if self.q_lora_rank is not None: + q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + else: + q_states = self.q_proj(hidden_states) + q_states = q_states.view(query_shape).transpose(1, 2) + q_pass, q_rot = torch.split( + q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + k_pass, k_rot = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + k_pass = ( + self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) + ) + k_pass, value_states = torch.split( + k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + + k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + + k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + + query_states = torch.cat((q_pass, q_rot), dim=-1) + key_states = torch.cat((k_pass, k_rot), dim=-1) + + if past_key_values is not None: + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx + ) + + if ( + self.config._attn_implementation == "flash_attention_2" + and self.q_head_dim != self.v_head_dim + ): + value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[ + self.config._attn_implementation + ] + + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + if ( + self.config._attn_implementation == "flash_attention_2" + and self.q_head_dim != self.v_head_dim + ): + attn_output = attn_output[:, :, :, : self.v_head_dim] + + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + if self.use_output_gate: + g = self.g_proj(hidden_states).sigmoid() + attn_output = attn_output * g + attn_output = self.o_proj(attn_output) + return attn_output + + +class KimiDeltaAttention(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.config = config + self.mode = "chunk" + + self.hidden_size = config.hidden_size + self.conv_size = config.linear_attn_config["short_conv_kernel_size"] + self.head_dim = config.linear_attn_config["head_dim"] + self.num_heads = config.linear_attn_config["num_heads"] + self.head_k_dim = self.head_dim + self.num_k_heads = self.num_heads + + self.layer_idx = layer_idx + + assert self.mode in [ + "chunk", + "fused_recurrent", + ], f"Not supported mode `{self.mode}`." + + projection_k_size = self.head_k_dim * self.num_k_heads + projection_size = self.head_dim * self.num_heads + + self.q_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False) + self.k_proj = nn.Linear(self.hidden_size, projection_k_size, bias=False) + self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + + self.q_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=self.conv_size, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=projection_size, + kernel_size=self.conv_size, + activation="silu", + ) + + self.A_log = torch.nn.Parameter( + torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)) + ) + + self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32)) + + self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False) + + self.use_full_rank_gate = config.linear_attn_config.get( + "use_full_rank_gate", False + ) + self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None) + if self.use_full_rank_gate: + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + else: + self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + + self.o_norm = FusedRMSNormGated( + self.head_dim, eps=config.rms_norm_eps, activation="sigmoid" + ) + self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + cache_params: KimiDynamicCache | None = None, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + if attention_mask.dim() != 2: + attention_mask = kwargs.get("padding_mask") + + if attention_mask is not None and attention_mask.dim() != 2: + raise ValueError( + "attention_mask must be a 0-1 matrix of shape [batch_size, seq_len] " + "(0 = padding). 3D masks are not supported here.", + ) + use_cache = cache_params is not None + batch_size, q_len, _ = hidden_states.shape + mode = "fused_recurrent" if use_cache and q_len == 1 else self.mode + if self.training: + assert mode == "chunk", "Only chunk mode is supported in training." + + cu_seqlens = kwargs.get("cu_seqlens") + indices = None + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis( + rearrange(hidden_states, "b s ... -> (b s) ..."), indices + ).unsqueeze(0) + + conv_state_q, conv_state_k, conv_state_v = None, None, None + recurrent_state = None + if cache_params is not None: + if cache_params.conv_states[self.layer_idx] is not None: + conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[ + self.layer_idx + ] + recurrent_state = cache_params.recurrent_states[self.layer_idx] + + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + q, conv_state_q = self.q_conv1d( + x=q_proj_states, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=k_proj_states, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=v_proj_states, + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + g = self.f_b_proj(self.f_a_proj(hidden_states)) + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) + beta = self.b_proj(hidden_states).float() + + q, k = map( + lambda x: rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim), (q, k) + ) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + if mode == "chunk": + o, recurrent_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=self.gate_lower_bound is not None, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + else: + o, recurrent_state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + transpose_state_layout=True, + cu_seqlens=cu_seqlens, + ) + if cache_params is not None: + cache_params.recurrent_states[self.layer_idx] = recurrent_state + cache_params.conv_states[self.layer_idx] = ( + conv_state_q, + conv_state_k, + conv_state_v, + ) + + if self.use_full_rank_gate: + g = self.g_proj(hidden_states) + else: + g = self.g_b_proj(self.g_a_proj(hidden_states)) + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, g) + + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o + + +class KimiMoEGate(nn.Module): + """ + MoEGate adapted from Deepseek-V3. + Parameter correspondences: + num_experts -> n_routed_experts + num_experts_per_token -> num_experts_per_tok + num_expert_group -> n_group + moe_router_activation_func -> scoring_func + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + + # topk selection algorithm + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim)), + ) + + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + import torch.nn.init as init + + init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states): + bsz, seq_len, h = hidden_states.shape + # compute gating score + hidden_states = hidden_states.view(-1, h) + logits = F.linear( + hidden_states.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + if self.moe_router_activation_func == "sigmoid": + scores = logits.sigmoid() + elif self.moe_router_activation_func == "softmax": + scores = logits.softmax(dim=1) + else: + raise NotImplementedError( + f"insupportable scoring function for MoE gating: {self.moe_router_activation_func}", + ) + + # select top-k experts + scores = scores.view(bsz * seq_len, -1) + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: + group_scores = ( + scores_for_choice.view(bsz * seq_len, self.num_expert_group, -1) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) # [n, num_expert_group] + group_idx = torch.topk( + group_scores, + k=self.topk_group, + dim=-1, + sorted=False, + )[1] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, num_expert_group] + group_mask.scatter_(1, group_idx, 1) # [n, num_expert_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + bsz * seq_len, + self.num_expert_group, + self.num_experts // self.num_expert_group, + ) + .reshape(bsz * seq_len, -1) + ) # [n, e] + tmp_scores = scores_for_choice.masked_fill( + ~score_mask.bool(), float("-inf") + ) # [n, e] + else: + tmp_scores = scores_for_choice + _, topk_idx = torch.topk( + tmp_scores, + k=self.top_k, + dim=-1, + sorted=False, + ) + topk_weight = scores.gather(1, topk_idx) + + # norm gate to sum 1 + if self.top_k > 1 and self.moe_renormalize: + denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 + topk_weight = topk_weight / denominator + # must multiply the scaling factor + topk_weight = topk_weight * self.routed_scaling_factor + + return topk_idx, topk_weight + + +class KimiSparseMoeBlock(nn.Module): + """ + Adapted from Deepseek-V3's MOE implementation + The namings are consistent with Kimi's version. + """ + + def __init__(self, config: KimiLinearConfig): + super().__init__() + self.config = config + self.hidden_dim = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_token + self.moe_renormalize = config.moe_renormalize + + self.use_latent_moe = ( + getattr(config, "routed_expert_hidden_size", None) is not None + ) + self.moe_hidden_size = ( + config.routed_expert_hidden_size + if self.use_latent_moe + else config.hidden_size + ) + self.latent_moe_use_norm = getattr(config, "latent_moe_use_norm", False) + + self.ep_size = 1 + self.experts_per_rank = config.num_experts + self.ep_rank = 0 + self.experts = nn.ModuleList( + [ + KimiBlockSparseMLP( + config, + hidden_size=self.moe_hidden_size, + intermediate_size=config.moe_intermediate_size, + ) + for _ in range(config.num_experts) + ], + ) + self.gate = KimiMoEGate(config) + if config.num_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.num_shared_experts + self.shared_experts = KimiMLP( + config=config, + intermediate_size=intermediate_size, + ) + + if self.use_latent_moe: + self.routed_expert_down_proj = nn.Linear( + config.hidden_size, + self.moe_hidden_size, + bias=False, + ) + self.routed_expert_up_proj = nn.Linear( + self.moe_hidden_size, + config.hidden_size, + bias=False, + ) + if self.latent_moe_use_norm: + self.routed_expert_norm = KimiRMSNorm( + self.moe_hidden_size, + eps=config.rms_norm_eps, + ) + + def forward(self, hidden_states): + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + + if self.use_latent_moe: + hidden_states = self.routed_expert_down_proj(hidden_states) + if not self.training: + if False:#if not self.training: + y = self.moe_infer(hidden_states, topk_idx, topk_weight) + else: + y = self.moe_train(hidden_states, topk_idx, topk_weight) + + if self.use_latent_moe: + if self.latent_moe_use_norm: + y = self.routed_expert_norm(y) + y = self.routed_expert_up_proj(y) + + y = y.view(*orig_shape) + + if self.config.num_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + def moe_train(self, x, topk_ids, topk_weight): + """Training-compatible MoE dispatch with gradient flow.""" + y = torch.zeros_like(x) + + with torch.no_grad(): + expert_mask = F.one_hot(topk_ids, self.num_experts).permute(2, 1, 0) + + for expert_idx, expert in enumerate(self.experts): + top_k_pos, token_indices = torch.where(expert_mask[expert_idx]) + + if get_calibrate_all_experts_flag(): + expert_out = expert(x)[token_indices] + else: + expert_out = expert(x[token_indices]) + + expert_weights = topk_weight[token_indices, top_k_pos, None] + y.index_add_(0, token_indices, (expert_out * expert_weights).to(y.dtype)) + + return y + + @torch.no_grad() + def moe_infer(self, x, topk_ids, topk_weight): + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + return final_out + + +class KimiDecoderLayer(nn.Module): + def __init__(self, config: KimiLinearConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.config = config + self.layer_idx = layer_idx + if config.is_kda_layer(layer_idx): + self.is_linear_attn = True + self.self_attn = KimiDeltaAttention(config=config, layer_idx=layer_idx) + elif config.is_mla: + self.is_linear_attn = False + self.self_attn = KimiMLAAttention(config=config, layer_idx=layer_idx) + else: + raise NotImplementedError + if ( + config.num_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % getattr(config, "moe_layer_freq", 1) == 0 + ): + self.block_sparse_moe = KimiSparseMoeBlock(config) + else: + self.mlp = KimiMLP(config) + self.input_layernorm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + # Attention residual + self.use_attn_residuals = ( + getattr(config, "attn_res_block_size", None) is not None + ) + if self.use_attn_residuals: + self.attn_res_block_size = config.attn_res_block_size + self.self_attention_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_res_norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attention_res_proj = nn.Linear(config.hidden_size, 1, bias=False) + self.mlp_res_proj = nn.Linear(config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + if self.use_attn_residuals: + return self._forward_attn_residual( + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + block_residual, + **kwargs, + ) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + def _forward_attn_residual( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + block_residual: torch.Tensor | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ): + batch_size, seq_len, hidden_size = hidden_states.shape + prefix_sum = hidden_states + is_block_start = self.layer_idx % self.attn_res_block_size == 0 + + hidden_states, block_residual = _attn_res_pre_update( + block_residual, + prefix_sum, + hidden_size, + batch_size, + seq_len, + self.self_attention_res_proj, + self.self_attention_res_norm, + is_block_start, + ) + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + if self.is_linear_attn is False: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + cache_params=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + + if is_block_start: + prefix_sum = hidden_states + else: + prefix_sum = prefix_sum + hidden_states + + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self.mlp_res_proj, + self.mlp_res_norm, + ).view(batch_size, seq_len, hidden_size) + + hidden_states = self.post_attention_layernorm(hidden_states) + if hasattr(self, "block_sparse_moe"): + hidden_states = self.block_sparse_moe(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + + prefix_sum = prefix_sum + hidden_states + + return prefix_sum, block_residual + + +class KimiPreTrainedModel(PreTrainedModel): + config_class = KimiLinearConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["KimiDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _can_record_outputs = { + "router_logits": OutputRecorder(KimiBlockSparseMLP, index=1), + "hidden_states": KimiDecoderLayer, + "attentions": KimiMLAAttention, + } + _is_stateful = True + + def _init_weights(self, module): + # HOTFIX: disk offloading attempts to initialize the meta tensors + # but this is bad programming: we shouldn't be initializing these + # params in the first place + # the init attempt attempts to get `module.weight`, which DNE for qmodels + + return + + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +def _apply_attn_res(prefix_sum, block_residual, proj, norm): + """ + prefix_sum: (num_tokens, hidden_size) + block_residual: (num_tokens, num_blocks, hidden_size) + """ + v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1) + v_float = v.float() + variance = v_float.pow(2).mean(-1, keepdim=True) + k = v_float * torch.rsqrt(variance + norm.variance_epsilon) + score_weight = norm.weight.float() * proj.weight.squeeze(0).float() + scores = (k * score_weight).sum(-1) + probs = scores.softmax(-1).unsqueeze(1) + hidden_states = torch.matmul(probs, v_float).squeeze(1) + return hidden_states.to(v.dtype) + + +@torch.fx.wrap +def _attn_res_pre_update( + block_residual, + prefix_sum, + hidden_size, + batch_size, + seq_len, + self_attention_res_proj, + self_attention_res_norm, + is_block_start, +): + hidden_states = prefix_sum + if block_residual is not None and block_residual.shape[1] > 0: + hidden_states = _apply_attn_res( + prefix_sum.view(-1, hidden_size), + block_residual, + self_attention_res_proj, + self_attention_res_norm, + ).view(batch_size, seq_len, hidden_size) + if is_block_start: + block_residual = torch.cat( + [block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1 + ) + return hidden_states, block_residual + + +class KimiLinearModel(KimiPreTrainedModel): + def __init__(self, config: KimiLinearConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = nn.ModuleList( + [ + KimiDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = KimiRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.use_attn_residuals = ( + getattr(config, "attn_res_block_size", None) is not None + ) + if self.use_attn_residuals: + self.output_attn_res_norm = KimiRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.output_attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False) + + from transformers.utils import is_flash_attn_2_available as _fa2_avail + + _requested = getattr(config, "_attn_implementation", None) + if _requested not in (None, "flash_attention_2") or not _fa2_avail(): + # Fall back gracefully when flash-attn2 is unavailable or a different impl is requested + if _requested == "flash_attention_2" and not _fa2_avail(): + logger.warning_once( + "flash_attention_2 requested but not available; falling back to sdpa." + ) + config._attn_implementation = ( + _requested if _requested not in (None, "flash_attention_2") else "eager" + ) + else: + config._attn_implementation = "flash_attention_2" + + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def _update_linear_attn_mask(self, attention_mask, cache_position): + """ + NOTE: Left-padding is used for linear attention mask. + No need for zeroing states when + 1. Cached forward + 2. Attending to all inputs + """ + linear_attn_mask = attention_mask + if cache_position[0] > 0 or ( + attention_mask is not None and torch.all(attention_mask == 1) + ): + linear_attn_mask = None + return linear_attn_mask + + @check_model_inputs + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) and (inputs_embeds is None): + raise ValueError( + "You must specify exactly one of input_ids or inputs_embeds" + ) + + # Get inputs_embeds + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = KimiDynamicCache(config=self.config) + + if cache_position is None: + past_seen_tokens = ( + past_key_values.get_seq_length() if past_key_values is not None else 0 + ) + cache_position: torch.Tensor = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position) + + hidden_states = inputs_embeds + if past_key_values is not None: + assert isinstance(past_key_values, KimiDynamicCache) + + block_residual = None + if self.use_attn_residuals: + block_residual = hidden_states.new_zeros( + hidden_states.shape[0] * hidden_states.shape[1], + 0, + hidden_states.shape[2], + ) + + for decoder_layer in self.layers: + layer_mask = ( + linear_attn_mask if decoder_layer.is_linear_attn else causal_mask + ) + + if self.use_attn_residuals: + hidden_states, block_residual = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + block_residual=block_residual, + **kwargs, + ) + else: + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + if self.use_attn_residuals: + hidden_states = self._apply_output_attn_res(hidden_states, block_residual) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + def _apply_output_attn_res(self, hidden_states, block_residual): + batch_size, seq_len, hidden_size = hidden_states.shape + return _apply_attn_res( + hidden_states.view(-1, hidden_size), + block_residual, + self.output_attn_res_proj, + self.output_attn_res_norm, + ).view(batch_size, seq_len, hidden_size) + + +class KimiLinearForCausalLM(KimiPreTrainedModel, GenerationMixin): + @classmethod + def _supports_default_dynamic_cache(cls) -> bool: + return False + + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = KimiLinearModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + @can_return_tuple + # @auto_docstring + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + generation_mode: bool | None = None, + return_dict: bool | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | CausalLMOutputWithPast: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + """ + + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + logits = outputs[0] + if generation_mode: + logits = logits[:, -1:] + logits = self.lm_head(logits) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py b/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py index 5bca540870..a9d00030cf 100644 --- a/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py +++ b/src/llmcompressor/modeling/kimi_k3/tokenization_kimi.py @@ -1,6 +1,3 @@ -# Differences from https://huggingface.co/moonshotai/Kimi-K3/tree/main: -# - Formatting only (no functional changes) - import os from logging import getLogger from pathlib import Path From 15e99ed80ad59d87a20ec866b1d6e4ccf03e0b31 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 18 Aug 2026 01:08:02 +0000 Subject: [PATCH 5/6] docs: add Kimi-K3 to key models, exclude modeling files from linting Add Kimi-K3 key models documentation page with NVFP4 example, update nav and index. Exclude kimi_k3 modeling files from ruff since they are vendored upstream files. Fix import ordering in kimi_k3_example.py. Co-Authored-By: Claude Opus 4.6 --- docs/.nav.yml | 3 + docs/key-models/index.md | 10 ++- docs/key-models/kimi-k3/index.md | 7 ++ docs/key-models/kimi-k3/nvfp4-example.md | 100 +++++++++++++++++++++ examples/quantizing_moe/kimi_k3_example.py | 2 +- pyproject.toml | 2 +- 6 files changed, 121 insertions(+), 3 deletions(-) create mode 100644 docs/key-models/kimi-k3/index.md create mode 100644 docs/key-models/kimi-k3/nvfp4-example.md diff --git a/docs/.nav.yml b/docs/.nav.yml index 3303806710..5a340c4ad2 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -20,6 +20,9 @@ nav: - key-models/kimi-k26/index.md - NVFP4 Example: key-models/kimi-k26/nvfp4-example.md - FP8 Block Example: key-models/kimi-k26/fp8-block-example.md + - Kimi-K3: + - key-models/kimi-k3/index.md + - NVFP4 Example: key-models/kimi-k3/nvfp4-example.md - Qwen3.5: - key-models/qwen3.5/index.md - NVFP4A16 VL Example: key-models/qwen3.5/nvfp4-vl-example.md diff --git a/docs/key-models/index.md b/docs/key-models/index.md index 35585fe70a..1b7c49d7a7 100644 --- a/docs/key-models/index.md +++ b/docs/key-models/index.md @@ -1,6 +1,6 @@ # Key Models -The following models are among the most commonly used with LLM Compressor: Llama 4, Qwen3.5, Qwen3.6, Kimi-K2, and Mistral Large 3. Each model page contains quantization examples with tested configurations and recommended parameters. +The following models are among the most commonly used with LLM Compressor: Llama 4, Qwen3.5, Qwen3.6, Kimi-K2, Kimi-K3, and Mistral Large 3. Each model page contains quantization examples with tested configurations and recommended parameters.
@@ -37,6 +37,14 @@ The following models are among the most commonly used with LLM Compressor: Llama [:octicons-arrow-right-24: Kimi-K2.6](kimi-k26/index.md) + - **Kimi-K3** + + --- + + Moonshot AI's Kimi-K3 multimodal model, quantized to NVFP4. + + [:octicons-arrow-right-24: Kimi-K3](kimi-k3/index.md) + - **Gemma 4** --- diff --git a/docs/key-models/kimi-k3/index.md b/docs/key-models/kimi-k3/index.md new file mode 100644 index 0000000000..637e7cd1b6 --- /dev/null +++ b/docs/key-models/kimi-k3/index.md @@ -0,0 +1,7 @@ +# Kimi K3 + +Quantization examples for the Kimi K3 model. + +## Examples + +- [NVFP4 Example](nvfp4-example.md) diff --git a/docs/key-models/kimi-k3/nvfp4-example.md b/docs/key-models/kimi-k3/nvfp4-example.md new file mode 100644 index 0000000000..729c17f951 --- /dev/null +++ b/docs/key-models/kimi-k3/nvfp4-example.md @@ -0,0 +1,100 @@ +## Kimi-K3 NVFP4 Example + +### Overview + +Kimi-K3 requires custom modeling files bundled with LLM Compressor, since it is not yet supported in Transformers. +The example below quantizes the model to NVFP4 using calibration data. + +The full example script can be found [here](../../../examples/quantizing_moe/kimi_k3_example.py). + +### Code Walkthrough + +```python +from compressed_tensors.quantization import QuantizationConfig +from transformers import AutoTokenizer + +from datasets import load_dataset +from llmcompressor import oneshot +from llmcompressor.modeling.kimi_k3 import KimiK3ForConditionalGeneration +from llmcompressor.modifiers.quantization import QuantizationModifier +from llmcompressor.utils import load_context + +MODEL_ID = "moonshotai/Kimi-K3" + +# Load quantization config from pretrained and add ignore patterns +# for modules that should not be quantized +qconfig = QuantizationConfig.from_pretrained(MODEL_ID) +qconfig.ignore += [ + "re:.*mlp_res_proj.*", + "re:.*self_attention_res_proj.*", + "re:.*routed_expert.*", + "re:.*output_attn_res_proj.*", +] + +# Load model with the modified quantization config +with load_context(KimiK3ForConditionalGeneration): + model = KimiK3ForConditionalGeneration.from_pretrained( + MODEL_ID, + quantization_config=qconfig, + device_map="auto", + torch_dtype="auto", + trust_remote_code=True, + ) +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + +DATASET_ID = "HuggingFaceH4/ultrachat_200k" +DATASET_SPLIT = "train_sft" +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 2048 + +# Load dataset and preprocess +ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]") +ds = ds.shuffle(seed=42) + + +def preprocess(example): + return { + "text": tokenizer.apply_chat_template( + example["messages"], + tokenize=False, + ) + } + + +ds = ds.map(preprocess) + + +def tokenize(sample): + return tokenizer( + sample["text"], + padding=False, + max_length=MAX_SEQUENCE_LENGTH, + truncation=True, + add_special_tokens=False, + ) + + +ds = ds.map(tokenize, remove_columns=ds.column_names) + +recipe = QuantizationModifier( + targets="Linear", + scheme="NVFP4", + ignore=[ + "lm_head", + r"re:.*block_sparse_moe\.gate", + "re:.*vision_tower.*", + ], +) + +oneshot( + model=model, + dataset=ds, + recipe=recipe, + max_seq_length=MAX_SEQUENCE_LENGTH, + num_calibration_samples=NUM_CALIBRATION_SAMPLES, +) + +SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4" +model.save_pretrained(SAVE_DIR) +tokenizer.save_pretrained(SAVE_DIR) +``` diff --git a/examples/quantizing_moe/kimi_k3_example.py b/examples/quantizing_moe/kimi_k3_example.py index de7756111b..559a6da4d4 100644 --- a/examples/quantizing_moe/kimi_k3_example.py +++ b/examples/quantizing_moe/kimi_k3_example.py @@ -1,7 +1,7 @@ from compressed_tensors.quantization import QuantizationConfig +from datasets import load_dataset from transformers import AutoTokenizer -from datasets import load_dataset from llmcompressor import oneshot from llmcompressor.modeling.kimi_k3 import KimiK3ForConditionalGeneration from llmcompressor.modifiers.quantization import QuantizationModifier diff --git a/pyproject.toml b/pyproject.toml index 7edb5d3446..7f8065a8e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" files = "src/llmcompressor" [tool.ruff] -extend-exclude = ["env", "src/llmcompressor/transformers/tracing/", "src/llmcompressor/version.py"] +extend-exclude = ["env", "src/llmcompressor/transformers/tracing/", "src/llmcompressor/modeling/kimi_k3/", "src/llmcompressor/version.py"] line-length = 88 lint.select = ["E", "F", "W", "I"] lint.extend-ignore = ["E203", "W605"] From f4fbc036b72643d8978f79667500f1d051401081 Mon Sep 17 00:00:00 2001 From: Kyle Sayers Date: Tue, 18 Aug 2026 01:23:38 +0000 Subject: [PATCH 6/6] docs: add Kimi-K3 FP8 Block example to key models Co-Authored-By: Claude Opus 4.6 --- docs/.nav.yml | 1 + docs/key-models/kimi-k3/fp8-block-example.md | 46 ++++++++++++++++++++ docs/key-models/kimi-k3/index.md | 1 + examples/model_free_ptq/kimi_k3_fp8_block.py | 33 ++++++++++++++ 4 files changed, 81 insertions(+) create mode 100644 docs/key-models/kimi-k3/fp8-block-example.md create mode 100644 examples/model_free_ptq/kimi_k3_fp8_block.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 5a340c4ad2..7e6d744609 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -23,6 +23,7 @@ nav: - Kimi-K3: - key-models/kimi-k3/index.md - NVFP4 Example: key-models/kimi-k3/nvfp4-example.md + - FP8 Block Example: key-models/kimi-k3/fp8-block-example.md - Qwen3.5: - key-models/qwen3.5/index.md - NVFP4A16 VL Example: key-models/qwen3.5/nvfp4-vl-example.md diff --git a/docs/key-models/kimi-k3/fp8-block-example.md b/docs/key-models/kimi-k3/fp8-block-example.md new file mode 100644 index 0000000000..6f441ff84b --- /dev/null +++ b/docs/key-models/kimi-k3/fp8-block-example.md @@ -0,0 +1,46 @@ +## Kimi-K3 FP8 Block Example + +### Overview + +This example uses `model_free_ptq` to quantize Kimi-K3 to FP8 block format without loading the full model into memory. +The original checkpoint ships pre-quantized, so a `CompressedTensorsDequantizer` is used to dequantize on the fly during conversion. + +The full example script can be found [here](../../../examples/model_free_ptq/kimi_k3_fp8_block.py). + +### Code Walkthrough + +```python +from compressed_tensors.entrypoints.convert import CompressedTensorsDequantizer + +from llmcompressor import model_free_ptq + +MODEL_ID = "moonshotai/Kimi-K3" +SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-BLOCK" + +# no attention because (q_proj|k_proj|v_proj|b_proj|f_a_proj) are all fused +# and `b_proj` has weight shape [96, 7168] which is not divisible by 128 +ignore = [ + "re:.*embed_tokens.*", + "re:.*self_attn.*", + "re:.*block_sparse_moe\.gate.*", + "re:.*self_attention_res_proj.*", + "re:.*mlp_res_proj.*", + "re:.*output_attn_res_proj.*", + "re:.*lm_head.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*", +] + +model_free_ptq( + model_stub=MODEL_ID, + save_directory=SAVE_DIR, + scheme="FP8_BLOCK", + ignore=ignore, + converter=CompressedTensorsDequantizer( + MODEL_ID, + ignore=ignore, + ), + max_workers=7, + device=[f"cuda:{i}" for i in range(7)], +) +``` diff --git a/docs/key-models/kimi-k3/index.md b/docs/key-models/kimi-k3/index.md index 637e7cd1b6..3a9b5f5d2a 100644 --- a/docs/key-models/kimi-k3/index.md +++ b/docs/key-models/kimi-k3/index.md @@ -5,3 +5,4 @@ Quantization examples for the Kimi K3 model. ## Examples - [NVFP4 Example](nvfp4-example.md) +- [FP8 Block Example](fp8-block-example.md) diff --git a/examples/model_free_ptq/kimi_k3_fp8_block.py b/examples/model_free_ptq/kimi_k3_fp8_block.py new file mode 100644 index 0000000000..38df672bdb --- /dev/null +++ b/examples/model_free_ptq/kimi_k3_fp8_block.py @@ -0,0 +1,33 @@ +from compressed_tensors.entrypoints.convert import CompressedTensorsDequantizer + +from llmcompressor import model_free_ptq + +MODEL_ID = "moonshotai/Kimi-K3" +SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-FP8-BLOCK" + +# no attention because (q_proj|k_proj|v_proj|b_proj|f_a_proj) are all fused +# and `b_proj` has weight shape [96, 7168] which is not divisible by 128 +ignore = [ + "re:.*embed_tokens.*", + "re:.*self_attn.*", + "re:.*block_sparse_moe\.gate.*", + "re:.*self_attention_res_proj.*", + "re:.*mlp_res_proj.*", + "re:.*output_attn_res_proj.*", + "re:.*lm_head.*", + "re:.*vision_tower.*", + "re:.*mm_projector.*", +] + +model_free_ptq( + model_stub=MODEL_ID, + save_directory=SAVE_DIR, + scheme="FP8_BLOCK", + ignore=ignore, + converter=CompressedTensorsDequantizer( + MODEL_ID, + ignore=ignore, + ), + max_workers=7, + device=[f"cuda:{i}" for i in range(7)], +)