zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

MoE backward improvements (#2335)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-07-07 17:59:53 -07:00
committed by GitHub
parent a4fcc893cd
commit 4a9b29a875
22 changed files with 1130 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
mlx/primitives.h
View File

@@ -526,6 +526,16 @@ class GatherMM : public UnaryPrimitive {
bool right_sorted_;
};
class SegmentedMM : public UnaryPrimitive {
public:
explicit SegmentedMM(Stream stream) : UnaryPrimitive(stream) {}
void eval_cpu(const std::vector<array>& inputs, array& out) override;
void eval_gpu(const std::vector<array>& inputs, array& out) override;
DEFINE_PRINT(SegmentedMM)
};
class BroadcastAxes : public UnaryPrimitive {
public:
explicit BroadcastAxes(Stream stream, std::vector<int> ignore_axes = {})