Replace recursion with an explicit stack

mlx_lm/models/switch_layers.py

@@ -53,12 +53,6 @@ class QuantizedSwitchLinear(nn.Module):
         # Freeze this model's parameters
         self.freeze()
 
-    def unfreeze(self, *args, **kwargs):
-        """Wrap unfreeze so that we unfreeze any layers we might contain but
-        our parameters will remain frozen."""
-        super().unfreeze(*args, **kwargs)
-        self.freeze(recurse=False)
-
     @property
     def input_dims(self):
         return self.scales.shape[2] * self.group_size

mlx_lm/quant/dwq.py

@@ -13,7 +13,8 @@ from mlx.utils import tree_flatten, tree_map
 from tqdm import tqdm
 
 from mlx_lm.tuner.datasets import load_dataset
-from mlx_lm.tuner.trainer import iterate_batches
+from mlx_lm.tuner.losses import kl_div_loss
+from mlx_lm.tuner.trainer import grad_checkpoint, iterate_batches
 from mlx_lm.tuner.utils import print_trainable_parameters
 from mlx_lm.utils import (
     fetch_from_hub,
@@ -43,6 +44,7 @@ def dwq_quantize(
     activation_layer_step: float = 0.25,
     activation_loss_weight: float = 1.0,
     dtype: mx.Dtype = mx.bfloat16,
+    gradient_checkpoint: bool = False,
 ):
     group = mx.distributed.init()
     world_size = group.size()
@@ -62,22 +64,22 @@ def dwq_quantize(
         model.layers[lid] = Catcher(model.layers[lid])
         q_model.layers[lid] = Catcher(q_model.layers[lid])
 
-    def log_norm(x):
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
+    if gradient_checkpoint:
+        grad_checkpoint(q_model.layers[0])
 
     def forward(model, inputs):
-        logprobs = log_norm(model(inputs).astype(mx.float32))
+        logits = model(inputs)
         extra_targets = [
             model.layers[lid].outputs.astype(mx.float32) for lid in layer_ids
         ]
         for lid in layer_ids:
             model.layers[lid].outputs = None
-        return logprobs, extra_targets
+        return logits, extra_targets
 
     def loss_fn(params, x, targets, extra_targets, lengths):
         q_model.update(tree_map(lambda x: x.astype(dtype), params))
-        logprobs, q_extra_targets = forward(q_model, x)
-        losses = nn.losses.kl_div_loss(logprobs, targets, reduction="none")
+        logits, q_extra_targets = forward(q_model, x)
+        losses = kl_div_loss(logits, targets)
         mask = mx.arange(1, 1 + targets.shape[1]) < lengths[:, 1:]
         ntoks = mask.sum()
         kl_loss = (mask * losses).sum() / ntoks
@@ -194,6 +196,11 @@ def main():
         default="allenai/tulu-3-sft-mixture",
         help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
     )
+    parser.add_argument(
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
+    )
     args = parser.parse_args()
 
     group = mx.distributed.init()
@@ -232,6 +239,7 @@ def main():
         calibration_data,
         batch_size=args.batch_size,
         max_seq_length=args.max_seq_length,
+        gradient_checkpoint=args.grad_checkpoint,
     )
     save(
         args.mlx_path,

mlx_lm/quant/dynamic_quant.py

@@ -8,10 +8,13 @@ import math
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
-from mlx.utils import tree_flatten, tree_map, tree_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
 from tqdm import tqdm
 
 from mlx_lm.quant.utils import load_data
+from mlx_lm.tuner.losses import kl_div_loss
+from mlx_lm.tuner.trainer import grad_checkpoint
+from mlx_lm.tuner.utils import get_total_parameters
 from mlx_lm.utils import (
     compute_bits_per_weight,
     fetch_from_hub,
@@ -22,6 +25,15 @@ from mlx_lm.utils import (
 )
 
 
+def make_quant_predicate(config):
+    def quant_predicate(p, m, _):
+        if not hasattr(m, "to_quantized"):
+            return False
+        return config.get(p, True)
+
+    return quant_predicate
+
+
 def eval_ppl(model, data, batch_size=8):
     all_loss = 0.0
     ntoks = 0
@@ -35,6 +47,26 @@ def eval_ppl(model, data, batch_size=8):
     return ppl
 
 
+def make_options(
+    low_bits, low_group_size, high_bits, high_group_size, include_bpw=True
+):
+    options = []
+    min_bpw = low_bits + 32 / low_group_size
+    max_bpw = high_bits + 32 / high_group_size
+    for b in range(low_bits, high_bits + 1):
+        for g in [32, 64, 128]:
+            cbpw = b + 32 / g
+            if b == 7 or not (min_bpw <= cbpw <= max_bpw):
+                continue
+            options.append({"bits": b, "group_size": g, "bpw": cbpw})
+    options.sort(key=lambda x: x["bpw"])
+    if not include_bpw:
+        for o in options:
+            o.pop("bpw")
+
+    return options
+
+
 def estimate_sensitivities(
     model,
     data,
@@ -42,17 +74,17 @@ def estimate_sensitivities(
     low_group_size,
     high_bits,
     high_group_size,
+    batch_size: int = 4,
+    gradient_accum_dtype: mx.Dtype = mx.float32,
+    gradient_checkpoint: bool = False,
 ):
-    batch_size = 4
-    layers = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module)
-    layers = {k: l for k, l in layers if hasattr(l, "to_quantized")}
-
-    q_model = copy.deepcopy(model)
-
     def qdq(w, bits, group_size):
         w, s, b = mx.quantize(w, bits=bits, group_size=group_size)
         return mx.dequantize(w, scales=s, biases=b, bits=bits, group_size=group_size)
 
+    layers = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module)
+    layers = {k: l for k, l in layers if hasattr(l, "to_quantized")}
+    q_model = copy.deepcopy(model)
     q_layers = copy.deepcopy(layers)
     for l in q_layers.values():
         l.weight = qdq(l.weight, low_bits, low_group_size)
@@ -62,34 +94,42 @@ def estimate_sensitivities(
     q_model.freeze()
     q_model.update_modules(tree_unflatten(list(q_layers.items())))
 
-    def log_norm(x):
-        x = x.astype(mx.float32)
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
-
     def loss_fn(batch, targets):
-        logprobs = log_norm(q_model(batch))
-        return nn.losses.kl_div_loss(logprobs, targets, reduction="mean")
+        return kl_div_loss(q_model(batch), targets).mean()
 
-    grad_accum = tree_map(lambda x: mx.zeros(x.shape), q_model.trainable_parameters())
+    if gradient_checkpoint:
+        grad_checkpoint(q_model.layers[0])
+
+    grad_accum = tree_map(
+        lambda x: mx.zeros(x.shape, dtype=gradient_accum_dtype),
+        q_model.trainable_parameters(),
+    )
     for e, s in tqdm(
         enumerate(range(0, len(data), batch_size)),
         total=len(data) // batch_size,
         desc="Estimating sensitivities",
     ):
         batch = data[s : s + batch_size]
-        targets = log_norm(model(batch))
+        targets = model(batch)
         mx.eval(targets)
         _, grads = nn.value_and_grad(q_model, loss_fn)(batch, targets)
         grad_accum = tree_map(lambda x, y: x + y, grad_accum, grads)
+        del grads
         mx.eval(grad_accum)
 
+    options = make_options(low_bits, low_group_size, high_bits, high_group_size)
+    current_bpw = options[0]["bpw"]
+
     def compute_sensitivity(gradient, low_q_weight, original_weight):
         n_batches = (len(data) + batch_size - 1) // batch_size
         gradient = gradient / n_batches
-        high_q_weight = qdq(original_weight, high_bits, high_group_size)
-        param_size = original_weight.size / 1e6
-        alignment = (gradient * (low_q_weight - high_q_weight)).sum()
-        return alignment / param_size
+        scores = [{"loss_change": 0, "extra_bits": 0}]
+        for opt in options[1:]:
+            extra_bits = (opt["bpw"] - current_bpw) * original_weight.size
+            other_weight = qdq(original_weight, opt["bits"], opt["group_size"])
+            loss_change = (gradient * (low_q_weight - other_weight)).sum()
+            scores.append({"loss_change": loss_change, "extra_bits": extra_bits})
+        return scores
 
     sensitivities = tree_map(
         compute_sensitivity,
@@ -99,11 +139,23 @@ def estimate_sensitivities(
     )
     mx.eval(sensitivities)
 
-    sensitivities = [(k[:-7], s.item()) for k, s in tree_flatten(sensitivities)]
+    sensitivities = [
+        (k.replace(".weight", ""), s.item() if isinstance(s, mx.array) else s)
+        for k, s in tree_flatten(sensitivities)
+    ]
 
     return sensitivities
 
 
+def compute_bit_budget(model, target_bpw):
+    model_bytes = tree_reduce(
+        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
+    )
+    model_params = get_total_parameters(model)
+
+    return model_params * target_bpw - model_bytes * 8
+
+
 def estimate_threshold(
     model,
     sensitivities,
@@ -113,35 +165,79 @@ def estimate_threshold(
     high_bits,
     high_group_size,
 ):
-    def predicate(p, m, high_threshold):
-        if not hasattr(m, "to_quantized"):
-            return False
-        if sensitivities[p] > high_threshold:
-            return {"bits": high_bits, "group_size": high_group_size}
-        return True
+    options = make_options(
+        low_bits, low_group_size, high_bits, high_group_size, include_bpw=False
+    )
+    sensitivities = tree_flatten(
+        tree_unflatten(list(sensitivities.items())),
+        is_leaf=lambda x: isinstance(x, list) and "loss_change" in x[0],
+    )
 
-    # Binary search for the threshold
-    sens_vals = list(sensitivities.values())
-    min_threshold = min(sens_vals)
-    max_threshold = max(sens_vals)
-    tolerance = 1e-3 * (max_threshold - min_threshold)
-    while (max_threshold - min_threshold) > tolerance:
-        mid = (max_threshold + min_threshold) / 2
-        class_predicate = lambda p, m: predicate(p, m, mid)
-        q_model = copy.deepcopy(model)
-        nn.quantize(
-            q_model,
-            group_size=low_group_size,
-            bits=low_bits,
-            class_predicate=class_predicate,
-        )
-        bpw = compute_bits_per_weight(q_model)
-        if bpw > target_bpw:
-            min_threshold = mid
-        else:
-            max_threshold = mid
+    q_model = copy.deepcopy(model)
+    nn.quantize(q_model, group_size=low_group_size, bits=low_bits)
+    budget = int(compute_bit_budget(q_model, target_bpw))
+    benefit_map = {}
 
-    return (max_threshold + min_threshold) / 2
+    def benefit(layer, option, budget):
+        if (layer, option, budget) in benefit_map:
+            return benefit_map[layer, option, budget]
+
+        stack = [(layer, option, budget)]
+        while stack:
+            layer, option, budget = stack[-1]
+
+            if budget <= 0 or layer < 0 or option < 0:
+                benefit_map[layer, option, budget] = 0
+                stack.pop()
+                continue
+
+            # We either not use this option
+            prev_layer = layer if option > 0 else layer - 1
+            prev_option = (option if option > 0 else len(options)) - 1
+            if (prev_layer, prev_option, budget) not in benefit_map:
+                stack.append((prev_layer, prev_option, budget))
+                continue
+            a = benefit_map[prev_layer, prev_option, budget]
+
+            # Or we use it so we have less budget for before
+            b = float("-inf")
+            info = sensitivities[layer][1][option]
+            prev_layer = layer - 1
+            prev_option = len(options) - 1
+            prev_budget = budget - info["extra_bits"]
+            if (
+                prev_layer,
+                prev_option,
+                prev_budget,
+            ) not in benefit_map and prev_budget >= 0:
+                stack.append((prev_layer, prev_option, prev_budget))
+                continue
+            if prev_budget >= 0:
+                b = benefit_map[prev_layer, prev_option, prev_budget]
+                b += info["loss_change"]
+
+            benefit_map[layer, option, budget] = max(a, b)
+            stack.pop()
+
+        return benefit_map[layer, option, budget]
+
+    def backtrack(layer, budget):
+        selected = []
+        while layer >= 0:
+            prev_benefit = benefit(layer - 1, len(options) - 1, budget)
+            option_benefits = [benefit(layer, i, budget) for i in range(len(options))]
+            idx, v = max(enumerate(option_benefits), key=lambda x: x[1] - prev_benefit)
+            info = sensitivities[layer][1][idx]
+            if v != 0:
+                budget -= info["extra_bits"]
+                selected.append((layer, idx))
+            layer -= 1
+        return selected[::-1]
+
+    selected = backtrack(len(sensitivities) - 1, budget)
+    config = {sensitivities[l][0]: options[i] for l, i in selected}
+
+    return config
 
 
 def main():
@@ -161,15 +257,25 @@ def main():
         "--target-bpw", type=float, default=5.0, help="Target bits per weight."
     )
     parser.add_argument("--low-bits", type=int, default=4)
-    parser.add_argument("--low-group-size", type=int, default=64)
+    parser.add_argument("--low-group-size", type=int, default=128)
     parser.add_argument("--high-bits", type=int, default=5)
-    parser.add_argument("--high-group-size", type=int, default=64)
+    parser.add_argument("--high-group-size", type=int, default=32)
     parser.add_argument(
         "--report-ppl",
         action="store_true",
         help="Compute the perplexity of the base and quantized models.",
     )
-
+    parser.add_argument(
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
+    )
+    parser.add_argument(
+        "--accumulation-dtype",
+        default="float32",
+        choices=["float32", "bfloat16"],
+        help="What type to use to accumulate the gradients for the sensitivities",
+    )
     args = parser.parse_args()
 
     group = mx.distributed.init()
@@ -186,6 +292,8 @@ def main():
             args.low_group_size,
             args.high_bits,
             args.high_group_size,
+            gradient_accum_dtype=getattr(mx, args.accumulation_dtype),
+            gradient_checkpoint=args.grad_checkpoint,
         )
         model_name = args.model.replace("/", "_")
         with open(f"{model_name}_sensitivities.json", "w") as fid:
@@ -204,7 +312,7 @@ def main():
         ppl = eval_ppl(model, data)
         print(f"Original PPL: {ppl:.3f}")
 
-    threshold = estimate_threshold(
+    quant_config = estimate_threshold(
         model,
         sensitivities,
         target_bpw=args.target_bpw,
@@ -214,19 +322,12 @@ def main():
         high_group_size=args.high_group_size,
     )
 
-    def quant_predicate(p, m, _):
-        if not hasattr(m, "to_quantized"):
-            return False
-        if sensitivities[p] > threshold:
-            return {"bits": args.high_bits, "group_size": args.high_group_size}
-        return True
-
     model, config = quantize_model(
         model,
         config,
         q_group_size=args.low_group_size,
         q_bits=args.low_bits,
-        quant_predicate=quant_predicate,
+        quant_predicate=make_quant_predicate(quant_config),
     )
 
     if args.report_ppl:
@@ -241,6 +342,7 @@ def main():
         config,
         hf_repo=hf_repo,
     )
+    print(f"Peak memory used: {mx.get_peak_memory() / 1000**3:.3f}GB")
 
 
 if __name__ == "__main__":

mlx_lm/tuner/losses.py

@@ -0,0 +1,378 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def _make_kl_forward_kernel():
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx;
+
+    float lse_q_minus_p;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q_minus_p = max_q + metal::fast::log(sum_exp_q) - lse_p;
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                kl += metal::fast::exp(vals_p[j] - lse_p) * (vals_p[j] - vals_q[j] + lse_q_minus_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                kl += metal::fast::exp(vals_p[j] - lse_p) * (vals_p[j] - vals_q[j] + lse_q_minus_p);
+            }
+        }
+
+        // Add the kl across the threadgroup
+        kl = simd_sum(kl);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id] = kl;
+        }
+        threadgroup_barrier(mem_flags::mem_none);
+        kl = shared[simd_lane_id];
+        kl = simd_sum(kl);
+
+        if (thread_index_in_threadgroup == 0) {
+            out[0] = static_cast<T>(kl);
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="kl_forward",
+        input_names=["logits_q", "logits_p"],
+        output_names=["out"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+def _make_kl_backward_kernel():
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx * V;
+    cotan += out_idx;
+
+    float lse_q;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q = max_q + metal::fast::log(sum_exp_q);
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl = 0;
+        float c = cotan[0];
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                out[offset + j] = static_cast<T>(
+                    c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p)));
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                if (offset + j < V) {
+                    out[offset + j] = static_cast<T>(
+                        c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p)));
+                }
+            }
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="kl_backward",
+        input_names=["logits_q", "logits_p", "cotan"],
+        output_names=["out"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+_kl_forward_kernel = _make_kl_forward_kernel()
+_kl_backward_kernel = _make_kl_backward_kernel()
+
+
+@mx.custom_function
+def _kl_div_loss(logits_q, logits_p):
+    n_outs = logits_q.size // logits_q.shape[-1]
+    dt = logits_q.dtype
+
+    return _kl_forward_kernel(
+        inputs=[logits_q, logits_p],
+        output_shapes=[logits_q.shape[:-1]],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, n_outs, 1),
+        threadgroup=(1024, 1, 1),
+    )[0]
+
+
+@_kl_div_loss.vjp
+def _kl_div_loss(primals, cotangent, output):
+    logits_q, logits_p = primals
+    dt = logits_q.dtype
+
+    dp = mx.zeros_like(logits_p)
+    dq = _kl_backward_kernel(
+        inputs=[logits_q, logits_p, cotangent],
+        output_shapes=[logits_q.shape],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, cotangent.size, 1),
+        threadgroup=(1024, 1, 1),
+    )[0]
+
+    return dq, dp
+
+
+def kl_div_loss(logits_q, logits_p):
+    if mx.metal.is_available():
+        return _kl_div_loss(logits_q, logits_p)
+    else:
+        return nn.losses.kl_div_loss(
+            logits_q - mx.logsumexp(logits_q, axis=-1, keepdims=True),
+            logits_p - mx.logsumexp(logits_p, axis=-1, keepdims=True),
+            axis=-1,
+            reduction="none",
+        )