Support KIMI K2 Thinking int4 checkpoint PTQ

CHANGELOG.rst

@@ -11,6 +11,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 **New Features**
 
 - Add standalone type inference option (``--use_standalone_type_inference``) in ONNX AutoCast as an alternative to ONNX's ``infer_shapes``. This experimental feature performs type-only inference without shape inference, useful as a workaround when shape inference fails or to avoid unnecessary shape inference overhead.
+- Add support for Kimi K2 Thinking model quantization from the original int4 checkpoint.
 
 0.41 (2026-01-19)
 ^^^^^^^^^^^^^^^^^

examples/llm_ptq/README.md

@@ -109,6 +109,7 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http
 | QWen3 MOE, Next <sup>6</sup> | ✅ | - | - | - | ✅ |
 | QwQ | ✅ | - | - | - | ✅ |
 | DeepSeek V3, R1, V3.1, V3.2<sup>7</sup> | - | - | - | - | ✅ |
+| Kimi K2 | - | - | - | - | ✅ |
 | T5 | ✅ | ✅ | ✅ | ✅ | - |
 | Whisper | ✅ | ❌ | ❌ | ❌ | - |
 

examples/llm_ptq/example_utils.py

@@ -187,6 +187,11 @@ def build_quant_cfg(
         quant_cfg["quant_cfg"]["model*.*attn*k_proj*"] = {"enable": False}
         quant_cfg["quant_cfg"]["model*.*attn*v_proj*"] = {"enable": False}
 
+    if model_type == "deepseek":
+        # Disable MLA quantization for accuracy.
+        quant_cfg["quant_cfg"]["*self_attn.q*"] = {"enable": False}
+        quant_cfg["quant_cfg"]["*self_attn.kv*"] = {"enable": False}
+
     return quant_cfg
 
 
@@ -346,6 +351,17 @@ def get_model(
                 device_map=device_map,
                 **model_kwargs,
             )
+        elif (
+            hasattr(hf_config, "quantization_config")
+            and hf_config.quantization_config.get("format", None) == "pack-quantized"
+        ):
+            torch_dtype = getattr(hf_config, "torch_dtype", torch.bfloat16)
+            model = AutoModelForCausalLM.from_pretrained(
+                ckpt_path,
+                device_map="auto",
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch_dtype,
+            )
         else:
             architecture = hf_config.architectures[0]
 
@@ -366,9 +382,9 @@ def get_model(
                 from_config = auto_model_module._from_config
 
             with init_empty_weights():
-                # When computing the device_map, assuming half precision by default,
+                # When computing the device_map, assuming bfloat16 precision by default,
                 # unless specified by the hf_config.
-                torch_dtype = getattr(hf_config, "torch_dtype", torch.float16)
+                torch_dtype = getattr(hf_config, "torch_dtype", torch.bfloat16)
                 model_kwargs2 = model_kwargs.copy()
                 if auto_model_module != AutoModelForCausalLM:
                     model_kwargs2.pop("trust_remote_code", None)

examples/llm_ptq/hf_ptq.py

@@ -575,7 +575,10 @@ def pre_quantize(
     ][0:1]
 
     # Generate preview before quantization
-    if is_nemotron_vl_model and tokenizer is not None:
+    if model_type == "deepseek":
+        # DeepSeek generation may go OOM, so we skip it
+        generated_ids_before_ptq = None
+    elif is_nemotron_vl_model and tokenizer is not None:
         generated_ids_before_ptq = run_nemotron_vl_preview(
             full_model,
             tokenizer,
@@ -618,7 +621,9 @@ def post_quantize(
     # Run some samples
     torch.cuda.empty_cache()
     generated_ids_after_ptq = None
-    if model_type != "llama4" and not is_nemotron_vl_model:
+    if generated_ids_before_ptq is None:
+        pass
+    elif model_type != "llama4" and not is_nemotron_vl_model:
         # Our fake quantizer may not be fully compatible with torch.compile.
         generated_ids_after_ptq = full_model.generate(preview_input_ids, max_new_tokens=100)
     elif is_nemotron_vl_model and tokenizer is not None:

examples/llm_ptq/requirements.txt

@@ -1,3 +1,4 @@
+compressed-tensors==0.12.0
 fire
 flash-attn>=2.6.0
 rouge_score>=0.1.2

examples/llm_ptq/scripts/huggingface_example.sh

@@ -53,9 +53,9 @@ esac
 IFS=","
 for qformat in $QFORMAT; do
     case $qformat in
-    fp8 | fp8_pc_pt | fp8_pb_wo | int8_wo | int8_sq | int4_awq | w4a8_awq | fp16 | bf16 | nvfp4 | nvfp4_awq | w4a8_nvfp4_fp8 | w4a8_mxfp4_fp8) ;;
+    fp8 | fp8_pc_pt | fp8_pb_wo | int8_wo | int8_sq | int4_awq | w4a8_awq | fp16 | bf16 | nvfp4 | nvfp4_awq | w4a8_nvfp4_fp8 | w4a8_mxfp4_fp8 | nvfp4_mlp_only) ;;
     *)
-        echo "Unknown quant argument: Expected one of: [fp8, fp8_pc_pt, fp8_pb_wo, int8_wo, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8]" >&2
+        echo "Unknown quant argument: Expected one of: [fp8, fp8_pc_pt, fp8_pb_wo, int8_wo, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, nvfp4_mlp_only]" >&2
         exit 1
         ;;
     esac

modelopt/torch/export/layer_utils.py

@@ -345,7 +345,8 @@ def is_moe(module: nn.Module) -> bool:
 
 def is_quantlinear(module: nn.Module) -> bool:
     """Returns whether the module is a quantized linear layer."""
-    return "QuantLinear" in type(module).__name__ and "lora" not in type(module).__name__.lower()
+    name = type(module).__name__
+    return ("QuantLinear" in name or "QuantCompressedLinear" in name) and "lora" not in name.lower()
 
 
 def dup_kv_weight(v: torch.Tensor, head_size: int, num_head: int, tp_size: int) -> torch.Tensor:

modelopt/torch/export/quant_utils.py

@@ -881,7 +881,13 @@ def postprocess_state_dict(
         "v_bmm_quantizer._bias_value": "v_proj.v_bias",
         "input_quantizer._pre_quant_scale": "pre_quant_scale",
     }
-    skip_keys = ["output_quantizer", "_amax", "_bias_value", "input_quantizer._pre_quant_scale"]
+    skip_keys = [
+        "output_quantizer",
+        "_amax",
+        "_bias_value",
+        "input_quantizer._pre_quant_scale",
+        "weight_shape",
+    ]
 
     # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment
     if is_modelopt_qlora:

modelopt/torch/export/unified_export_hf.py

@@ -490,6 +490,8 @@ def _export_quantized_weight(
     if weight_scale is not None:
         sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale)
 
+    torch.cuda.empty_cache()
+
 
 def _process_quantized_modules(
     model: nn.Module,
@@ -525,6 +527,8 @@ def _process_quantized_modules(
         if is_modelopt_qlora and (hasattr(sub_module, "base_layer")):
             continue
 
+        if hasattr(sub_module, "weight_packed"):
+            sub_module.unpack_weight()
         if get_quantization_format(sub_module) != QUANTIZATION_NONE:
             if is_quantlinear(sub_module):
                 with fsdp2_aware_weight_update(model, sub_module, reshard=False):
@@ -742,6 +746,10 @@ def export_hf_checkpoint(
 
             hf_quant_config = convert_hf_quant_config_format(hf_quant_config)
 
+        # Remove hf_quantizer from model so post_state_dict can be exported.
+        if getattr(model, "hf_quantizer", None) is not None:
+            model.hf_quantizer = None
+
         # Save model
         model.save_pretrained(
             export_dir, state_dict=post_state_dict, save_modelopt_state=save_modelopt_state

modelopt/torch/quantization/plugins/huggingface.py

@@ -22,6 +22,8 @@
 from typing import TYPE_CHECKING
 
 import torch
+from torch import Tensor
+from torch.nn.functional import linear
 
 try:
     from torch.distributed.tensor import Shard
@@ -665,6 +667,32 @@ def top_k(self, value):
         self.router.moe_top_k = value
 
 
+class _QuantCompressedLinear(QuantModule):
+    def _setup(self):
+        self.input_quantizer = TensorQuantizer()
+        self.weight_quantizer = TensorQuantizer()
+
+    def forward(self, input: Tensor) -> Tensor:
+        from compressed_tensors.quantization import QuantizationStatus
+
+        if self.quantization_status == QuantizationStatus.COMPRESSED:
+            weight_data = self.compressor.decompress_module(self)
+        else:
+            weight_data = self.weight
+
+        return linear(self.input_quantizer(input), self.weight_quantizer(weight_data), self.bias)
+
+    def unpack_weight(self):
+        from compressed_tensors.quantization import QuantizationStatus
+
+        if self.quantization_status == QuantizationStatus.COMPRESSED:
+            self.weight = nn.Parameter(self.compressor.decompress_module(self), requires_grad=False)
+        if hasattr(self, "weight_packed"):
+            del self.weight_packed
+        if hasattr(self, "weight_scale"):
+            del self.weight_scale
+
+
 try:
     from transformers.models.llama4.modeling_llama4 import Llama4TextExperts, Llama4TextMoe
 
@@ -740,6 +768,16 @@ def top_k(self, value):
 except ImportError:
     pass
 
+try:
+    from compressed_tensors.linear.compressed_linear import CompressedLinear
+
+    if CompressedLinear not in QuantModuleRegistry:
+        QuantModuleRegistry.register({CompressedLinear: "hf.CompressedLinear"})(
+            _QuantCompressedLinear
+        )
+except ImportError:
+    pass
+
 try:
     from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import (
         Qwen3VLMoeTextExperts,