Rebase onto upstream main: take upstream 1-bit (#3161), re-apply the fork as 10 patches (gemma4-qat#282)

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,59 @@ 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;
+};
+
+// Shared by all copies of one kernel function. Holds at most
+// `metal_kernel_max_cached_variants` entries and is cleared when full. Callers
+// keep an owning pointer to their entry, so clearing is safe at any time.
+struct VariantCache {
+  std::mutex mutex;
+  std::unordered_map<std::string, std::shared_ptr<const KernelVariant>>
+      variants;
+};
+
+// Encodes each call argument that the name and source depend on. Template
+// arguments include their kind and name: `int` 1 and `bool` true give the same
+// name but a different signature.
+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 +325,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 +356,86 @@ 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;
+    {
+      std::lock_guard 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 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";
+        }
+      }
+      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});
+      std::lock_guard lock(cache->mutex);
+      if (auto it = cache->variants.find(key); it != cache->variants.end()) {
+        variant = it->second;
+      } else {
+        if (cache->variants.size() >= metal_kernel_max_cached_variants) {
+          cache->variants.clear();
+        }
+        cache->variants.emplace(std::move(key), variant);
+      }
     }
 
     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 +444,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 +454,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/quantized.cpp

@@ -943,8 +943,71 @@ 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; pass biases (e.g. derive_biases) on CPU.");
+    }
     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/cuda/custom_kernel.cpp

@@ -310,14 +310,13 @@ void CustomKernel::eval_gpu(
   // Compile the custom kernel
   std::string kernel_name =
       (is_precompiled_) ? name_ : "mlx::core::cu::" + name_;
-  std::string module_name =
-      fmt::format("{}_{:x}", name_, std::hash<std::string>{}(source_));
+  std::string module_name = fmt::format("{}_{:x}", name_, source_hash_);
   cu::JitModule& mod = cu::get_jit_module(
       encoder.device(),
       module_name,
       [&]() {
         return std::make_tuple(
-            is_precompiled_, source_, std::vector{kernel_name});
+            is_precompiled_, *source_, std::vector{kernel_name});
       },
       false);
 

mlx/backend/cuda/eval.cpp

@@ -83,7 +83,7 @@ void finalize(Stream s) {
   cu::get_command_encoder(s).commit();
 }
 
-void synchronize(Stream s) {
+void synchronize(Stream s, bool) {
   nvtx3::scoped_range r("gpu::synchronize");
   cu::get_command_encoder(s).synchronize();
 }

mlx/backend/cuda/primitives.cpp

@@ -49,6 +49,7 @@ NO_GPU_MULTI(SVD)
 NO_GPU(Inverse)
 NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
+NO_GPU_USE_FALLBACK(fast::SpecDecodeVerify)
 
 namespace fast {
 NO_GPU_MULTI(GatedDeltaUpdate)

mlx/backend/cuda/quantized/quantized.cpp

@@ -24,6 +24,10 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   if (inputs.size() > 3) {
     biases = inputs[3];
   }
+  if (mode_ == QuantizationMode::Affine && !biases) {
+    throw std::runtime_error(
+        "[QuantizedMatmul::eval_gpu] Bias-free affine is Metal-only for now.");
+  }
 
   auto supports = [&](auto&& f) {
     return f(

mlx/backend/gpu/eval.h

@@ -15,7 +15,8 @@ void new_stream(Stream s);
 void new_thread_unsafe_stream(Stream s);
 void eval(array& arr);
 void finalize(Stream s);
-void synchronize(Stream s);
+// explicit_sync is false for internal flushes (e.g. eval error recovery).
+void synchronize(Stream s, bool explicit_sync = true);
 void clear_streams();
 
 } // namespace mlx::core::gpu