665 lines
25 KiB
C++
665 lines
25 KiB
C++
#include "fused_moe.h"
|
|
|
|
struct FusedParams {
|
|
int experts;
|
|
int hidden;
|
|
int inter;
|
|
int group_size;
|
|
int bits;
|
|
int k;
|
|
};
|
|
|
|
struct MappedProj {
|
|
void* w{MAP_FAILED};
|
|
void* s{MAP_FAILED};
|
|
void* b{MAP_FAILED};
|
|
size_t wn{0}, sn{0}, bn{0};
|
|
int wfd{-1}, sfd{-1}, bfd{-1};
|
|
};
|
|
|
|
struct ProjectionShape {
|
|
int out_dim;
|
|
int in_dim;
|
|
size_t weight_per;
|
|
size_t scale_per;
|
|
};
|
|
|
|
static size_t file_size(int fd) {
|
|
struct stat st {};
|
|
if (fstat(fd, &st) != 0) {
|
|
throw std::runtime_error("fstat failed");
|
|
}
|
|
return static_cast<size_t>(st.st_size);
|
|
}
|
|
|
|
static void* map_file(const std::string& path, size_t& nbytes, int& fd) {
|
|
fd = open(path.c_str(), O_RDONLY);
|
|
if (fd < 0) {
|
|
throw std::runtime_error("failed to open " + path);
|
|
}
|
|
nbytes = file_size(fd);
|
|
void* p = mmap(nullptr, nbytes, PROT_READ, MAP_PRIVATE, fd, 0);
|
|
if (p == MAP_FAILED) {
|
|
throw std::runtime_error("failed to mmap " + path);
|
|
}
|
|
return p;
|
|
}
|
|
|
|
static MappedProj map_proj(const std::string& dir, int layer, const std::string& proj) {
|
|
MappedProj m;
|
|
std::string layer_name = "layer" + std::string(layer < 10 ? "0" : "") + std::to_string(layer);
|
|
std::string base = dir + "/" + layer_name + "." + proj;
|
|
m.w = map_file(base + ".weight.bin", m.wn, m.wfd);
|
|
m.s = map_file(base + ".scales.bin", m.sn, m.sfd);
|
|
m.b = map_file(base + ".biases.bin", m.bn, m.bfd);
|
|
return m;
|
|
}
|
|
|
|
static void unmap_proj(MappedProj& m) {
|
|
if (m.w != MAP_FAILED) munmap(m.w, m.wn);
|
|
if (m.s != MAP_FAILED) munmap(m.s, m.sn);
|
|
if (m.b != MAP_FAILED) munmap(m.b, m.bn);
|
|
if (m.wfd >= 0) close(m.wfd);
|
|
if (m.sfd >= 0) close(m.sfd);
|
|
if (m.bfd >= 0) close(m.bfd);
|
|
}
|
|
|
|
static ProjectionShape proj_shape(int out_dim, int in_dim, int group, int bits) {
|
|
int words = (in_dim * bits) / 32;
|
|
int groups = in_dim / group;
|
|
return ProjectionShape{
|
|
out_dim,
|
|
in_dim,
|
|
static_cast<size_t>(out_dim) * words * sizeof(uint32_t),
|
|
static_cast<size_t>(out_dim) * groups * sizeof(uint16_t),
|
|
};
|
|
}
|
|
|
|
static uint16_t fp32_to_bf16(float x) {
|
|
uint32_t u = 0;
|
|
std::memcpy(&u, &x, sizeof(u));
|
|
return static_cast<uint16_t>(u >> 16);
|
|
}
|
|
|
|
static void fill_synthetic(
|
|
MTL::Buffer* wbuf,
|
|
MTL::Buffer* sbuf,
|
|
MTL::Buffer* bbuf,
|
|
uint32_t seed) {
|
|
std::mt19937 rng(seed);
|
|
auto* w = reinterpret_cast<uint32_t*>(wbuf->contents());
|
|
size_t wn = wbuf->length() / sizeof(uint32_t);
|
|
for (size_t i = 0; i < wn; ++i) w[i] = rng();
|
|
auto* s = reinterpret_cast<uint16_t*>(sbuf->contents());
|
|
auto* b = reinterpret_cast<uint16_t*>(bbuf->contents());
|
|
size_t sn = sbuf->length() / sizeof(uint16_t);
|
|
uint16_t scale = fp32_to_bf16(0.001f);
|
|
uint16_t bias = fp32_to_bf16(-0.032f);
|
|
for (size_t i = 0; i < sn; ++i) {
|
|
s[i] = scale;
|
|
b[i] = bias;
|
|
}
|
|
}
|
|
|
|
static void copy_projection(
|
|
const MappedProj& src,
|
|
const ProjectionShape& shape,
|
|
const std::vector<int>& ids,
|
|
MTL::Buffer* wbuf,
|
|
MTL::Buffer* sbuf,
|
|
MTL::Buffer* bbuf) {
|
|
auto* wd = static_cast<char*>(wbuf->contents());
|
|
auto* sd = static_cast<char*>(sbuf->contents());
|
|
auto* bd = static_cast<char*>(bbuf->contents());
|
|
auto* ws = static_cast<const char*>(src.w);
|
|
auto* ss = static_cast<const char*>(src.s);
|
|
auto* bs = static_cast<const char*>(src.b);
|
|
for (size_t local = 0; local < ids.size(); ++local) {
|
|
size_t expert = static_cast<size_t>(ids[local]);
|
|
std::memcpy(wd + local * shape.weight_per, ws + expert * shape.weight_per, shape.weight_per);
|
|
std::memcpy(sd + local * shape.scale_per, ss + expert * shape.scale_per, shape.scale_per);
|
|
std::memcpy(bd + local * shape.scale_per, bs + expert * shape.scale_per, shape.scale_per);
|
|
}
|
|
}
|
|
|
|
static std::string metal_source() {
|
|
return R"(
|
|
#include <metal_stdlib>
|
|
using namespace metal;
|
|
struct FusedParams { int experts; int hidden; int inter; int group_size; int bits; int k; };
|
|
inline float bf16_to_float(ushort v) { uint u = uint(v) << 16; return as_type<float>(u); }
|
|
inline float qvalue(const device uint* w, const device ushort* s, const device ushort* b,
|
|
uint expert, uint row, uint col, uint out_dim, uint in_dim,
|
|
uint group_size, uint bits) {
|
|
uint words_per_row = (in_dim * bits) / 32;
|
|
uint groups_per_row = in_dim / group_size;
|
|
uint bit_offset = col * bits;
|
|
uint word_idx = bit_offset / 32;
|
|
uint shift = bit_offset % 32;
|
|
uint base = (expert * out_dim + row) * words_per_row;
|
|
uint q = w[base + word_idx] >> shift;
|
|
if (shift + bits > 32) q |= w[base + word_idx + 1] << (32 - shift);
|
|
q &= (1u << bits) - 1u;
|
|
uint sb = (expert * out_dim + row) * groups_per_row + col / group_size;
|
|
return float(q) * bf16_to_float(s[sb]) + bf16_to_float(b[sb]);
|
|
}
|
|
kernel void synthetic_moe(
|
|
const device float* x [[buffer(0)]],
|
|
device float* y [[buffer(1)]],
|
|
const device float* scores [[buffer(2)]],
|
|
constant FusedParams& p [[buffer(3)]],
|
|
uint gid [[thread_position_in_grid]]) {
|
|
uint token = gid / p.hidden;
|
|
uint col = gid % p.hidden;
|
|
float acc = 0.0f;
|
|
for (int j = 0; j < p.k; ++j) {
|
|
acc += scores[token * p.k + j];
|
|
}
|
|
y[gid] = x[token * p.hidden + col] * acc;
|
|
}
|
|
kernel void fused_moe_pairs(
|
|
const device float* x [[buffer(0)]],
|
|
const device uint* gate_w [[buffer(1)]], const device ushort* gate_s [[buffer(2)]], const device ushort* gate_b [[buffer(3)]],
|
|
const device uint* up_w [[buffer(4)]], const device ushort* up_s [[buffer(5)]], const device ushort* up_b [[buffer(6)]],
|
|
const device uint* down_w [[buffer(7)]], const device ushort* down_s [[buffer(8)]], const device ushort* down_b [[buffer(9)]],
|
|
device float* y [[buffer(10)]], constant FusedParams& p [[buffer(11)]],
|
|
uint tid [[thread_position_in_threadgroup]], uint gid [[thread_position_in_grid]]) {
|
|
constexpr uint block_size = 512;
|
|
constexpr uint lanes_per_row = 16;
|
|
constexpr uint rows_per_step = block_size / lanes_per_row;
|
|
uint expert = gid / block_size;
|
|
if (expert >= uint(p.experts)) return;
|
|
uint token = expert / p.k;
|
|
uint local_row = tid / lanes_per_row;
|
|
uint row_lane = tid % lanes_per_row;
|
|
threadgroup float act[1024];
|
|
threadgroup float gate_part[block_size];
|
|
threadgroup float up_part[block_size];
|
|
for (uint row_base = 0; row_base < uint(p.inter); row_base += rows_per_step) {
|
|
uint row = row_base + local_row;
|
|
float gate_acc = 0.0f, up_acc = 0.0f;
|
|
if (row < uint(p.inter)) {
|
|
for (uint col = row_lane; col < uint(p.hidden); col += lanes_per_row) {
|
|
float xv = x[token * p.hidden + col];
|
|
gate_acc += qvalue(gate_w, gate_s, gate_b, expert, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv;
|
|
up_acc += qvalue(up_w, up_s, up_b, expert, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv;
|
|
}
|
|
}
|
|
gate_part[tid] = gate_acc;
|
|
up_part[tid] = up_acc;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
|
|
if (row_lane < stride) { gate_part[tid] += gate_part[tid + stride]; up_part[tid] += up_part[tid + stride]; }
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (row_lane == 0 && row < uint(p.inter)) {
|
|
float g = gate_part[tid];
|
|
act[row] = g * (1.0f / (1.0f + exp(-g))) * up_part[tid];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
for (uint row_base = 0; row_base < uint(p.hidden); row_base += rows_per_step) {
|
|
uint row = row_base + local_row;
|
|
float acc = 0.0f;
|
|
if (row < uint(p.hidden)) {
|
|
for (uint col = row_lane; col < uint(p.inter); col += lanes_per_row) {
|
|
acc += qvalue(down_w, down_s, down_b, expert, row, col, p.hidden, p.inter, p.group_size, p.bits) * act[col];
|
|
}
|
|
}
|
|
gate_part[tid] = acc;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
|
|
if (row_lane < stride) gate_part[tid] += gate_part[tid + stride];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (row_lane == 0 && row < uint(p.hidden)) {
|
|
y[expert * p.hidden + row] = gate_part[tid];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
}
|
|
kernel void fused_moe_slots_pairs(
|
|
const device float* x [[buffer(0)]],
|
|
const device uint* gate_w [[buffer(1)]], const device ushort* gate_s [[buffer(2)]], const device ushort* gate_b [[buffer(3)]],
|
|
const device uint* up_w [[buffer(4)]], const device ushort* up_s [[buffer(5)]], const device ushort* up_b [[buffer(6)]],
|
|
const device uint* down_w [[buffer(7)]], const device ushort* down_s [[buffer(8)]], const device ushort* down_b [[buffer(9)]],
|
|
device float* y [[buffer(10)]], constant FusedParams& p [[buffer(11)]],
|
|
const device uint* local_slots [[buffer(12)]],
|
|
uint tid [[thread_position_in_threadgroup]], uint gid [[thread_position_in_grid]]) {
|
|
constexpr uint block_size = 512;
|
|
constexpr uint lanes_per_row = 16;
|
|
constexpr uint rows_per_step = block_size / lanes_per_row;
|
|
uint active_idx = gid / block_size;
|
|
if (active_idx >= uint(p.experts)) return;
|
|
uint slot = local_slots[active_idx];
|
|
uint token = active_idx / p.k;
|
|
uint local_row = tid / lanes_per_row;
|
|
uint row_lane = tid % lanes_per_row;
|
|
threadgroup float act[1024];
|
|
threadgroup float gate_part[block_size];
|
|
threadgroup float up_part[block_size];
|
|
for (uint row_base = 0; row_base < uint(p.inter); row_base += rows_per_step) {
|
|
uint row = row_base + local_row;
|
|
float gate_acc = 0.0f, up_acc = 0.0f;
|
|
if (row < uint(p.inter)) {
|
|
for (uint col = row_lane; col < uint(p.hidden); col += lanes_per_row) {
|
|
float xv = x[token * p.hidden + col];
|
|
gate_acc += qvalue(gate_w, gate_s, gate_b, slot, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv;
|
|
up_acc += qvalue(up_w, up_s, up_b, slot, row, col, p.inter, p.hidden, p.group_size, p.bits) * xv;
|
|
}
|
|
}
|
|
gate_part[tid] = gate_acc;
|
|
up_part[tid] = up_acc;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
|
|
if (row_lane < stride) { gate_part[tid] += gate_part[tid + stride]; up_part[tid] += up_part[tid + stride]; }
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (row_lane == 0 && row < uint(p.inter)) {
|
|
float g = gate_part[tid];
|
|
act[row] = g * (1.0f / (1.0f + exp(-g))) * up_part[tid];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
for (uint row_base = 0; row_base < uint(p.hidden); row_base += rows_per_step) {
|
|
uint row = row_base + local_row;
|
|
float acc = 0.0f;
|
|
if (row < uint(p.hidden)) {
|
|
for (uint col = row_lane; col < uint(p.inter); col += lanes_per_row) {
|
|
acc += qvalue(down_w, down_s, down_b, slot, row, col, p.hidden, p.inter, p.group_size, p.bits) * act[col];
|
|
}
|
|
}
|
|
gate_part[tid] = acc;
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
for (uint stride = lanes_per_row / 2; stride > 0; stride >>= 1) {
|
|
if (row_lane < stride) gate_part[tid] += gate_part[tid + stride];
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
if (row_lane == 0 && row < uint(p.hidden)) {
|
|
y[active_idx * p.hidden + row] = gate_part[tid];
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
}
|
|
kernel void reduce_pairs(
|
|
const device float* pair_y [[buffer(0)]],
|
|
const device float* scores [[buffer(1)]],
|
|
device float* out [[buffer(2)]],
|
|
constant FusedParams& p [[buffer(3)]],
|
|
uint gid [[thread_position_in_grid]]) {
|
|
uint token = gid / p.hidden;
|
|
uint col = gid % p.hidden;
|
|
float acc = 0.0f;
|
|
for (int j = 0; j < p.k; ++j) {
|
|
uint pair = token * p.k + uint(j);
|
|
acc += pair_y[pair * p.hidden + col] * scores[pair];
|
|
}
|
|
out[gid] = acc;
|
|
}
|
|
)";
|
|
}
|
|
|
|
class FusedMoePrimitive : public mx::Primitive {
|
|
public:
|
|
FusedMoePrimitive(
|
|
mx::Stream stream,
|
|
std::string compute_dir,
|
|
int layer,
|
|
int hidden,
|
|
int inter,
|
|
int group,
|
|
int bits,
|
|
int num_experts,
|
|
bool synthetic,
|
|
std::vector<int> expert_ids)
|
|
: Primitive(stream),
|
|
compute_dir_(std::move(compute_dir)),
|
|
layer_(layer),
|
|
hidden_(hidden),
|
|
inter_(inter),
|
|
group_(group),
|
|
bits_(bits),
|
|
num_experts_(num_experts),
|
|
synthetic_(synthetic),
|
|
expert_ids_(std::move(expert_ids)) {
|
|
if (!synthetic_) {
|
|
gate_ = map_proj(compute_dir_, layer_, "gate_proj");
|
|
up_ = map_proj(compute_dir_, layer_, "up_proj");
|
|
down_ = map_proj(compute_dir_, layer_, "down_proj");
|
|
}
|
|
}
|
|
|
|
~FusedMoePrimitive() override {
|
|
if (!synthetic_) {
|
|
unmap_proj(gate_);
|
|
unmap_proj(up_);
|
|
unmap_proj(down_);
|
|
}
|
|
}
|
|
|
|
const char* name() const override { return "FusedMoePrimitive"; }
|
|
|
|
void eval_cpu(const std::vector<mx::array>&, std::vector<mx::array>&) override {
|
|
throw std::runtime_error("FusedMoePrimitive only supports GPU evaluation");
|
|
}
|
|
|
|
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
|
|
auto& x = inputs[0];
|
|
auto& scores = inputs[1];
|
|
auto& out = outputs[0];
|
|
int active = static_cast<int>(expert_ids_.size());
|
|
int tokens = static_cast<int>(x.size() / hidden_);
|
|
int k = active / std::max(1, tokens);
|
|
if (active <= 0 || tokens <= 0 || active != static_cast<int>(scores.size())) {
|
|
throw std::runtime_error("FusedMoePrimitive shape mismatch");
|
|
}
|
|
out.set_data(mx::allocator::malloc(out.nbytes()));
|
|
auto& d = mx::metal::device(stream().device);
|
|
auto& enc = mx::metal::get_command_encoder(stream());
|
|
auto lib = d.get_library("native_moe_mlx_ext", []() { return metal_source(); });
|
|
if (synthetic_) {
|
|
auto synthetic_kernel = d.get_kernel("synthetic_moe", lib);
|
|
enc.set_compute_pipeline_state(synthetic_kernel);
|
|
enc.set_input_array(x, 0);
|
|
enc.set_output_array(out, 1);
|
|
enc.set_input_array(scores, 2);
|
|
FusedParams params{active, hidden_, inter_, group_, bits_, k};
|
|
enc.set_bytes(params, 3);
|
|
enc.dispatch_threads(MTL::Size(out.size(), 1, 1), MTL::Size(std::min<size_t>(out.size(), 256), 1, 1));
|
|
return;
|
|
}
|
|
auto fused = d.get_kernel("fused_moe_pairs", lib);
|
|
auto reduce = d.get_kernel("reduce_pairs", lib);
|
|
|
|
ProjectionShape gu = proj_shape(inter_, hidden_, group_, bits_);
|
|
ProjectionShape down = proj_shape(hidden_, inter_, group_, bits_);
|
|
size_t gu_w_bytes = static_cast<size_t>(active) * gu.weight_per;
|
|
size_t gu_s_bytes = static_cast<size_t>(active) * gu.scale_per;
|
|
size_t down_w_bytes = static_cast<size_t>(active) * down.weight_per;
|
|
size_t down_s_bytes = static_cast<size_t>(active) * down.scale_per;
|
|
buffers_.clear();
|
|
auto make_buffer = [&](size_t nbytes) -> MTL::Buffer* {
|
|
buffers_.push_back(NS::TransferPtr(d.mtl_device()->newBuffer(nbytes, MTL::ResourceStorageModeShared)));
|
|
return buffers_.back().get();
|
|
};
|
|
MTL::Buffer* gw = make_buffer(gu_w_bytes);
|
|
MTL::Buffer* gs = make_buffer(gu_s_bytes);
|
|
MTL::Buffer* gb = make_buffer(gu_s_bytes);
|
|
MTL::Buffer* uw = make_buffer(gu_w_bytes);
|
|
MTL::Buffer* us = make_buffer(gu_s_bytes);
|
|
MTL::Buffer* ub = make_buffer(gu_s_bytes);
|
|
MTL::Buffer* dw = make_buffer(down_w_bytes);
|
|
MTL::Buffer* ds = make_buffer(down_s_bytes);
|
|
MTL::Buffer* db = make_buffer(down_s_bytes);
|
|
if (synthetic_) {
|
|
fill_synthetic(gw, gs, gb, 1);
|
|
fill_synthetic(uw, us, ub, 2);
|
|
fill_synthetic(dw, ds, db, 3);
|
|
} else {
|
|
copy_projection(gate_, gu, expert_ids_, gw, gs, gb);
|
|
copy_projection(up_, gu, expert_ids_, uw, us, ub);
|
|
copy_projection(down_, down, expert_ids_, dw, ds, db);
|
|
}
|
|
|
|
auto pair_out = mx::zeros(mx::Shape{active, hidden_}, mx::float32, stream());
|
|
pair_out.set_data(mx::allocator::malloc(pair_out.nbytes()));
|
|
|
|
enc.set_compute_pipeline_state(fused);
|
|
enc.set_input_array(x, 0);
|
|
enc.set_buffer(gw, 1);
|
|
enc.set_buffer(gs, 2);
|
|
enc.set_buffer(gb, 3);
|
|
enc.set_buffer(uw, 4);
|
|
enc.set_buffer(us, 5);
|
|
enc.set_buffer(ub, 6);
|
|
enc.set_buffer(dw, 7);
|
|
enc.set_buffer(ds, 8);
|
|
enc.set_buffer(db, 9);
|
|
enc.set_output_array(out, 10);
|
|
FusedParams params{active, hidden_, inter_, group_, bits_, k};
|
|
enc.set_bytes(params, 11);
|
|
enc.set_input_array(scores, 12);
|
|
enc.add_temporary(x);
|
|
enc.add_temporary(scores);
|
|
enc.dispatch_threads(MTL::Size(static_cast<size_t>(active) * 512, 1, 1), MTL::Size(512, 1, 1));
|
|
}
|
|
|
|
private:
|
|
std::string compute_dir_;
|
|
int layer_;
|
|
int hidden_;
|
|
int inter_;
|
|
int group_;
|
|
int bits_;
|
|
int num_experts_;
|
|
bool synthetic_;
|
|
std::vector<int> expert_ids_;
|
|
MappedProj gate_;
|
|
MappedProj up_;
|
|
MappedProj down_;
|
|
std::vector<NS::SharedPtr<MTL::Buffer>> buffers_;
|
|
};
|
|
|
|
class StagedFusedMoePrimitive : public mx::Primitive {
|
|
public:
|
|
StagedFusedMoePrimitive(mx::Stream stream, int hidden, int inter, int group, int bits)
|
|
: Primitive(stream), hidden_(hidden), inter_(inter), group_(group), bits_(bits) {}
|
|
|
|
const char* name() const override { return "StagedFusedMoePrimitive"; }
|
|
|
|
void eval_cpu(const std::vector<mx::array>&, std::vector<mx::array>&) override {
|
|
throw std::runtime_error("StagedFusedMoePrimitive only supports GPU evaluation");
|
|
}
|
|
|
|
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
|
|
auto& x = inputs[0];
|
|
auto& scores = inputs[1];
|
|
auto& out = outputs[0];
|
|
int tokens = static_cast<int>(x.size() / hidden_);
|
|
int active = static_cast<int>(scores.size());
|
|
int k = active / std::max(1, tokens);
|
|
if (active <= 0 || tokens <= 0 || active != static_cast<int>(scores.size())) {
|
|
throw std::runtime_error("StagedFusedMoePrimitive shape mismatch");
|
|
}
|
|
out.set_data(mx::allocator::malloc(out.nbytes()));
|
|
auto& d = mx::metal::device(stream().device);
|
|
auto& enc = mx::metal::get_command_encoder(stream());
|
|
auto lib = d.get_library("native_moe_mlx_ext", []() { return metal_source(); });
|
|
auto fused = d.get_kernel("fused_moe_pairs", lib);
|
|
auto reduce = d.get_kernel("reduce_pairs", lib);
|
|
auto pair_out = mx::zeros(mx::Shape{active, hidden_}, mx::float32, stream());
|
|
pair_out.set_data(mx::allocator::malloc(pair_out.nbytes()));
|
|
enc.set_compute_pipeline_state(fused);
|
|
enc.set_input_array(x, 0);
|
|
enc.set_input_array(inputs[2], 1);
|
|
enc.set_input_array(inputs[3], 2);
|
|
enc.set_input_array(inputs[4], 3);
|
|
enc.set_input_array(inputs[5], 4);
|
|
enc.set_input_array(inputs[6], 5);
|
|
enc.set_input_array(inputs[7], 6);
|
|
enc.set_input_array(inputs[8], 7);
|
|
enc.set_input_array(inputs[9], 8);
|
|
enc.set_input_array(inputs[10], 9);
|
|
enc.set_output_array(pair_out, 10);
|
|
FusedParams params{active, hidden_, inter_, group_, bits_, k};
|
|
enc.set_bytes(params, 11);
|
|
enc.dispatch_threads(MTL::Size(static_cast<size_t>(active) * 512, 1, 1), MTL::Size(512, 1, 1));
|
|
enc.set_compute_pipeline_state(reduce);
|
|
enc.set_input_array(pair_out, 0);
|
|
enc.set_input_array(scores, 1);
|
|
enc.set_output_array(out, 2);
|
|
enc.set_bytes(params, 3);
|
|
enc.add_temporary(pair_out);
|
|
enc.dispatch_threads(MTL::Size(out.size(), 1, 1), MTL::Size(std::min<size_t>(out.size(), 256), 1, 1));
|
|
}
|
|
|
|
private:
|
|
int hidden_;
|
|
int inter_;
|
|
int group_;
|
|
int bits_;
|
|
};
|
|
|
|
class SlotFusedMoePrimitive : public mx::Primitive {
|
|
public:
|
|
SlotFusedMoePrimitive(mx::Stream stream, int hidden, int inter, int group, int bits)
|
|
: Primitive(stream), hidden_(hidden), inter_(inter), group_(group), bits_(bits) {}
|
|
|
|
const char* name() const override { return "SlotFusedMoePrimitive"; }
|
|
|
|
void eval_cpu(const std::vector<mx::array>&, std::vector<mx::array>&) override {
|
|
throw std::runtime_error("SlotFusedMoePrimitive only supports GPU evaluation");
|
|
}
|
|
|
|
void eval_gpu(const std::vector<mx::array>& inputs, std::vector<mx::array>& outputs) override {
|
|
auto& x = inputs[0];
|
|
auto& local_slots = inputs[1];
|
|
auto& scores = inputs[2];
|
|
auto& out = outputs[0];
|
|
int tokens = static_cast<int>(x.size() / hidden_);
|
|
int active = static_cast<int>(scores.size());
|
|
int k = active / std::max(1, tokens);
|
|
if (active <= 0 || tokens <= 0 || active != static_cast<int>(local_slots.size())) {
|
|
throw std::runtime_error("SlotFusedMoePrimitive shape mismatch");
|
|
}
|
|
out.set_data(mx::allocator::malloc(out.nbytes()));
|
|
auto& d = mx::metal::device(stream().device);
|
|
auto& enc = mx::metal::get_command_encoder(stream());
|
|
auto lib = d.get_library("native_moe_mlx_ext", []() { return metal_source(); });
|
|
auto fused = d.get_kernel("fused_moe_slots_pairs", lib);
|
|
auto reduce = d.get_kernel("reduce_pairs", lib);
|
|
auto pair_out = mx::zeros(mx::Shape{active, hidden_}, mx::float32, stream());
|
|
pair_out.set_data(mx::allocator::malloc(pair_out.nbytes()));
|
|
enc.set_compute_pipeline_state(fused);
|
|
enc.set_input_array(x, 0);
|
|
enc.set_input_array(inputs[3], 1);
|
|
enc.set_input_array(inputs[4], 2);
|
|
enc.set_input_array(inputs[5], 3);
|
|
enc.set_input_array(inputs[6], 4);
|
|
enc.set_input_array(inputs[7], 5);
|
|
enc.set_input_array(inputs[8], 6);
|
|
enc.set_input_array(inputs[9], 7);
|
|
enc.set_input_array(inputs[10], 8);
|
|
enc.set_input_array(inputs[11], 9);
|
|
enc.set_output_array(pair_out, 10);
|
|
FusedParams params{active, hidden_, inter_, group_, bits_, k};
|
|
enc.set_bytes(params, 11);
|
|
enc.set_input_array(local_slots, 12);
|
|
enc.dispatch_threads(MTL::Size(static_cast<size_t>(active) * 512, 1, 1), MTL::Size(512, 1, 1));
|
|
enc.set_compute_pipeline_state(reduce);
|
|
enc.set_input_array(pair_out, 0);
|
|
enc.set_input_array(scores, 1);
|
|
enc.set_output_array(out, 2);
|
|
enc.set_bytes(params, 3);
|
|
enc.add_temporary(pair_out);
|
|
enc.dispatch_threads(MTL::Size(out.size(), 1, 1), MTL::Size(std::min<size_t>(out.size(), 256), 1, 1));
|
|
}
|
|
|
|
private:
|
|
int hidden_;
|
|
int inter_;
|
|
int group_;
|
|
int bits_;
|
|
};
|
|
mx::array fused_moe(
|
|
const mx::array& x,
|
|
const mx::array& expert_ids,
|
|
const mx::array& scores,
|
|
const std::string& compute_dir,
|
|
int layer,
|
|
int hidden,
|
|
int inter,
|
|
int group,
|
|
int bits,
|
|
int num_experts,
|
|
bool synthetic,
|
|
mx::StreamOrDevice s = {}) {
|
|
if (!synthetic) {
|
|
throw std::runtime_error(
|
|
"real mmap staging in MLX Primitive needs MLX-managed staging buffers; use synthetic for bridge tests");
|
|
}
|
|
auto ids = expert_ids;
|
|
ids.eval();
|
|
std::vector<int> expert_vec(ids.size());
|
|
const uint32_t* idp = ids.data<uint32_t>();
|
|
for (size_t i = 0; i < ids.size(); ++i) expert_vec[i] = static_cast<int>(idp[i]);
|
|
auto out_shape = x.shape();
|
|
out_shape.back() = hidden;
|
|
return mx::array(
|
|
out_shape,
|
|
mx::float32,
|
|
std::make_shared<FusedMoePrimitive>(
|
|
mx::to_stream(s),
|
|
compute_dir,
|
|
layer,
|
|
hidden,
|
|
inter,
|
|
group,
|
|
bits,
|
|
num_experts,
|
|
synthetic,
|
|
std::move(expert_vec)),
|
|
std::vector<mx::array>{x, scores});
|
|
}
|
|
|
|
mx::array fused_moe_staged(
|
|
const mx::array& x,
|
|
const mx::array& scores,
|
|
const mx::array& gate_w,
|
|
const mx::array& gate_s,
|
|
const mx::array& gate_b,
|
|
const mx::array& up_w,
|
|
const mx::array& up_s,
|
|
const mx::array& up_b,
|
|
const mx::array& down_w,
|
|
const mx::array& down_s,
|
|
const mx::array& down_b,
|
|
int hidden,
|
|
int inter,
|
|
int group,
|
|
int bits,
|
|
mx::StreamOrDevice s = {}) {
|
|
auto out_shape = x.shape();
|
|
out_shape.back() = hidden;
|
|
return mx::array(
|
|
out_shape,
|
|
mx::float32,
|
|
std::make_shared<StagedFusedMoePrimitive>(mx::to_stream(s), hidden, inter, group, bits),
|
|
std::vector<mx::array>{
|
|
x, scores,
|
|
gate_w, gate_s, gate_b,
|
|
up_w, up_s, up_b,
|
|
down_w, down_s, down_b});
|
|
}
|
|
|
|
mx::array fused_moe_slots(
|
|
const mx::array& x,
|
|
const mx::array& local_slots,
|
|
const mx::array& scores,
|
|
const mx::array& gate_w,
|
|
const mx::array& gate_s,
|
|
const mx::array& gate_b,
|
|
const mx::array& up_w,
|
|
const mx::array& up_s,
|
|
const mx::array& up_b,
|
|
const mx::array& down_w,
|
|
const mx::array& down_s,
|
|
const mx::array& down_b,
|
|
int hidden,
|
|
int inter,
|
|
int group,
|
|
int bits,
|
|
mx::StreamOrDevice s = {}) {
|
|
auto out_shape = x.shape();
|
|
out_shape.back() = hidden;
|
|
return mx::array(
|
|
out_shape,
|
|
mx::float32,
|
|
std::make_shared<SlotFusedMoePrimitive>(mx::to_stream(s), hidden, inter, group, bits),
|
|
std::vector<mx::array>{
|
|
x, local_slots, scores,
|
|
gate_w, gate_s, gate_b,
|
|
up_w, up_s, up_b,
|
|
down_w, down_s, down_b});
|
|
}
|