Rebase the fork onto v0.32.4; qmv_quad ragged-tile scale/bias read (gemma4-qat#302)

docs/src/python/fast.rst

@@ -13,6 +13,7 @@ Fast
   cross_entropy
   rope
   scaled_dot_product_attention
+  spec_decode_verify
   metal_kernel
   cuda_kernel
   precompiled_cuda_kernel

mlx/backend/common/metal_kernel.cpp

@@ -1,7 +1,9 @@
 // Copyright © 2024 Apple Inc.
 
 #include <iostream>
+#include <mutex>
 #include <sstream>
+#include <unordered_map>
 
 #include "mlx/backend/common/compiled.h"
 #include "mlx/backend/metal/metal.h"
@@ -217,6 +219,57 @@ std::string make_template_hash(const std::string& template_def) {
   return template_hash;
 }
 
+// The name, source and source hash of one kernel variant.
+struct KernelVariant {
+  std::string name;
+  std::shared_ptr<const std::string> source;
+  size_t source_hash;
+};
+
+// One per kernel function, cleared when full. Callers own their entry, so a
+// clear is safe at any time.
+struct VariantCache {
+  std::mutex mutex;
+  std::unordered_map<std::string, std::shared_ptr<const KernelVariant>>
+      variants;
+};
+
+// Encodes each argument that the name and source depend on. Template arguments
+// add their kind, because `int` 1 and `bool` true give the same name.
+std::string variant_key(
+    const std::vector<array>& inputs,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args) {
+  std::string key;
+  auto append_bytes = [&key](auto v) {
+    key.append(reinterpret_cast<const char*>(&v), sizeof(v));
+  };
+  for (const auto& arr : inputs) {
+    key += static_cast<char>(arr.dtype().val());
+    key += arr.ndim() == 0                     ? 's'
+        : arr.size() < max_constant_array_size ? 'c'
+                                               : 'd';
+  }
+  for (const auto& dtype : output_dtypes) {
+    key += static_cast<char>(dtype.val());
+  }
+  for (const auto& [name, arg] : template_args) {
+    key += static_cast<char>(arg.index());
+    append_bytes(name.size());
+    key += name;
+    std::visit(
+        [&](auto v) {
+          if constexpr (std::is_same_v<decltype(v), Dtype>) {
+            key += static_cast<char>(v.val());
+          } else {
+            append_bytes(v);
+          }
+        },
+        arg);
+  }
+  return key;
+}
+
 } // namespace
 
 CustomKernelFunction metal_kernel(
@@ -270,7 +323,10 @@ CustomKernelFunction metal_kernel(
     }
   }
 
+  auto cache = std::make_shared<VariantCache>();
+
   return [=,
+          cache = std::move(cache),
           shape_infos = std::move(shape_infos),
           attributes = std::move(attributes)](
              const std::vector<array>& inputs,
@@ -298,60 +354,80 @@ CustomKernelFunction metal_kernel(
 
     auto s = resolve_metal_kernel_stream(s_);
 
-    std::string kernel_name = "custom_kernel_" + name;
-    std::string template_def = "";
-    if (!template_args.empty()) {
-      template_def = write_template(template_args);
-      auto template_hash = make_template_hash(template_def);
-      kernel_name += "_";
-      kernel_name += template_hash;
+    auto key = variant_key(inputs, output_dtypes, template_args);
+    std::shared_ptr<const KernelVariant> variant;
+    // Build a missing variant under the lock, so concurrent first callers
+    // build it once.
+    std::unique_lock lock(cache->mutex);
+    if (auto it = cache->variants.find(key); it != cache->variants.end()) {
+      variant = it->second;
     }
+    if (!variant) {
+      std::string kernel_name = "custom_kernel_" + name;
+      std::string template_def = "";
+      if (!template_args.empty()) {
+        template_def = write_template(template_args);
+        auto template_hash = make_template_hash(template_def);
+        kernel_name += "_";
+        kernel_name += template_hash;
+      }
 
-    // The generated source depends on the dtypes of the inputs and outputs
-    // and on how each input is passed (see `write_signature`). Include them
-    // in the kernel name so that a given name always maps to the same source.
-    for (const auto& arr : inputs) {
-      kernel_name += "_";
-      kernel_name += get_type_string(arr.dtype());
-      if (arr.ndim() == 0) {
-        kernel_name += "s";
-      } else if (arr.size() < max_constant_array_size) {
-        kernel_name += "c";
+      // The source depends on the input and output dtypes and on how each
+      // input is passed, so the name includes them.
+      for (const auto& arr : inputs) {
+        kernel_name += "_";
+        kernel_name += get_type_string(arr.dtype());
+        if (arr.ndim() == 0) {
+          kernel_name += "s";
+        } else if (arr.size() < max_constant_array_size) {
+          kernel_name += "c";
+        }
+      }
+      for (const auto& dtype : output_dtypes) {
+        kernel_name += "_";
+        kernel_name += get_type_string(dtype);
       }
-    }
-    for (const auto& dtype : output_dtypes) {
-      kernel_name += "_";
-      kernel_name += get_type_string(dtype);
-    }
 
-    std::string kernel_source = write_signature(
-        kernel_name,
-        header,
-        source,
-        input_names,
-        inputs,
-        output_names,
-        output_dtypes,
-        template_args,
-        attributes,
-        shape_infos,
-        atomic_outputs);
+      std::string kernel_source = write_signature(
+          kernel_name,
+          header,
+          source,
+          input_names,
+          inputs,
+          output_names,
+          output_dtypes,
+          template_args,
+          attributes,
+          shape_infos,
+          atomic_outputs);
 
-    if (!template_args.empty()) {
-      template_def = kernel_name + template_def;
-      kernel_source += "\ntemplate [[host_name(\"";
-      kernel_source += kernel_name;
-      kernel_source += "\")]] [[kernel]] decltype(";
-      kernel_source += template_def;
-      kernel_source += ") ";
-      kernel_source += template_def;
-      kernel_source += ";\n";
+      if (!template_args.empty()) {
+        template_def = kernel_name + template_def;
+        kernel_source += "\ntemplate [[host_name(\"";
+        kernel_source += kernel_name;
+        kernel_source += "\")]] [[kernel]] decltype(";
+        kernel_source += template_def;
+        kernel_source += ") ";
+        kernel_source += template_def;
+        kernel_source += ";\n";
+      }
+
+      auto source_hash = std::hash<std::string>{}(kernel_source);
+      variant = std::make_shared<const KernelVariant>(KernelVariant{
+          std::move(kernel_name),
+          std::make_shared<const std::string>(std::move(kernel_source)),
+          source_hash});
+      if (cache->variants.size() >= metal_kernel_max_cached_variants) {
+        cache->variants.clear();
+      }
+      cache->variants.emplace(std::move(key), variant);
     }
+    lock.unlock();
 
     if (verbose) {
       std::cout << "Generated source code for `" << name << "`:" << std::endl
                 << "```" << std::endl
-                << kernel_source << std::endl
+                << *variant->source << std::endl
                 << "```" << std::endl;
     }
 
@@ -360,8 +436,8 @@ CustomKernelFunction metal_kernel(
         std::move(output_dtypes),
         std::make_shared<CustomKernel>(
             s,
-            std::move(kernel_name),
-            std::move(kernel_source),
+            variant->name,
+            variant->source,
             grid,
             threadgroup,
             shape_infos,
@@ -370,7 +446,8 @@ CustomKernelFunction metal_kernel(
             std::vector<ScalarArg>{},
             false,
             0,
-            compile_options.serialize()),
+            compile_options.serialize(),
+            variant->source_hash),
         std::move(inputs));
   };
 }

mlx/backend/common/quantized.h

@@ -1,5 +1,11 @@
 // Copyright © 2026 Apple Inc.
 
+#pragma once
+
+#include <optional>
+
+#include "mlx/array.h"
+
 namespace mlx::core {
 
 inline constexpr short get_pack_factor(int bits, int wsize = 8) {
@@ -11,4 +17,14 @@ inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
   return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
 }
 
+// Implied bias (affine mode): a 0-d `biases` array holds a single factor f
+// standing for the per-group bias `scales * T(f)` in the scales' dtype T.
+inline bool is_implied_bias(const array& biases) {
+  return biases.ndim() == 0;
+}
+
+inline bool is_implied_bias(const std::optional<array>& biases) {
+  return biases.has_value() && is_implied_bias(*biases);
+}
+
 } // namespace mlx::core

mlx/backend/cpu/CMakeLists.txt

@@ -64,6 +64,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/spec_decode.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/luf.cpp

mlx/backend/cpu/quantized.cpp

@@ -943,8 +943,72 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
   encoder.set_input_array(scales);
   encoder.set_output_array(out);
   if (mode_ == QuantizationMode::Affine) {
+    if (inputs.size() < 4) {
+      throw std::runtime_error(
+          "[QuantizedMatmul::eval_cpu] Bias-free affine is Metal-only for "
+          "now; on CPU pass biases as the 0-d implied-bias factor (-0.5 for "
+          "1-bit, -1 for 2-bit).");
+    }
     auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
     encoder.set_input_array(biases);
+    if (is_implied_bias(biases)) {
+      // Implied bias: expand the factor f into a temporary `scales * T(f)`
+      // for the per-group-bias kernels.
+      array full_biases(scales.shape(), scales.dtype(), nullptr, {});
+      full_biases.set_data(allocator::malloc(full_biases.nbytes()));
+      encoder.add_temporary(full_biases);
+      encoder.set_output_array(full_biases);
+      encoder.dispatch([out = array::unsafe_weak_copy(out),
+                        x = array::unsafe_weak_copy(x),
+                        w = array::unsafe_weak_copy(w),
+                        scales = array::unsafe_weak_copy(scales),
+                        factor = array::unsafe_weak_copy(biases),
+                        full_biases = array::unsafe_weak_copy(full_biases),
+                        group_size_ = group_size_,
+                        bits_ = bits_,
+                        transpose_ = transpose_]() mutable {
+        float f;
+        switch (factor.dtype()) {
+          case float32:
+            f = factor.data<float>()[0];
+            break;
+          case float16:
+            f = static_cast<float>(factor.data<float16_t>()[0]);
+            break;
+          case bfloat16:
+            f = static_cast<float>(factor.data<bfloat16_t>()[0]);
+            break;
+          default:
+            throw std::runtime_error(
+                "[QuantizedMatmul::eval_cpu] Implied-bias factor must be a "
+                "real floating type.");
+        }
+        auto fill = [&](auto* dst, const auto* src) {
+          using T = std::remove_pointer_t<decltype(dst)>;
+          for (size_t i = 0; i < scales.size(); ++i) {
+            dst[i] = static_cast<T>(f * static_cast<float>(src[i]));
+          }
+        };
+        switch (scales.dtype()) {
+          case float32:
+            fill(full_biases.data<float>(), scales.data<float>());
+            break;
+          case float16:
+            fill(full_biases.data<float16_t>(), scales.data<float16_t>());
+            break;
+          case bfloat16:
+            fill(full_biases.data<bfloat16_t>(), scales.data<bfloat16_t>());
+            break;
+          default:
+            throw std::runtime_error(
+                "[QuantizedMatmul::eval_cpu] Only real floating scales are "
+                "supported with an implied bias.");
+        }
+        _qmm_dispatch(
+            out, x, w, scales, full_biases, group_size_, bits_, transpose_);
+      });
+      return;
+    }
     encoder.dispatch([out = array::unsafe_weak_copy(out),
                       x = array::unsafe_weak_copy(x),
                       w = array::unsafe_weak_copy(w),

mlx/backend/cpu/spec_decode.cpp

@@ -0,0 +1,58 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/allocator.h"
+#include "mlx/backend/cpu/encoder.h"
+#include "mlx/fast_primitives.h"
+
+namespace mlx::core::fast {
+
+void SpecDecodeVerify::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  auto& draft = inputs[0]; // [B, K] int32
+  auto& target = inputs[1]; // [B, K+1] int32
+  auto& n_accepted = outputs[0]; // [B] int32
+  auto& committed = outputs[1]; // [B, K+1] int32
+  n_accepted.set_data(allocator::malloc(n_accepted.nbytes()));
+  committed.set_data(allocator::malloc(committed.nbytes()));
+  if (n_accepted.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(draft);
+  encoder.set_input_array(target);
+  encoder.set_output_array(n_accepted);
+  encoder.set_output_array(committed);
+  encoder.dispatch([draft = array::unsafe_weak_copy(draft),
+                    target = array::unsafe_weak_copy(target),
+                    n_accepted = array::unsafe_weak_copy(n_accepted),
+                    committed = array::unsafe_weak_copy(committed)]() mutable {
+    int B = draft.shape(0);
+    int K = draft.shape(1);
+    auto ds = draft.strides();
+    auto ts = target.strides();
+    auto* d = draft.data<int32_t>();
+    auto* t = target.data<int32_t>();
+    auto* n_out = n_accepted.data<int32_t>();
+    auto* c_out = committed.data<int32_t>();
+    for (int b = 0; b < B; ++b) {
+      auto dt = [&](int j) { return d[b * ds[0] + j * ds[1]]; };
+      auto tt = [&](int j) { return t[b * ts[0] + j * ts[1]]; };
+      int n = K;
+      for (int j = 0; j < K; ++j) {
+        if (dt(j) != tt(j)) {
+          n = j;
+          break;
+        }
+      }
+      n_out[b] = n;
+      int32_t* c = c_out + b * (K + 1);
+      for (int j = 0; j < K + 1; ++j) {
+        c[j] = (j < n) ? dt(j) : ((j == n) ? tt(n) : 0);
+      }
+    }
+  });
+}
+
+} // namespace mlx::core::fast