feat: add Kimi-K3 model definition and quantization example

docs/.nav.yml

@@ -20,6 +20,10 @@ 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
+      - 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

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.
 
 <div class="grid cards" markdown>
 
@@ -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**
 
       ---

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)],
+)
+```

docs/key-models/kimi-k3/index.md

@@ -0,0 +1,8 @@
+# Kimi K3
+
+Quantization examples for the Kimi K3 model.
+
+## Examples
+
+- [NVFP4 Example](nvfp4-example.md)
+- [FP8 Block Example](fp8-block-example.md)

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)
+```

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)],
+)

examples/quantizing_moe/kimi_k3_example.py

@@ -0,0 +1,87 @@
+from compressed_tensors.quantization import QuantizationConfig
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+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)

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"]

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 "

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",

src/llmcompressor/modeling/kimi_k3/__init__.py

@@ -0,0 +1 @@
+from .modeling_kimi_k3 import KimiK3ForConditionalGeneration