chore: import upstream snapshot with attribution

2026-07-13 12:34:46 +08:00
commit f4e68ed970
84 changed files with 14896 additions and 0 deletions
@@ -0,0 +1,46 @@
+// Per-group INT8 GEMM/MV primitive header — symmetric quant (no bias)
+#pragma once
+
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+
+#include <string>
+
+namespace cider {
+
+namespace mx = mlx::core;
+
+// Per-group INT8 GEMM (prefill) / MV (decode) — symmetric quantization
+// Inputs:
+//   x: [M, K] float16/bfloat16 — activation
+//   w: [N, K] int8 — per-group symmetric quantized weight
+//   scale_w: [N, num_groups] float32 — per-group weight scales
+//   group_size: 64, 128, or 256
+class PerGroupLinear : public mx::Primitive {
+public:
+  PerGroupLinear(mx::Stream s, const std::string &kernel_dir, int group_size)
+      : mx::Primitive(s), kernel_dir_(kernel_dir), group_size_(group_size) {}
+
+  void eval_cpu(const std::vector<mx::array> &,
+                std::vector<mx::array> &) override {
+    throw std::runtime_error("PerGroupLinear: CPU not supported");
+  }
+  void eval_gpu(const std::vector<mx::array> &inputs,
+                std::vector<mx::array> &outputs) override;
+
+  const char *name() const override { return "PerGroupLinear"; }
+  bool is_equivalent(const mx::Primitive &other) const override { return true; }
+
+private:
+  std::string kernel_dir_;
+  int group_size_;
+};
+
+// Python-facing function
+mx::array pergroup_linear(const mx::array &x, const mx::array &w,
+                          const mx::array &scale_w, const mx::array &bias,
+                          const mx::array &new_bias,
+                          int group_size, const std::string &kernel_dir,
+                          mx::StreamOrDevice s = {});
+
+} // namespace cider
@@ -0,0 +1,143 @@
+// Cider SDPA — Custom Primitive for v9 optimized attention
+//
+// Dispatches 1-pass or 2-pass SDPA kernels via MLX's CommandEncoder.
+// All kernel selection (1pass vs 2pass, blocks value) is decided by the
+// Python caller and passed as arguments.
+
+#pragma once
+
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+
+#include <string>
+
+namespace cider {
+
+namespace mx = mlx::core;
+
+// ── 1-pass SDPA Primitive ────────────────────────────────────────
+// Inputs: [Q(B*H, D), K(...), V(...)]
+// Output: [O(B*H, D)]
+class SDPAVector : public mx::Primitive {
+ public:
+  SDPAVector(mx::Stream stream,
+             const std::string& kernel_dir,
+             int gqa_factor,
+             int N,
+             size_t k_head_stride,
+             size_t k_seq_stride,
+             size_t v_head_stride,
+             size_t v_seq_stride,
+             float scale)
+      : mx::Primitive(stream),
+        kernel_dir_(kernel_dir),
+        gqa_factor_(gqa_factor),
+        N_(N),
+        k_head_stride_(k_head_stride),
+        k_seq_stride_(k_seq_stride),
+        v_head_stride_(v_head_stride),
+        v_seq_stride_(v_seq_stride),
+        scale_(scale) {}
+
+  void eval_cpu(const std::vector<mx::array>& inputs,
+                std::vector<mx::array>& outputs) override {
+    throw std::runtime_error("SDPAVector: CPU not supported");
+  }
+
+  void eval_gpu(const std::vector<mx::array>& inputs,
+                std::vector<mx::array>& outputs) override;
+
+  const char* name() const override { return "CiderSDPAVector"; }
+  bool is_equivalent(const mx::Primitive& other) const override { return true; }
+
+ private:
+  std::string kernel_dir_;
+  int gqa_factor_;
+  int N_;
+  size_t k_head_stride_;
+  size_t k_seq_stride_;
+  size_t v_head_stride_;
+  size_t v_seq_stride_;
+  float scale_;
+};
+
+// ── 2-pass SDPA Primitive ────────────────────────────────────────
+// Inputs: [Q(B*H, D), K(...), V(...)]
+// Output: [O(B*H, D)]
+// Internally allocates partials/sums/maxs as temporaries.
+class SDPAVector2Pass : public mx::Primitive {
+ public:
+  SDPAVector2Pass(mx::Stream stream,
+                  const std::string& kernel_dir,
+                  int gqa_factor,
+                  int N,
+                  int blocks,
+                  int num_kv_heads,
+                  int batch_size,
+                  size_t k_head_stride,
+                  size_t k_seq_stride,
+                  size_t v_head_stride,
+                  size_t v_seq_stride,
+                  float scale)
+      : mx::Primitive(stream),
+        kernel_dir_(kernel_dir),
+        gqa_factor_(gqa_factor),
+        N_(N),
+        blocks_(blocks),
+        num_kv_heads_(num_kv_heads),
+        batch_size_(batch_size),
+        k_head_stride_(k_head_stride),
+        k_seq_stride_(k_seq_stride),
+        v_head_stride_(v_head_stride),
+        v_seq_stride_(v_seq_stride),
+        scale_(scale) {}
+
+  void eval_cpu(const std::vector<mx::array>& inputs,
+                std::vector<mx::array>& outputs) override {
+    throw std::runtime_error("SDPAVector2Pass: CPU not supported");
+  }
+
+  void eval_gpu(const std::vector<mx::array>& inputs,
+                std::vector<mx::array>& outputs) override;
+
+  const char* name() const override { return "CiderSDPAVector2Pass"; }
+  bool is_equivalent(const mx::Primitive& other) const override { return true; }
+
+ private:
+  std::string kernel_dir_;
+  int gqa_factor_;
+  int N_;
+  int blocks_;
+  int num_kv_heads_;
+  int batch_size_;
+  size_t k_head_stride_;
+  size_t k_seq_stride_;
+  size_t v_head_stride_;
+  size_t v_seq_stride_;
+  float scale_;
+};
+
+// ── Public API ───────────────────────────────────────────────────
+
+// 1-pass SDPA (short sequences)
+mx::array cider_sdpa_1pass(
+    const mx::array& queries,   // [B*H, D]
+    const mx::array& keys,      // [B*Hkv, N, D]
+    const mx::array& values,    // [B*Hkv, N, D]
+    int gqa_factor,
+    float scale,
+    const std::string& kernel_dir,
+    mx::StreamOrDevice s = {});
+
+// 2-pass SDPA (long sequences)
+mx::array cider_sdpa_2pass(
+    const mx::array& queries,   // [B, H, 1, D]
+    const mx::array& keys,      // [B, Hkv, N, D]
+    const mx::array& values,    // [B, Hkv, N, D]
+    int gqa_factor,
+    int blocks,
+    float scale,
+    const std::string& kernel_dir,
+    mx::StreamOrDevice s = {});
+
+}  // namespace cider
@@ -0,0 +1,55 @@
+// W4A8 Linear as mlx Custom Primitive.
+// Packed INT4 weights × INT8 activations via TensorOps.
+
+#pragma once
+
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+
+#include <string>
+
+namespace cider {
+
+namespace mx = mlx::core;
+
+// Inputs: [x(M,K) float16, packed_w(K/2,N) uint8, scale_w(N) float32]
+// Output: [y(M,N) float16]
+//
+// Weight layout: [K/2, N] uint8 — packed INT4 symmetric (zero_point=8)
+//   high nibble = even k index, low nibble = odd k index
+// scale_w: per-column scale (includes group scale pre-folded)
+class W4A8Linear : public mx::Primitive {
+ public:
+  explicit W4A8Linear(mx::Stream stream, const std::string& kernel_dir)
+      : mx::Primitive(stream), kernel_dir_(kernel_dir) {}
+
+  void eval_cpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override {
+    throw std::runtime_error("W4A8Linear: CPU not supported");
+  }
+
+  void eval_gpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+
+  const char* name() const override {
+    return "W4A8Linear";
+  }
+
+  bool is_equivalent(const mx::Primitive& other) const override {
+    return true;
+  }
+
+ private:
+  std::string kernel_dir_;
+};
+
+mx::array w4a8_linear(
+    const mx::array& x,          // [M, K] float16
+    const mx::array& packed_w,   // [K/2, N] uint8 (packed INT4)
+    const mx::array& scale_w,    // [N] float32
+    const std::string& kernel_dir,
+    mx::StreamOrDevice s = {});
+
+}  // namespace cider
@@ -0,0 +1,102 @@
+// W8A8 Linear as mlx Custom Primitive.
+// Dispatches quantize + INT8 matmul (prefill) or FP MV (decode) via mlx's CommandEncoder.
+
+#pragma once
+
+#include "mlx/ops.h"
+#include "mlx/primitives.h"
+
+#include <mutex>
+#include <string>
+#include <unordered_map>
+
+namespace MTL {
+class ComputePipelineState;
+}
+
+namespace cider {
+
+namespace mx = mlx::core;
+
+// ── Custom Primitive ─────────────────────────────────────────────
+// Inputs: [x(M,K) float16, w(N,K) int8, scale_w(N) float32, bias(N) float16]
+// Output: [y(M,N) float16]
+//
+// M > 1: quantize activation + INT8 GEMM (prefill)
+// M == 1: FP activation × dequant weight MV (decode)
+class W8A8Linear : public mx::Primitive {
+ public:
+  explicit W8A8Linear(mx::Stream stream, const std::string& kernel_dir)
+      : mx::Primitive(stream), kernel_dir_(kernel_dir) {}
+
+  void eval_cpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override {
+    throw std::runtime_error("W8A8Linear: CPU not supported");
+  }
+
+  void eval_gpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+
+  const char* name() const override {
+    return "W8A8Linear";
+  }
+
+  bool is_equivalent(const mx::Primitive& other) const override {
+    return true;
+  }
+
+ private:
+  std::string kernel_dir_;
+};
+
+
+// ── Raw INT32 Matmul Primitive ───────────────────────────────────
+// Inputs: [A(M,K) int8, B(K,N) int8]
+// Output: [C(M,N) int32]   (bit-exact, no dequant)
+class Int8MatMulInt32 : public mx::Primitive {
+ public:
+  explicit Int8MatMulInt32(mx::Stream stream, const std::string& kernel_dir)
+      : mx::Primitive(stream), kernel_dir_(kernel_dir) {}
+
+  void eval_cpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override {
+    throw std::runtime_error("Int8MatMulInt32: CPU not supported");
+  }
+
+  void eval_gpu(
+      const std::vector<mx::array>& inputs,
+      std::vector<mx::array>& outputs) override;
+
+  const char* name() const override {
+    return "Int8MatMulInt32";
+  }
+
+  bool is_equivalent(const mx::Primitive& other) const override {
+    return true;
+  }
+
+ private:
+  std::string kernel_dir_;
+};
+
+// ── Public API ───────────────────────────────────────────────────
+mx::array perchannel_linear(
+    const mx::array& x,       // [M, K] float16
+    const mx::array& w,       // [N, K] int8
+    const mx::array& scale_w, // [N] float32
+    const mx::array& bias,    // [N] float16
+    const std::string& kernel_dir,
+    mx::StreamOrDevice s = {});
+
+
+// ── Raw INT32 matmul (for bit-exact testing) ─────────────────────
+mx::array int8_matmul_int32(
+    const mx::array& a,       // [M, K] int8
+    const mx::array& b,       // [K, N] int8
+    const std::string& kernel_dir,
+    mx::StreamOrDevice s = {});
+
+}  // namespace cider