Fix for kimi k2 (#593)

* fix for kimi k2

* actually dequant

* use native int4

Author: awni

mlx_lm/benchmark.py

@@ -6,6 +6,7 @@
 
 from mlx_lm import batch_generate, load, stream_generate
 from mlx_lm.generate import DEFAULT_MODEL
+from mlx_lm.utils import pipeline_load
 
 
 def setup_arg_parser():
@@ -56,11 +57,21 @@ def main():
     args = parser.parse_args()
     mx.random.seed(0)
 
+    group = mx.distributed.init()
+    rank = group.rank()
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
     model_path = args.model or DEFAULT_MODEL
 
-    model, tokenizer, config = load(
-        args.model, return_config=True, tokenizer_config={"trust_remote_code": True}
-    )
+    if group.size() > 1:
+        model, tokenizer, config = pipeline_load(args.model, return_config=True)
+    else:
+        model, tokenizer, config = load(
+            args.model, return_config=True, tokenizer_config={"trust_remote_code": True}
+        )
 
     # Empty to avoid early stopping
     tokenizer._eos_token_ids = {}
@@ -89,26 +100,26 @@ def batch_bench():
     else:
         _bench = batch_bench
 
-    print("Running warmup..")
+    rprint("Running warmup..")
     _bench()
 
     report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
-    print(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
+    rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     responses = []
     for i in range(args.num_trials):
         response = _bench()
         responses.append(response)
         results = [(k, getattr(response, k)) for k in report_keys]
         results = [f"{k}={v:.3f}" for k, v in results]
-        print(f"Trial {i+1}:  " + ", ".join(results))
+        rprint(f"Trial {i+1}:  " + ", ".join(results))
 
     def avg(k):
         vals = (getattr(response, k) for response in responses)
         return sum(vals) / args.num_trials
 
     results = [(k, avg(k)) for k in report_keys]
     results = [f"{k}={v:.3f}" for k, v in results]
-    print(f"Averages: " + ", ".join(results))
+    rprint(f"Averages: " + ", ".join(results))
 
 
 if __name__ == "__main__":

mlx_lm/models/deepseek_v3.py

@@ -466,7 +466,22 @@ def dequant(weight, scale_inv):
             )
             return weight[:m, :n].astype(dtype)
 
-        # Dequantize
+        # Remap for int4
+        new_weights = {}
+        for k, v in weights.items():
+            if k.endswith("weight_shape"):
+                base = k.replace("weight_shape", "")
+                new_weights[base + "weight"] = weights[base + "weight_packed"].view(
+                    mx.uint32
+                )
+                s = weights[base + "weight_scale"]
+                new_weights[base + "scales"] = s
+                new_weights[base + "biases"] = -8 * s
+            elif not (k.endswith("weight_scale") or k.endswith("weight_packed")):
+                new_weights[k] = v
+        weights = new_weights
+
+        # Dequantize fp8
         new_weights = {}
         for k, v in weights.items():
             if "weight_scale_inv" in k:

mlx_lm/utils.py

@@ -238,6 +238,11 @@ def class_predicate(p, m):
             config["quantization"] = quantization
             config["quantization_config"] = quantization
             _quantize(quantization)
+        elif quant_method == "compressed-tensors":
+            quantization = {"group_size": 32, "bits": 4, "mode": "affine"}
+            config["quantization"] = quantization
+            config["quantization_config"] = quantization
+            _quantize(quantization)
 
     model.load_weights(list(weights.items()), strict=strict)