24 lines
1.3 KiB
C++
24 lines
1.3 KiB
C++
// fused MoE 计算：合成/真实 mmap staging、staged(预切片权重)、slot(local slot 索引)三条工厂。
|
||
// 实现见 fused_moe.cpp；这里只暴露给 bindings 的自由函数声明（默认参数放在声明处）。
|
||
#pragma once
|
||
#include "../common.h"
|
||
|
||
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);
|
||
|
||
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);
|
||
|
||
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);
|