chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
@@ -0,0 +1,314 @@
|
||||
// Per-group INT8 GEMM primitive implementation
|
||||
//
|
||||
// Dispatches:
|
||||
// 1. quantize_per_token: x[M,K] float16 → a_int8[M,K] int8 + scale_a[M]
|
||||
// float32
|
||||
// 2. pergroup_int8_gemm_gXX: A[M,K] int8 × B[N,K] int8 → C[M,N] float16
|
||||
// with per-group scale_w[N, num_groups] and per-token scale_a[M]
|
||||
|
||||
#include "pergroup_primitive.h"
|
||||
#include "mlx/backend/metal/device.h"
|
||||
#include "mlx/backend/metal/utils.h"
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
|
||||
#include <fstream>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace cider {
|
||||
|
||||
using namespace mlx::core;
|
||||
|
||||
// ── Pipeline cache for per-group kernels ────────────────────────
|
||||
class PerGroupPipelineCache {
|
||||
public:
|
||||
static PerGroupPipelineCache &instance() {
|
||||
static PerGroupPipelineCache cache;
|
||||
return cache;
|
||||
}
|
||||
|
||||
void ensure_init(const std::string &kernel_dir) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (initialized_ && kernel_dir_ == kernel_dir)
|
||||
return;
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *mtl_device = dev.mtl_device();
|
||||
pergroup_lib_ =
|
||||
compile_source(mtl_device, kernel_dir + "/pergroup_int8_gemm.metal");
|
||||
mv_lib_ =
|
||||
compile_source(mtl_device, kernel_dir + "/pergroup_int8_mv.metal");
|
||||
quantize_lib_ =
|
||||
compile_source(mtl_device, kernel_dir + "/w8a8_quantize.metal");
|
||||
kernel_dir_ = kernel_dir;
|
||||
pipelines_.clear();
|
||||
initialized_ = true;
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *get(const std::string &name) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
auto it = pipelines_.find(name);
|
||||
if (it != pipelines_.end()) {
|
||||
return it->second;
|
||||
}
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *pso = make_pipeline(dev.mtl_device(), name);
|
||||
pipelines_[name] = pso;
|
||||
return pso;
|
||||
}
|
||||
|
||||
private:
|
||||
PerGroupPipelineCache() = default;
|
||||
bool initialized_ = false;
|
||||
std::string kernel_dir_;
|
||||
std::unordered_map<std::string, MTL::ComputePipelineState *> pipelines_;
|
||||
MTL::Library *pergroup_lib_ = nullptr;
|
||||
MTL::Library *mv_lib_ = nullptr;
|
||||
MTL::Library *quantize_lib_ = nullptr;
|
||||
std::mutex mutex_;
|
||||
|
||||
static std::string read_file(const std::string &path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.is_open()) {
|
||||
throw std::runtime_error("Cannot open: " + path);
|
||||
}
|
||||
std::stringstream ss;
|
||||
ss << f.rdbuf();
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
MTL::Library *compile_source(MTL::Device *mtl_device,
|
||||
const std::string &source_path) {
|
||||
std::string source = read_file(source_path);
|
||||
@autoreleasepool {
|
||||
NSString *src = [NSString stringWithUTF8String:source.c_str()];
|
||||
MTLCompileOptions *opts = [[MTLCompileOptions alloc] init];
|
||||
opts.languageVersion = MTLLanguageVersion4_0;
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLLibrary> lib_objc = [device_objc newLibraryWithSource:src
|
||||
options:opts
|
||||
error:&error];
|
||||
if (!lib_objc) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("Metal compile failed (" + source_path +
|
||||
"): " + err);
|
||||
}
|
||||
return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc);
|
||||
}
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device,
|
||||
const std::string &name) {
|
||||
@autoreleasepool {
|
||||
NSString *fn_name = [NSString stringWithUTF8String:name.c_str()];
|
||||
// Try pergroup lib, then mv lib, then quantize lib
|
||||
id<MTLFunction> func = nil;
|
||||
id<MTLLibrary> lib_objc = (__bridge id<MTLLibrary>)pergroup_lib_;
|
||||
func = [lib_objc newFunctionWithName:fn_name];
|
||||
if (!func) {
|
||||
lib_objc = (__bridge id<MTLLibrary>)mv_lib_;
|
||||
func = [lib_objc newFunctionWithName:fn_name];
|
||||
}
|
||||
if (!func) {
|
||||
lib_objc = (__bridge id<MTLLibrary>)quantize_lib_;
|
||||
func = [lib_objc newFunctionWithName:fn_name];
|
||||
}
|
||||
if (!func) {
|
||||
throw std::runtime_error("Kernel not found: " + name);
|
||||
}
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLComputePipelineState> pso =
|
||||
[device_objc newComputePipelineStateWithFunction:func error:&error];
|
||||
if (!pso) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("Pipeline failed for " + name + ": " + err);
|
||||
}
|
||||
return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain(
|
||||
pso);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void PerGroupLinear::eval_gpu(const std::vector<array> &inputs,
|
||||
std::vector<array> &outputs) {
|
||||
auto &x = inputs[0]; // [M, K] float16
|
||||
auto &w = inputs[1]; // [N, K] int8
|
||||
auto &scale_w = inputs[2]; // [N, num_groups] float32
|
||||
auto &bias = inputs[3]; // [N] float16
|
||||
auto &new_bias = inputs[4]; // [N, num_groups] float32 (asymmetric correction)
|
||||
auto &out = outputs[0]; // [M, N] float16
|
||||
|
||||
uint32_t M = static_cast<uint32_t>(x.shape(0));
|
||||
uint32_t N = static_cast<uint32_t>(w.shape(0));
|
||||
uint32_t K = static_cast<uint32_t>(w.shape(1));
|
||||
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
|
||||
auto &cache = PerGroupPipelineCache::instance();
|
||||
cache.ensure_init(kernel_dir_);
|
||||
|
||||
auto &s = stream();
|
||||
auto &enc = metal::get_command_encoder(s);
|
||||
|
||||
if (M == 1) {
|
||||
// ── Decode path: per-group MV kernel ──
|
||||
// No activation quantization needed; directly use FP16 activation
|
||||
std::string kname;
|
||||
if (group_size_ == 64) {
|
||||
kname = "pergroup_int8_mv_g64";
|
||||
} else if (group_size_ == 128) {
|
||||
kname = "pergroup_int8_mv_g128";
|
||||
} else {
|
||||
kname = "pergroup_int8_mv_g256";
|
||||
}
|
||||
|
||||
auto *pso = cache.get(kname);
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
|
||||
constexpr uint32_t TOTAL_ROWS = 8; // NUM_SIMDGROUPS(2) * RESULTS_PER_SG(4)
|
||||
uint32_t threadgroups = (N + TOTAL_ROWS - 1) / TOTAL_ROWS;
|
||||
uint32_t threads = 64; // 2 simdgroups x 32 threads
|
||||
|
||||
enc.set_input_array(x, 0); // [1, K] float16 → pass as [K]
|
||||
enc.set_input_array(w, 1); // [N, K] int8
|
||||
enc.set_output_array(out, 2); // [1, N] float16 → write as [N]
|
||||
enc.set_input_array(scale_w, 3); // [N, num_groups] float32
|
||||
enc.set_bytes(N, 4);
|
||||
enc.set_bytes(K, 5);
|
||||
enc.set_input_array(bias, 6); // [N] float16
|
||||
enc.set_input_array(new_bias, 7); // [N, num_groups] float32 (correction)
|
||||
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(threadgroups, 1, 1),
|
||||
MTL::Size::Make(threads, 1, 1));
|
||||
} else {
|
||||
// ── Prefill path: per-group GEMM with activation quantization ──
|
||||
|
||||
// Scratch: activation quantization
|
||||
size_t a_bytes = static_cast<size_t>(M) * K;
|
||||
size_t sa_bytes = static_cast<size_t>(M) * sizeof(float);
|
||||
|
||||
array a_int8({static_cast<int>(M), static_cast<int>(K)}, int8, nullptr, {});
|
||||
a_int8.set_data(allocator::malloc(a_bytes));
|
||||
|
||||
array sa({static_cast<int>(M)}, float32, nullptr, {});
|
||||
sa.set_data(allocator::malloc(sa_bytes));
|
||||
|
||||
// DISPATCH 1: quantize_per_token
|
||||
{
|
||||
auto *pso = cache.get("quantize_per_token");
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
enc.set_input_array(x, 0);
|
||||
enc.set_output_array(a_int8, 1);
|
||||
enc.set_output_array(sa, 2);
|
||||
enc.set_bytes(M, 3);
|
||||
enc.set_bytes(K, 4);
|
||||
uint32_t tg = std::min(256u, std::max(32u, ((K + 31) / 32) * 32));
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(M, 1, 1),
|
||||
MTL::Size::Make(tg, 1, 1));
|
||||
}
|
||||
|
||||
enc.barrier();
|
||||
|
||||
// DISPATCH 2: per-group GEMM
|
||||
{
|
||||
std::string kname;
|
||||
bool use_small = (M <= 64);
|
||||
uint32_t BM, BN, threads;
|
||||
|
||||
if (group_size_ == 64) {
|
||||
kname = use_small ? "pergroup_int8_gemm_g64_small"
|
||||
: "pergroup_int8_gemm_g64";
|
||||
} else if (group_size_ == 128) {
|
||||
kname = use_small ? "pergroup_int8_gemm_g128_small"
|
||||
: "pergroup_int8_gemm_g128";
|
||||
} else {
|
||||
kname = use_small ? "pergroup_int8_gemm_g256_small"
|
||||
: "pergroup_int8_gemm_g256";
|
||||
}
|
||||
|
||||
BM = use_small ? 32 : 128;
|
||||
BN = 128;
|
||||
threads = use_small ? 128 : 512;
|
||||
|
||||
auto *pso = cache.get(kname);
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
|
||||
uint32_t tiles_n = (N + BN - 1) / BN;
|
||||
uint32_t tiles_m = (M + BM - 1) / BM;
|
||||
uint32_t swizzle_log;
|
||||
if (tiles_m <= 3) {
|
||||
swizzle_log = 0;
|
||||
} else if (tiles_m <= 6) {
|
||||
swizzle_log = 1;
|
||||
} else {
|
||||
swizzle_log = 2;
|
||||
}
|
||||
|
||||
uint32_t tile = 1u << swizzle_log;
|
||||
uint32_t grid_x = tiles_n * tile;
|
||||
uint32_t grid_y = (tiles_m + tile - 1) / tile;
|
||||
|
||||
enc.set_input_array(a_int8, 0);
|
||||
enc.set_input_array(w, 1);
|
||||
enc.set_output_array(out, 2);
|
||||
enc.set_bytes(M, 3);
|
||||
enc.set_bytes(N, 4);
|
||||
enc.set_bytes(K, 5);
|
||||
enc.set_input_array(sa, 6);
|
||||
enc.set_input_array(scale_w, 7);
|
||||
enc.set_bytes(swizzle_log, 8);
|
||||
enc.set_bytes(tiles_m, 9);
|
||||
enc.set_bytes(tiles_n, 10);
|
||||
enc.set_input_array(bias, 11);
|
||||
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(grid_x, grid_y, 1),
|
||||
MTL::Size::Make(threads, 1, 1));
|
||||
}
|
||||
|
||||
enc.add_temporary(a_int8);
|
||||
enc.add_temporary(sa);
|
||||
}
|
||||
}
|
||||
|
||||
array pergroup_linear(const array &x, const array &w, const array &scale_w,
|
||||
const array &bias, const array &new_bias, int group_size,
|
||||
const std::string &kernel_dir, StreamOrDevice s) {
|
||||
if (x.ndim() != 2) {
|
||||
throw std::invalid_argument("pergroup_linear: x must be 2D [M,K]");
|
||||
}
|
||||
if (w.ndim() != 2) {
|
||||
throw std::invalid_argument("pergroup_linear: w must be 2D [N,K]");
|
||||
}
|
||||
if (scale_w.ndim() != 2) {
|
||||
throw std::invalid_argument(
|
||||
"pergroup_linear: scale_w must be 2D [N,num_groups]");
|
||||
}
|
||||
if (group_size != 64 && group_size != 128 && group_size != 256) {
|
||||
throw std::invalid_argument(
|
||||
"pergroup_linear: group_size must be 64, 128, or 256");
|
||||
}
|
||||
|
||||
int M = x.shape(0);
|
||||
int N = w.shape(0);
|
||||
auto stream = to_stream(s);
|
||||
|
||||
auto result =
|
||||
array({M, N}, float16,
|
||||
std::make_shared<PerGroupLinear>(stream, kernel_dir, group_size),
|
||||
{astype(x, float16, stream), astype(w, int8, stream),
|
||||
astype(scale_w, float32, stream), astype(bias, float16, stream),
|
||||
astype(new_bias, float32, stream)});
|
||||
|
||||
if (x.dtype() == bfloat16) {
|
||||
return astype(result, bfloat16, stream);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace cider
|
||||
@@ -0,0 +1,49 @@
|
||||
#include <nanobind/nanobind.h>
|
||||
#include <nanobind/stl/optional.h>
|
||||
#include <nanobind/stl/string.h>
|
||||
#include <nanobind/stl/variant.h>
|
||||
|
||||
#include "mlx/ops.h"
|
||||
#include "pergroup_primitive.h"
|
||||
#include "sdpa_primitive.h"
|
||||
#include "w4a8_primitive.h"
|
||||
#include "w8a8_primitive.h"
|
||||
|
||||
namespace nb = nanobind;
|
||||
using namespace nb::literals;
|
||||
namespace mx = mlx::core;
|
||||
|
||||
NB_MODULE(_cider_prim, m) {
|
||||
m.doc() = "cider: W8A8 + W4A8 INT8 TensorOps + SDPA primitives for Apple M5+";
|
||||
|
||||
m.def("perchannel_linear", &cider::perchannel_linear, "x"_a, "w"_a, "scale_w"_a,
|
||||
"bias"_a, "kernel_dir"_a, nb::kw_only(), "stream"_a = nb::none(),
|
||||
"W8A8 quantized linear: y = dequant(quant_a(x) @ w_int8) + bias");
|
||||
|
||||
m.def("w4a8_linear", &cider::w4a8_linear, "x"_a, "packed_w"_a, "scale_w"_a,
|
||||
"kernel_dir"_a, nb::kw_only(), "stream"_a = nb::none(),
|
||||
"W4A8 quantized linear: y = dequant(quant_a(x) @ unpack4(w))");
|
||||
|
||||
m.def("int8_matmul_int32", &cider::int8_matmul_int32, "a"_a, "b"_a,
|
||||
"kernel_dir"_a, nb::kw_only(), "stream"_a = nb::none(),
|
||||
"Raw INT8xINT8->INT32 matmul (bit-exact, no dequant)");
|
||||
|
||||
m.def("pergroup_linear", &cider::pergroup_linear, "x"_a, "w"_a, "scale_w"_a,
|
||||
"bias"_a, "new_bias"_a, "group_size"_a, "kernel_dir"_a, nb::kw_only(),
|
||||
"stream"_a = nb::none(),
|
||||
"Per-group INT8 linear with bias: prefill GEMM or decode MV with "
|
||||
"per-group scales");
|
||||
|
||||
// ── SDPA ──
|
||||
m.def("cider_sdpa_1pass", &cider::cider_sdpa_1pass,
|
||||
"queries"_a, "keys"_a, "values"_a,
|
||||
"gqa_factor"_a, "scale"_a, "kernel_dir"_a,
|
||||
nb::kw_only(), "stream"_a = nb::none(),
|
||||
"Cider v9 SDPA 1-pass (short sequences)");
|
||||
|
||||
m.def("cider_sdpa_2pass", &cider::cider_sdpa_2pass,
|
||||
"queries"_a, "keys"_a, "values"_a,
|
||||
"gqa_factor"_a, "blocks"_a, "scale"_a, "kernel_dir"_a,
|
||||
nb::kw_only(), "stream"_a = nb::none(),
|
||||
"Cider v9 SDPA 2-pass with contiguous chunks + TILE=4 tiling");
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
// sdpa_primitive.mm — Cider v9 SDPA Custom Primitive implementation
|
||||
//
|
||||
// Dispatches precompiled Metal kernels for optimized attention.
|
||||
// Two primitives: SDPAVector (1-pass) and SDPAVector2Pass (2-pass).
|
||||
|
||||
#include "sdpa_primitive.h"
|
||||
#include "mlx/backend/metal/device.h"
|
||||
#include "mlx/backend/metal/utils.h"
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
|
||||
#include <fstream>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace cider {
|
||||
|
||||
using namespace mlx::core;
|
||||
|
||||
// ── Pipeline cache for SDPA kernels ─────────────────────────────
|
||||
class SDPAPipelineCache {
|
||||
public:
|
||||
static SDPAPipelineCache &instance() {
|
||||
static SDPAPipelineCache cache;
|
||||
return cache;
|
||||
}
|
||||
|
||||
void ensure_init(const std::string &kernel_dir) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (initialized_ && kernel_dir_ == kernel_dir)
|
||||
return;
|
||||
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *mtl_device = dev.mtl_device();
|
||||
|
||||
// Read and compile the SDPA metal source
|
||||
std::string source = read_file(kernel_dir + "/cider_sdpa_vector.metal");
|
||||
|
||||
// Prepend the include path — the header is in the same directory
|
||||
// We inline the header by reading it and prepending
|
||||
std::string header = read_file(kernel_dir + "/cider_sdpa_vector.h");
|
||||
|
||||
// Replace the #include directive with the actual header content
|
||||
std::string combined;
|
||||
combined.reserve(header.size() + source.size() + 256);
|
||||
// Add metal_stdlib include
|
||||
combined += "#include <metal_stdlib>\n";
|
||||
combined += "#include <metal_simdgroup>\n";
|
||||
combined += "using namespace metal;\n\n";
|
||||
combined += "// === Inlined cider_sdpa_vector.h ===\n";
|
||||
combined += header;
|
||||
combined += "\n// === End inlined header ===\n\n";
|
||||
|
||||
// Append the .metal source but strip its own #include lines
|
||||
std::istringstream iss(source);
|
||||
std::string line;
|
||||
while (std::getline(iss, line)) {
|
||||
if (line.find("#include") != std::string::npos)
|
||||
continue;
|
||||
if (line.find("using namespace metal") != std::string::npos)
|
||||
continue;
|
||||
combined += line + "\n";
|
||||
}
|
||||
|
||||
lib_ = compile_source(mtl_device, combined);
|
||||
kernel_dir_ = kernel_dir;
|
||||
pipelines_.clear();
|
||||
initialized_ = true;
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *get(const std::string &name) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
auto it = pipelines_.find(name);
|
||||
if (it != pipelines_.end())
|
||||
return it->second;
|
||||
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *pso = make_pipeline(dev.mtl_device(), name);
|
||||
pipelines_[name] = pso;
|
||||
return pso;
|
||||
}
|
||||
|
||||
private:
|
||||
SDPAPipelineCache() = default;
|
||||
bool initialized_ = false;
|
||||
std::string kernel_dir_;
|
||||
std::unordered_map<std::string, MTL::ComputePipelineState *> pipelines_;
|
||||
MTL::Library *lib_ = nullptr;
|
||||
std::mutex mutex_;
|
||||
|
||||
static std::string read_file(const std::string &path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.is_open()) {
|
||||
throw std::runtime_error("Cannot open: " + path);
|
||||
}
|
||||
std::stringstream ss;
|
||||
ss << f.rdbuf();
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
MTL::Library *compile_source(MTL::Device *mtl_device,
|
||||
const std::string &source) {
|
||||
@autoreleasepool {
|
||||
NSString *src = [NSString stringWithUTF8String:source.c_str()];
|
||||
MTLCompileOptions *opts = [[MTLCompileOptions alloc] init];
|
||||
opts.languageVersion = MTLLanguageVersion3_1;
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLLibrary> lib_objc = [device_objc newLibraryWithSource:src
|
||||
options:opts
|
||||
error:&error];
|
||||
if (!lib_objc) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("SDPA Metal compile failed: " + err);
|
||||
}
|
||||
return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc);
|
||||
}
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device,
|
||||
const std::string &name) {
|
||||
@autoreleasepool {
|
||||
NSString *fn_name = [NSString stringWithUTF8String:name.c_str()];
|
||||
id<MTLLibrary> lib_objc = (__bridge id<MTLLibrary>)lib_;
|
||||
id<MTLFunction> func = [lib_objc newFunctionWithName:fn_name];
|
||||
if (!func) {
|
||||
throw std::runtime_error("SDPA kernel not found: " + name);
|
||||
}
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLComputePipelineState> pso =
|
||||
[device_objc newComputePipelineStateWithFunction:func error:&error];
|
||||
if (!pso) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("SDPA pipeline failed for " + name + ": " +
|
||||
err);
|
||||
}
|
||||
return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain(
|
||||
pso);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// ── Helper: dtype string ────────────────────────────────────────
|
||||
static std::string dtype_str(Dtype dt) {
|
||||
switch (dt) {
|
||||
case float32:
|
||||
return "float32";
|
||||
case float16:
|
||||
return "float16";
|
||||
case bfloat16:
|
||||
return "bfloat16";
|
||||
default:
|
||||
throw std::runtime_error("SDPA: unsupported dtype");
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// SDPAVector (1-pass) eval_gpu
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
void SDPAVector::eval_gpu(const std::vector<array> &inputs,
|
||||
std::vector<array> &outputs) {
|
||||
auto &queries = inputs[0]; // [B*H, D]
|
||||
auto &keys = inputs[1];
|
||||
auto &values = inputs[2];
|
||||
auto &out = outputs[0];
|
||||
|
||||
int D = queries.shape(-1);
|
||||
int total_heads = queries.shape(0);
|
||||
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
|
||||
auto &cache = SDPAPipelineCache::instance();
|
||||
cache.ensure_init(kernel_dir_);
|
||||
|
||||
// Kernel name: cider_sdpa_vector_{type}_{D}
|
||||
std::string kname = "cider_sdpa_vector_" + dtype_str(queries.dtype()) + "_" +
|
||||
std::to_string(D);
|
||||
auto *pso = cache.get(kname);
|
||||
|
||||
auto &s = stream();
|
||||
auto &enc = metal::get_command_encoder(s);
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
|
||||
enc.set_input_array(queries, 0);
|
||||
enc.set_input_array(keys, 1);
|
||||
enc.set_input_array(values, 2);
|
||||
enc.set_output_array(out, 3);
|
||||
enc.set_bytes(gqa_factor_, 4);
|
||||
enc.set_bytes(N_, 5);
|
||||
enc.set_bytes(k_head_stride_, 6);
|
||||
enc.set_bytes(k_seq_stride_, 7);
|
||||
enc.set_bytes(v_head_stride_, 8);
|
||||
enc.set_bytes(v_seq_stride_, 9);
|
||||
enc.set_bytes(scale_, 10);
|
||||
|
||||
// 1-pass: grid=(total_heads, 1, 1), group=(1024, 1, 1)
|
||||
MTL::Size grid_dims(total_heads, 1, 1);
|
||||
MTL::Size group_dims(1024, 1, 1);
|
||||
enc.dispatch_threadgroups(grid_dims, group_dims);
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// SDPAVector2Pass eval_gpu
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
void SDPAVector2Pass::eval_gpu(const std::vector<array> &inputs,
|
||||
std::vector<array> &outputs) {
|
||||
auto &queries = inputs[0]; // [B, H, 1, D]
|
||||
auto &keys = inputs[1]; // [B, Hkv, N, D]
|
||||
auto &values = inputs[2]; // [B, Hkv, N, D]
|
||||
auto &out = outputs[0]; // [B, H, 1, D]
|
||||
|
||||
int D = queries.shape(-1);
|
||||
int num_q_heads = num_kv_heads_ * gqa_factor_;
|
||||
int total_q_heads = batch_size_ * num_q_heads;
|
||||
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
|
||||
auto &cache = SDPAPipelineCache::instance();
|
||||
cache.ensure_init(kernel_dir_);
|
||||
|
||||
auto &s = stream();
|
||||
auto &enc = metal::get_command_encoder(s);
|
||||
|
||||
// ── Allocate intermediates ──
|
||||
int partial_total = total_q_heads * blocks_;
|
||||
size_t partial_bytes =
|
||||
static_cast<size_t>(partial_total) * D * size_of(queries.dtype());
|
||||
size_t stats_bytes = static_cast<size_t>(partial_total) * sizeof(float);
|
||||
|
||||
array intermediate({partial_total, D}, queries.dtype(), nullptr, {});
|
||||
array sums({partial_total}, float32, nullptr, {});
|
||||
array maxs({partial_total}, float32, nullptr, {});
|
||||
intermediate.set_data(allocator::malloc(intermediate.nbytes()));
|
||||
sums.set_data(allocator::malloc(sums.nbytes()));
|
||||
maxs.set_data(allocator::malloc(maxs.nbytes()));
|
||||
enc.add_temporary(intermediate);
|
||||
enc.add_temporary(sums);
|
||||
enc.add_temporary(maxs);
|
||||
|
||||
// ── Pass 1 ──
|
||||
std::string kname1 = "cider_sdpa_vector_2pass_1_" +
|
||||
dtype_str(queries.dtype()) + "_" + std::to_string(D) +
|
||||
"_b" + std::to_string(blocks_);
|
||||
auto *pso1 = cache.get(kname1);
|
||||
enc.set_compute_pipeline_state(pso1);
|
||||
|
||||
enc.set_input_array(queries, 0);
|
||||
enc.set_input_array(keys, 1);
|
||||
enc.set_input_array(values, 2);
|
||||
enc.set_output_array(intermediate, 3);
|
||||
enc.set_output_array(sums, 4);
|
||||
enc.set_output_array(maxs, 5);
|
||||
// buffer 6 is skipped (matching MLX convention)
|
||||
enc.set_bytes(N_, 7);
|
||||
enc.set_bytes(k_head_stride_, 8);
|
||||
enc.set_bytes(k_seq_stride_, 9);
|
||||
enc.set_bytes(v_head_stride_, 10);
|
||||
enc.set_bytes(v_seq_stride_, 11);
|
||||
enc.set_bytes(scale_, 12);
|
||||
|
||||
// grid=(num_kv_heads, batch_size, blocks), group=(32, gqa_factor, 1)
|
||||
MTL::Size grid1(num_kv_heads_, batch_size_, blocks_);
|
||||
MTL::Size group1(32, gqa_factor_, 1);
|
||||
enc.dispatch_threadgroups(grid1, group1);
|
||||
|
||||
// ── Pass 2 ──
|
||||
std::string kname2 = "cider_sdpa_vector_2pass_2_" +
|
||||
dtype_str(queries.dtype()) + "_" + std::to_string(D) +
|
||||
"_b" + std::to_string(blocks_);
|
||||
auto *pso2 = cache.get(kname2);
|
||||
enc.set_compute_pipeline_state(pso2);
|
||||
|
||||
enc.set_input_array(intermediate, 0);
|
||||
enc.set_input_array(sums, 1);
|
||||
enc.set_input_array(maxs, 2);
|
||||
enc.set_output_array(out, 3);
|
||||
|
||||
// grid=(total_q_heads, 1, 1), group=(1024, 1, 1)
|
||||
MTL::Size grid2(total_q_heads, 1, 1);
|
||||
MTL::Size group2(1024, 1, 1);
|
||||
enc.dispatch_threadgroups(grid2, group2);
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// Public API
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
|
||||
array cider_sdpa_1pass(const array &queries, const array &keys,
|
||||
const array &values, int gqa_factor, float scale,
|
||||
const std::string &kernel_dir, StreamOrDevice s) {
|
||||
auto stream = to_stream(s);
|
||||
|
||||
int total_heads = queries.shape(0);
|
||||
int D = queries.shape(-1);
|
||||
int N = keys.shape(-2);
|
||||
|
||||
// Compute strides from key/value shapes
|
||||
// keys: contiguous [B*Hkv, N, D]
|
||||
size_t k_head_stride = static_cast<size_t>(N) * D;
|
||||
size_t k_seq_stride = D;
|
||||
size_t v_head_stride = k_head_stride;
|
||||
size_t v_seq_stride = D;
|
||||
|
||||
return array({total_heads, D}, queries.dtype(),
|
||||
std::make_shared<SDPAVector>(stream, kernel_dir, gqa_factor, N,
|
||||
k_head_stride, k_seq_stride,
|
||||
v_head_stride, v_seq_stride, scale),
|
||||
{queries, keys, values});
|
||||
}
|
||||
|
||||
array cider_sdpa_2pass(const array &queries, const array &keys,
|
||||
const array &values, int gqa_factor, int blocks,
|
||||
float scale, const std::string &kernel_dir,
|
||||
StreamOrDevice s) {
|
||||
auto stream = to_stream(s);
|
||||
|
||||
int B = queries.shape(0);
|
||||
int H = queries.shape(1);
|
||||
int D = queries.shape(-1);
|
||||
int Hkv = keys.shape(1);
|
||||
int N = keys.shape(2);
|
||||
|
||||
// Strides for 4D layout [B, Hkv, N, D]
|
||||
size_t k_head_stride = static_cast<size_t>(N) * D;
|
||||
size_t k_seq_stride = D;
|
||||
size_t v_head_stride = k_head_stride;
|
||||
size_t v_seq_stride = D;
|
||||
|
||||
auto out_shape = queries.shape();
|
||||
|
||||
return array(out_shape, queries.dtype(),
|
||||
std::make_shared<SDPAVector2Pass>(
|
||||
stream, kernel_dir, gqa_factor, N, blocks, Hkv, B,
|
||||
k_head_stride, k_seq_stride, v_head_stride, v_seq_stride,
|
||||
scale),
|
||||
{queries, keys, values});
|
||||
}
|
||||
|
||||
} // namespace cider
|
||||
@@ -0,0 +1,285 @@
|
||||
// W4A8 Linear V1 — Packed INT4 weights × INT8 activations
|
||||
// Reuses W8A8's scratch buffer pool for activation quantization.
|
||||
// Kernel: unpack W4→INT8 in threadgroup mem, then TensorOps matmul2d.
|
||||
|
||||
#include "w4a8_primitive.h"
|
||||
#include "mlx/backend/metal/device.h"
|
||||
#include "mlx/backend/metal/utils.h"
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
|
||||
#include <fstream>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace cider {
|
||||
|
||||
using namespace mlx::core;
|
||||
|
||||
struct TileLarge {
|
||||
static constexpr uint32_t BM = 128, BN = 128, THREADS = 512;
|
||||
};
|
||||
struct TileSmall {
|
||||
static constexpr uint32_t BM = 32, BN = 128, THREADS = 128;
|
||||
};
|
||||
|
||||
static std::string read_file(const std::string &path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.is_open()) {
|
||||
throw std::runtime_error("Cannot open: " + path);
|
||||
}
|
||||
std::stringstream ss;
|
||||
ss << f.rdbuf();
|
||||
f.close();
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
// Scratch buffer pool (same pattern as W8A8)
|
||||
struct ScratchPool {
|
||||
allocator::Buffer a_int8_buf{nullptr};
|
||||
size_t a_int8_size = 0;
|
||||
allocator::Buffer scale_a_buf{nullptr};
|
||||
size_t scale_a_size = 0;
|
||||
|
||||
allocator::Buffer get_a_int8(size_t needed) {
|
||||
if (a_int8_buf.ptr() && a_int8_size >= needed) {
|
||||
return a_int8_buf;
|
||||
}
|
||||
if (a_int8_buf.ptr()) {
|
||||
allocator::free(a_int8_buf);
|
||||
}
|
||||
a_int8_buf = allocator::malloc(needed);
|
||||
a_int8_size = needed;
|
||||
return a_int8_buf;
|
||||
}
|
||||
allocator::Buffer get_scale_a(size_t needed) {
|
||||
if (scale_a_buf.ptr() && scale_a_size >= needed) {
|
||||
return scale_a_buf;
|
||||
}
|
||||
if (scale_a_buf.ptr()) {
|
||||
allocator::free(scale_a_buf);
|
||||
}
|
||||
scale_a_buf = allocator::malloc(needed);
|
||||
scale_a_size = needed;
|
||||
return scale_a_buf;
|
||||
}
|
||||
~ScratchPool() {
|
||||
if (a_int8_buf.ptr()) {
|
||||
allocator::free(a_int8_buf);
|
||||
}
|
||||
if (scale_a_buf.ptr()) {
|
||||
allocator::free(scale_a_buf);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
class W4A8PipelineCache {
|
||||
public:
|
||||
static W4A8PipelineCache &instance() {
|
||||
static W4A8PipelineCache cache;
|
||||
return cache;
|
||||
}
|
||||
|
||||
void ensure_init(const std::string &kernel_dir) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (initialized_ && kernel_dir_ == kernel_dir) {
|
||||
return;
|
||||
}
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *mtl_device = dev.mtl_device();
|
||||
w4a8_lib_ = compile_source(mtl_device, kernel_dir + "/w4a8_matmul.metal");
|
||||
quantize_lib_ =
|
||||
compile_source(mtl_device, kernel_dir + "/w8a8_quantize.metal");
|
||||
kernel_dir_ = kernel_dir;
|
||||
pipelines_.clear();
|
||||
initialized_ = true;
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *get(const std::string &name) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
auto it = pipelines_.find(name);
|
||||
if (it != pipelines_.end()) {
|
||||
return it->second;
|
||||
}
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *pso = make_pipeline(dev.mtl_device(), name);
|
||||
pipelines_[name] = pso;
|
||||
return pso;
|
||||
}
|
||||
|
||||
ScratchPool &scratch() { return scratch_pool_; }
|
||||
|
||||
private:
|
||||
W4A8PipelineCache() = default;
|
||||
bool initialized_ = false;
|
||||
std::string kernel_dir_;
|
||||
std::unordered_map<std::string, MTL::ComputePipelineState *> pipelines_;
|
||||
MTL::Library *w4a8_lib_ = nullptr;
|
||||
MTL::Library *quantize_lib_ = nullptr;
|
||||
ScratchPool scratch_pool_;
|
||||
std::mutex mutex_;
|
||||
|
||||
MTL::Library *compile_source(MTL::Device *mtl_device,
|
||||
const std::string &source_path) {
|
||||
std::string source = read_file(source_path);
|
||||
@autoreleasepool {
|
||||
NSString *src = [NSString stringWithUTF8String:source.c_str()];
|
||||
MTLCompileOptions *opts = [[MTLCompileOptions alloc] init];
|
||||
opts.languageVersion = MTLLanguageVersion4_0;
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLLibrary> lib_objc = [device_objc newLibraryWithSource:src
|
||||
options:opts
|
||||
error:&error];
|
||||
if (!lib_objc) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("Metal compile failed (" + source_path +
|
||||
"): " + err);
|
||||
}
|
||||
return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc);
|
||||
}
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device,
|
||||
const std::string &name) {
|
||||
@autoreleasepool {
|
||||
NSString *fn_name = [NSString stringWithUTF8String:name.c_str()];
|
||||
// Try w4a8 lib first, then quantize lib
|
||||
id<MTLLibrary> lib_objc = (__bridge id<MTLLibrary>)w4a8_lib_;
|
||||
id<MTLFunction> func = [lib_objc newFunctionWithName:fn_name];
|
||||
if (!func) {
|
||||
lib_objc = (__bridge id<MTLLibrary>)quantize_lib_;
|
||||
func = [lib_objc newFunctionWithName:fn_name];
|
||||
}
|
||||
if (!func) {
|
||||
throw std::runtime_error("Kernel not found: " + name);
|
||||
}
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLComputePipelineState> pso =
|
||||
[device_objc newComputePipelineStateWithFunction:func error:&error];
|
||||
if (!pso) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("Pipeline failed for " + name + ": " + err);
|
||||
}
|
||||
return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain(
|
||||
pso);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void W4A8Linear::eval_gpu(const std::vector<array> &inputs,
|
||||
std::vector<array> &outputs) {
|
||||
auto &x = inputs[0]; // [M, K] float16
|
||||
auto &packed_w = inputs[1]; // [K/2, N] uint8
|
||||
auto &scale_w = inputs[2]; // [N] float32
|
||||
auto &out = outputs[0]; // [M, N] float16
|
||||
|
||||
uint32_t M = static_cast<uint32_t>(x.shape(0));
|
||||
uint32_t K = static_cast<uint32_t>(x.shape(1));
|
||||
uint32_t N = static_cast<uint32_t>(packed_w.shape(1));
|
||||
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
|
||||
auto &cache = W4A8PipelineCache::instance();
|
||||
cache.ensure_init(kernel_dir_);
|
||||
auto &pool = cache.scratch();
|
||||
|
||||
size_t a_bytes = static_cast<size_t>(M) * K;
|
||||
size_t sa_bytes = static_cast<size_t>(M) * sizeof(float);
|
||||
|
||||
allocator::Buffer a_int8_buf = pool.get_a_int8(a_bytes);
|
||||
allocator::Buffer sa_buf = pool.get_scale_a(sa_bytes);
|
||||
|
||||
auto noop = [](allocator::Buffer) {};
|
||||
array a_int8(a_int8_buf, {static_cast<int>(M), static_cast<int>(K)}, int8,
|
||||
noop);
|
||||
array sa(sa_buf, {static_cast<int>(M)}, float32, noop);
|
||||
|
||||
auto &s = stream();
|
||||
// auto& dev removed — use free function get_command_encoder
|
||||
auto &enc = metal::get_command_encoder(s);
|
||||
|
||||
// DISPATCH 1: quantize activations (reuse w8a8_quantize.metal)
|
||||
{
|
||||
auto *pso = cache.get("quantize_per_token");
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
enc.set_input_array(x, 0);
|
||||
enc.set_output_array(a_int8, 1);
|
||||
enc.set_output_array(sa, 2);
|
||||
enc.set_bytes(M, 3);
|
||||
enc.set_bytes(K, 4);
|
||||
uint32_t tg = std::min(256u, std::max(32u, ((K + 31) / 32) * 32));
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(M, 1, 1),
|
||||
MTL::Size::Make(tg, 1, 1));
|
||||
}
|
||||
|
||||
enc.barrier();
|
||||
|
||||
// DISPATCH 2: W4A8 matmul with fused dequant
|
||||
{
|
||||
bool use_small = (M <= 64);
|
||||
const char *kname = use_small ? "w4a8_matmul_fused_dequant_small"
|
||||
: "w4a8_matmul_fused_dequant";
|
||||
uint32_t BM = use_small ? TileSmall::BM : TileLarge::BM;
|
||||
uint32_t BN = use_small ? TileSmall::BN : TileLarge::BN;
|
||||
uint32_t threads = use_small ? TileSmall::THREADS : TileLarge::THREADS;
|
||||
|
||||
auto *pso = cache.get(kname);
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
|
||||
uint32_t tiles_n = (N + BN - 1) / BN;
|
||||
uint32_t tiles_m = (M + BM - 1) / BM;
|
||||
uint32_t swizzle_log;
|
||||
if (tiles_m <= 3) {
|
||||
swizzle_log = 0;
|
||||
} else if (tiles_m <= 6) {
|
||||
swizzle_log = 1;
|
||||
} else {
|
||||
swizzle_log = 2;
|
||||
}
|
||||
uint32_t tile = 1u << swizzle_log;
|
||||
uint32_t grid_x = tiles_n * tile;
|
||||
uint32_t grid_y = (tiles_m + tile - 1) / tile;
|
||||
|
||||
enc.set_input_array(a_int8, 0); // A: INT8 activations
|
||||
enc.set_input_array(packed_w, 1); // B: packed W4
|
||||
enc.set_output_array(out, 2); // C: FP16 output
|
||||
enc.set_bytes(M, 3);
|
||||
enc.set_bytes(N, 4);
|
||||
enc.set_bytes(K, 5);
|
||||
enc.set_input_array(sa, 6); // scale_a
|
||||
enc.set_input_array(scale_w, 7); // scale_w
|
||||
enc.set_bytes(swizzle_log, 8);
|
||||
enc.set_bytes(tiles_m, 9);
|
||||
enc.set_bytes(tiles_n, 10);
|
||||
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(grid_x, grid_y, 1),
|
||||
MTL::Size::Make(threads, 1, 1));
|
||||
}
|
||||
}
|
||||
|
||||
array w4a8_linear(const array &x, const array &packed_w, const array &scale_w,
|
||||
const std::string &kernel_dir, StreamOrDevice s) {
|
||||
if (x.ndim() != 2) {
|
||||
throw std::invalid_argument("w4a8_linear: x must be 2D [M,K]");
|
||||
}
|
||||
if (packed_w.ndim() != 2) {
|
||||
throw std::invalid_argument("w4a8_linear: packed_w must be 2D [K/2,N]");
|
||||
}
|
||||
|
||||
int M = x.shape(0);
|
||||
int N = packed_w.shape(1);
|
||||
auto stream = to_stream(s);
|
||||
|
||||
return array(
|
||||
{M, N}, float16, std::make_shared<W4A8Linear>(stream, kernel_dir),
|
||||
{astype(x, float16, stream), astype(packed_w, mlx::core::uint8, stream),
|
||||
astype(scale_w, float32, stream)});
|
||||
}
|
||||
|
||||
} // namespace cider
|
||||
@@ -0,0 +1,344 @@
|
||||
// W8A8 Linear V9 — Unified prefill (GEMM) + decode (FP MV) with bias
|
||||
//
|
||||
// M > 1: quantize activation → INT8 GEMM → y = acc * scale_a * scale_w + bias
|
||||
// M == 1: FP MV → y = scale_w * sum(float(w_int8) * float(x)) + bias
|
||||
// Per-channel: single scale per row, no group switching overhead.
|
||||
|
||||
#include "w8a8_primitive.h"
|
||||
#include "mlx/backend/metal/device.h"
|
||||
#include "mlx/backend/metal/utils.h"
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <Metal/Metal.h>
|
||||
|
||||
#include <fstream>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
#include <unordered_map>
|
||||
|
||||
namespace cider {
|
||||
|
||||
using namespace mlx::core;
|
||||
|
||||
struct TileLarge {
|
||||
static constexpr uint32_t BM = 128, BN = 128, THREADS = 512;
|
||||
};
|
||||
struct TileSmall {
|
||||
static constexpr uint32_t BM = 32, BN = 128, THREADS = 128;
|
||||
};
|
||||
|
||||
static std::string read_file(const std::string &path) {
|
||||
std::ifstream f(path);
|
||||
if (!f.is_open()) {
|
||||
throw std::runtime_error("Cannot open: " + path);
|
||||
}
|
||||
std::stringstream ss;
|
||||
ss << f.rdbuf();
|
||||
f.close();
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
// ── Pipeline cache ──────────────────────────────────────────────
|
||||
class PipelineCache {
|
||||
public:
|
||||
static PipelineCache &instance() {
|
||||
static PipelineCache cache;
|
||||
return cache;
|
||||
}
|
||||
|
||||
void ensure_init(const std::string &kernel_dir) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
if (initialized_ && kernel_dir_ == kernel_dir) {
|
||||
return;
|
||||
}
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *mtl_device = dev.mtl_device();
|
||||
matmul_lib_ = compile_source(mtl_device, kernel_dir + "/w8a8_matmul.metal");
|
||||
quantize_lib_ = compile_source(mtl_device, kernel_dir + "/w8a8_quantize.metal");
|
||||
mv_lib_ = compile_source(mtl_device, kernel_dir + "/w8a8_int8_mv.metal");
|
||||
kernel_dir_ = kernel_dir;
|
||||
pipelines_.clear();
|
||||
initialized_ = true;
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *get(const std::string &name) {
|
||||
std::lock_guard<std::mutex> lock(mutex_);
|
||||
auto it = pipelines_.find(name);
|
||||
if (it != pipelines_.end()) {
|
||||
return it->second;
|
||||
}
|
||||
auto &dev = metal::device(mlx::core::Device::gpu);
|
||||
auto *pso = make_pipeline(dev.mtl_device(), name);
|
||||
pipelines_[name] = pso;
|
||||
return pso;
|
||||
}
|
||||
|
||||
private:
|
||||
PipelineCache() = default;
|
||||
bool initialized_ = false;
|
||||
std::string kernel_dir_;
|
||||
std::unordered_map<std::string, MTL::ComputePipelineState *> pipelines_;
|
||||
MTL::Library *matmul_lib_ = nullptr;
|
||||
MTL::Library *quantize_lib_ = nullptr;
|
||||
MTL::Library *mv_lib_ = nullptr;
|
||||
std::mutex mutex_;
|
||||
|
||||
MTL::Library *compile_source(MTL::Device *mtl_device,
|
||||
const std::string &source_path) {
|
||||
std::string source = read_file(source_path);
|
||||
@autoreleasepool {
|
||||
NSString *src = [NSString stringWithUTF8String:source.c_str()];
|
||||
MTLCompileOptions *opts = [[MTLCompileOptions alloc] init];
|
||||
opts.languageVersion = MTLLanguageVersion4_0;
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLLibrary> lib_objc = [device_objc newLibraryWithSource:src
|
||||
options:opts
|
||||
error:&error];
|
||||
if (!lib_objc) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("Metal compile failed (" + source_path +
|
||||
"): " + err);
|
||||
}
|
||||
return (__bridge MTL::Library *)(void *)CFBridgingRetain(lib_objc);
|
||||
}
|
||||
}
|
||||
|
||||
MTL::ComputePipelineState *make_pipeline(MTL::Device *mtl_device,
|
||||
const std::string &name) {
|
||||
@autoreleasepool {
|
||||
NSString *fn_name = [NSString stringWithUTF8String:name.c_str()];
|
||||
id<MTLFunction> func = nil;
|
||||
for (auto *lib : {matmul_lib_, quantize_lib_, mv_lib_}) {
|
||||
if (!lib) continue;
|
||||
id<MTLLibrary> lib_objc = (__bridge id<MTLLibrary>)lib;
|
||||
func = [lib_objc newFunctionWithName:fn_name];
|
||||
if (func) break;
|
||||
}
|
||||
if (!func) {
|
||||
throw std::runtime_error("Kernel not found: " + name);
|
||||
}
|
||||
NSError *error = nil;
|
||||
id<MTLDevice> device_objc = (__bridge id<MTLDevice>)mtl_device;
|
||||
id<MTLComputePipelineState> pso =
|
||||
[device_objc newComputePipelineStateWithFunction:func error:&error];
|
||||
if (!pso) {
|
||||
std::string err =
|
||||
error ? [[error localizedDescription] UTF8String] : "Unknown error";
|
||||
throw std::runtime_error("Pipeline failed for " + name + ": " + err);
|
||||
}
|
||||
return (__bridge MTL::ComputePipelineState *)(void *)CFBridgingRetain(
|
||||
pso);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void W8A8Linear::eval_gpu(const std::vector<array> &inputs,
|
||||
std::vector<array> &outputs) {
|
||||
auto &x = inputs[0]; // [M, K] float16
|
||||
auto &w = inputs[1]; // [N, K] int8
|
||||
auto &scale_w = inputs[2]; // [N] float32
|
||||
auto &bias = inputs[3]; // [N] float16
|
||||
auto &out = outputs[0]; // [M, N] float16
|
||||
|
||||
uint32_t M = static_cast<uint32_t>(x.shape(0));
|
||||
uint32_t N = static_cast<uint32_t>(w.shape(0));
|
||||
uint32_t K = static_cast<uint32_t>(w.shape(1));
|
||||
|
||||
out.set_data(allocator::malloc(out.nbytes()));
|
||||
|
||||
auto &cache = PipelineCache::instance();
|
||||
cache.ensure_init(kernel_dir_);
|
||||
|
||||
auto &s = stream();
|
||||
auto &enc = metal::get_command_encoder(s);
|
||||
|
||||
if (M == 1) {
|
||||
// ── Decode path: FP MV (no activation quantization) ──
|
||||
// w8a8_int8_mv.metal buffer layout:
|
||||
// x[0], W[1], y[2], scale_w[3], N[4], K[5], bias[6]
|
||||
auto *pso = cache.get("w8a8_int8_mv");
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
enc.set_input_array(x, 0); // [1, K] float16
|
||||
enc.set_input_array(w, 1); // [N, K] int8
|
||||
enc.set_output_array(out, 2); // [1, N] float16
|
||||
enc.set_input_array(scale_w, 3); // [N] float32
|
||||
enc.set_bytes(N, 4);
|
||||
enc.set_bytes(K, 5);
|
||||
enc.set_input_array(bias, 6); // [N] float16
|
||||
|
||||
constexpr uint32_t NUM_SIMDGROUPS = 2;
|
||||
constexpr uint32_t RESULTS_PER_SG = 4;
|
||||
constexpr uint32_t SIMD_SIZE = 32;
|
||||
uint32_t rows_per_tg = NUM_SIMDGROUPS * RESULTS_PER_SG; // 16
|
||||
uint32_t grid_n = (N + rows_per_tg - 1) / rows_per_tg;
|
||||
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(grid_n, 1, 1),
|
||||
MTL::Size::Make(SIMD_SIZE * NUM_SIMDGROUPS, 1, 1));
|
||||
} else {
|
||||
// ── Prefill path: quantize + INT8 GEMM ──
|
||||
size_t a_bytes = static_cast<size_t>(M) * K;
|
||||
size_t sa_bytes = static_cast<size_t>(M) * sizeof(float);
|
||||
|
||||
array a_int8({static_cast<int>(M), static_cast<int>(K)}, int8, nullptr, {});
|
||||
a_int8.set_data(allocator::malloc(a_bytes));
|
||||
|
||||
array sa({static_cast<int>(M)}, float32, nullptr, {});
|
||||
sa.set_data(allocator::malloc(sa_bytes));
|
||||
|
||||
// DISPATCH 1: quantize_per_token
|
||||
{
|
||||
auto *pso = cache.get("quantize_per_token");
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
enc.set_input_array(x, 0);
|
||||
enc.set_output_array(a_int8, 1);
|
||||
enc.set_output_array(sa, 2);
|
||||
enc.set_bytes(M, 3);
|
||||
enc.set_bytes(K, 4);
|
||||
uint32_t tg = std::min(256u, std::max(32u, ((K + 31) / 32) * 32));
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(M, 1, 1),
|
||||
MTL::Size::Make(tg, 1, 1));
|
||||
}
|
||||
|
||||
enc.barrier();
|
||||
|
||||
// DISPATCH 2: matmul_dequant
|
||||
{
|
||||
bool use_small = (M <= 64);
|
||||
const char *kname = use_small ? "w8a8_matmul_fused_dequant_small"
|
||||
: "w8a8_matmul_fused_dequant";
|
||||
uint32_t BM = use_small ? TileSmall::BM : TileLarge::BM;
|
||||
uint32_t BN = use_small ? TileSmall::BN : TileLarge::BN;
|
||||
uint32_t threads = use_small ? TileSmall::THREADS : TileLarge::THREADS;
|
||||
|
||||
auto *pso = cache.get(kname);
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
|
||||
uint32_t tiles_n = (N + BN - 1) / BN;
|
||||
uint32_t tiles_m = (M + BM - 1) / BM;
|
||||
uint32_t swizzle_log;
|
||||
if (tiles_m <= 3) {
|
||||
swizzle_log = 0;
|
||||
} else if (tiles_m <= 6) {
|
||||
swizzle_log = 1;
|
||||
} else {
|
||||
swizzle_log = 2;
|
||||
}
|
||||
uint32_t tile = 1u << swizzle_log;
|
||||
uint32_t grid_x = tiles_n * tile;
|
||||
uint32_t grid_y = (tiles_m + tile - 1) / tile;
|
||||
|
||||
enc.set_input_array(a_int8, 0);
|
||||
enc.set_input_array(w, 1);
|
||||
enc.set_output_array(out, 2);
|
||||
enc.set_bytes(M, 3);
|
||||
enc.set_bytes(N, 4);
|
||||
enc.set_bytes(K, 5);
|
||||
enc.set_input_array(sa, 6);
|
||||
enc.set_input_array(scale_w, 7);
|
||||
enc.set_bytes(swizzle_log, 8);
|
||||
enc.set_bytes(tiles_m, 9);
|
||||
enc.set_bytes(tiles_n, 10);
|
||||
enc.set_input_array(bias, 11);
|
||||
|
||||
enc.dispatch_threadgroups(MTL::Size::Make(grid_x, grid_y, 1),
|
||||
MTL::Size::Make(threads, 1, 1));
|
||||
}
|
||||
|
||||
enc.add_temporary(a_int8);
|
||||
enc.add_temporary(sa);
|
||||
}
|
||||
}
|
||||
|
||||
array perchannel_linear(const array &x, const array &w, const array &scale_w,
|
||||
const array &bias, const std::string &kernel_dir,
|
||||
StreamOrDevice s) {
|
||||
if (x.ndim() != 2) {
|
||||
throw std::invalid_argument("perchannel_linear: x must be 2D [M,K]");
|
||||
}
|
||||
if (w.ndim() != 2) {
|
||||
throw std::invalid_argument("perchannel_linear: w must be 2D [N,K]");
|
||||
}
|
||||
|
||||
int M = x.shape(0);
|
||||
int N = w.shape(0);
|
||||
auto stream = to_stream(s);
|
||||
|
||||
auto result =
|
||||
array({M, N}, float16, std::make_shared<W8A8Linear>(stream, kernel_dir),
|
||||
{astype(x, float16, stream), astype(w, int8, stream),
|
||||
astype(scale_w, float32, stream), astype(bias, float16, stream)});
|
||||
if (x.dtype() == bfloat16) {
|
||||
return astype(result, bfloat16, stream);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// ── Int8MatMulInt32 primitive (raw INT32 output, no dequant) ────
|
||||
|
||||
void Int8MatMulInt32::eval_gpu(const std::vector<mx::array> &inputs,
|
||||
std::vector<mx::array> &outputs) {
|
||||
auto &a = inputs[0]; // [M, K] int8
|
||||
auto &b = inputs[1]; // [N, K] int8
|
||||
|
||||
uint32_t M = static_cast<uint32_t>(a.shape(0));
|
||||
uint32_t N = static_cast<uint32_t>(b.shape(0));
|
||||
uint32_t K = static_cast<uint32_t>(b.shape(1));
|
||||
|
||||
auto &out = outputs[0];
|
||||
out.set_data(mx::allocator::malloc(out.nbytes()));
|
||||
|
||||
auto &cache = PipelineCache::instance();
|
||||
cache.ensure_init(kernel_dir_);
|
||||
|
||||
bool use_small = (M <= 64);
|
||||
auto *pso = use_small ? cache.get("int8_matmul_int32_small")
|
||||
: cache.get("int8_matmul_int32");
|
||||
|
||||
uint32_t BM = use_small ? 32 : 128;
|
||||
uint32_t BN = 128;
|
||||
uint32_t threads = use_small ? 128 : 512;
|
||||
|
||||
uint32_t tiles_n = (N + BN - 1) / BN;
|
||||
uint32_t tiles_m = (M + BM - 1) / BM;
|
||||
uint32_t swizzle_log;
|
||||
if (tiles_m <= 3) {
|
||||
swizzle_log = 0;
|
||||
} else if (tiles_m <= 6) {
|
||||
swizzle_log = 1;
|
||||
} else {
|
||||
swizzle_log = 2;
|
||||
}
|
||||
uint32_t tile = 1u << swizzle_log;
|
||||
uint32_t grid_x = tiles_n * tile;
|
||||
uint32_t grid_y = (tiles_m + tile - 1) / tile;
|
||||
|
||||
auto &s = stream();
|
||||
auto &enc = metal::get_command_encoder(s);
|
||||
enc.set_compute_pipeline_state(pso);
|
||||
enc.set_input_array(a, 0);
|
||||
enc.set_input_array(b, 1);
|
||||
enc.set_output_array(out, 2);
|
||||
enc.set_bytes(&M, sizeof(M), 3);
|
||||
enc.set_bytes(&N, sizeof(N), 4);
|
||||
enc.set_bytes(&K, sizeof(K), 5);
|
||||
enc.set_bytes(&swizzle_log, sizeof(swizzle_log), 6);
|
||||
enc.set_bytes(&tiles_m, sizeof(tiles_m), 7);
|
||||
enc.set_bytes(&tiles_n, sizeof(tiles_n), 8);
|
||||
enc.dispatch_threadgroups(MTL::Size(grid_x, grid_y, 1),
|
||||
MTL::Size(threads, 1, 1));
|
||||
}
|
||||
|
||||
mx::array int8_matmul_int32(const mx::array &a, const mx::array &b,
|
||||
const std::string &kernel_dir,
|
||||
mx::StreamOrDevice s) {
|
||||
uint32_t M = static_cast<uint32_t>(a.shape(0));
|
||||
uint32_t N = static_cast<uint32_t>(b.shape(0));
|
||||
return mx::array(
|
||||
{static_cast<int>(M), static_cast<int>(N)}, mx::int32,
|
||||
std::make_shared<Int8MatMulInt32>(mx::to_stream(s), kernel_dir),
|
||||
{astype(a, int8, mx::to_stream(s)), astype(b, int8, mx::to_stream(s))});
|
||||
}
|
||||
|
||||
} // namespace cider
|
||||
Reference in New Issue
Block a user