wehub-resource-sync
103 lines
3.0 KiB
C++
103 lines
3.0 KiB
C++
// 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
|