Skip to content

fix: sync bug fixes from the WasmEdge MLX plugin #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
grorge123 merged 3 commits into from
Nov 9, 2026
Merged

fix: sync bug fixes from the WasmEdge MLX plugin #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8a6fc19
fix: sync bug fixes from the WasmEdge MLX plugin
grorge123 Nov 9, 2026
c329760
ci: run macOS build on macos-latest
grorge123 Nov 9, 2026
fe49765
build: drop unused fmt dependency
grorge123 Nov 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/build.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ jobs:
cxx: clang++
mlx_args: -DMLX_BUILD_CPU=ON -DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=OFF
- name: macOS
os: macos-14
os: macos-latest
cc: clang
cxx: clang++
mlx_args: -DMLX_BUILD_CPU=ON -DMLX_BUILD_METAL=ON -DMLX_BUILD_CUDA=OFF
Expand Down
5 changes: 0 additions & 5 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,6 @@ endif()

include(FetchContent)

FetchContent_Declare(
fmt
GIT_REPOSITORY https://github.com/fmtlib/fmt
GIT_TAG e69e5f977d458f2650bb346dadf2ad30c5320281)
FetchContent_MakeAvailable(fmt)
find_package(MLX CONFIG REQUIRED)
find_package(ZLIB REQUIRED)

Expand Down
20 changes: 11 additions & 9 deletions mlx/base.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,24 +9,25 @@

namespace mlx::core::nn {

mx::array &Module::registerParameter(const std::string &Name, mx::array &&W) {
Parameters.insert({Name, W});
return Parameters.at(Name);
mx::array &Module::registerParameter(const std::string &ParamName,
mx::array &&W) {
Parameters.insert({ParamName, W});
return Parameters.at(ParamName);
}
void Module::update(std::unordered_map<std::string, mx::array> Parameters) {
for (auto &[K, V] : Parameters) {
void Module::update(std::unordered_map<std::string, mx::array> NewParameters) {
for (auto &[K, V] : NewParameters) {
apply(K, V);
}
}
std::shared_ptr<nn::Module> Module::toQuantized(
int GroupSize, int Bits, const std::string &Prefix,
const std::unordered_map<std::string, mx::array> &Parameters) {
const std::unordered_map<std::string, mx::array> &LoadedWeights) {
auto NewPrefix = Prefix + Name + (Prefix.empty() && Name.empty() ? "" : ".");
for (auto &[K, V] : Submodules) {
if (V->hasQuantize()) {
auto Weights = V->Parameters.find("weight");
if (Weights != V->Parameters.end() && !Parameters.empty()) {
if (Parameters.count(NewPrefix + V->Name + ".scales") == 0) {
if (Weights != V->Parameters.end() && !LoadedWeights.empty()) {
if (LoadedWeights.count(NewPrefix + V->Name + ".scales") == 0) {
continue;
}
}
Expand All @@ -36,7 +37,8 @@ std::shared_ptr<nn::Module> Module::toQuantized(
}
}
V = V->toQuantized(GroupSize, Bits,
Prefix + Name + (Name.empty() ? "" : "."), Parameters);
Prefix + Name + (Name.empty() ? "" : "."),
LoadedWeights);
}
return shared_from_this();
}
Expand Down
8 changes: 4 additions & 4 deletions mlx/base.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,15 @@ class Module : public std::enable_shared_from_this<Module> {
std::string Name;
std::unordered_map<std::string, mx::array> Parameters{};
std::unordered_map<std::string, std::shared_ptr<Module>> Submodules{};
mx::array &registerParameter(const std::string &Name, mx::array &&W);
mx::array &registerParameter(const std::string &ParamName, mx::array &&W);
std::unordered_map<std::string, mx::array>
getWeigts(const std::string &Prefix = "model");
virtual std::shared_ptr<nn::Module> toQuantized(
int GroupSize = 64, int Bits = 4, const std::string &Prefix = "",
const std::unordered_map<std::string, mx::array> &Parameters = {});
const std::unordered_map<std::string, mx::array> &LoadedWeights = {});
virtual bool hasQuantize() { return false; }
void update(std::unordered_map<std::string, mx::array> Parameters);
void apply(std::string Key, const mx::array &Parameters);
void update(std::unordered_map<std::string, mx::array> NewParameters);
void apply(std::string Key, const mx::array &Value);
template <typename T>
void registerModule(const std::string &ModuleName,
const std::shared_ptr<T> &M) {
Expand Down
4 changes: 2 additions & 2 deletions mlx/embedding.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@ class Embedding : public Module {
mx::array asLinear(const mx::array &Input);
std::shared_ptr<nn::Module>
toQuantized(int GroupSize = 64, int Bits = 4, const std::string &Prefix = "",
const std::unordered_map<std::string, mx::array> &Parameters = {})
override;
const std::unordered_map<std::string, mx::array> &LoadedWeights =
{}) override;

virtual bool hasQuantize() override { return true; }
};
Expand Down
4 changes: 2 additions & 2 deletions mlx/linear.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ class Linear : public Module {
virtual mx::array forward(mx::array Input);
std::shared_ptr<nn::Module>
toQuantized(int GroupSize = 64, int Bits = 4, const std::string &Prefix = "",
const std::unordered_map<std::string, mx::array> &Parameters = {})
override;
const std::unordered_map<std::string, mx::array> &LoadedWeights =
{}) override;

virtual bool hasQuantize() override { return true; }
};
Expand Down
40 changes: 20 additions & 20 deletions mlx/pooling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,8 @@ mx::array slidingWindows(const mx::array &X,
if (X.ndim() < 3) {
spdlog::error(
"To extract sliding windows at least 1 spatial dimension (3 total) is "
"needed but the input only has " +
std::to_string(X.ndim()) + " dimensions.");
"needed but the input only has {} dimensions.",
X.ndim());
assumingUnreachable();
}
auto Shape = X.shape();
Expand Down Expand Up @@ -111,13 +111,12 @@ mx::array slidingWindows(const mx::array &X,
}
return mx::as_strided(X, FinalShape, FinalStrides, 0);
}
Pool::Pool(
const std::function<mx::array(const mx::array &, const std::vector<int> &)>
&PoolingFunction,
const std::vector<int> &KernelSize, const std::vector<int> &Stride,
const std::vector<std::pair<int, int>> &Padding, int PaddingValue)
: PoolingFunction(PoolingFunction), KernelSize(KernelSize), Stride(Stride),
Padding(Padding), PaddingValue(PaddingValue) {
Pool::Pool(const std::function<mx::array(const mx::array &,
const std::vector<int> &)> &PoolingFn,
const std::vector<int> &KernelSizes, const std::vector<int> &Strides,
const std::vector<std::pair<int, int>> &Paddings, int PadValue)
: PoolingFunction(PoolingFn), KernelSize(KernelSizes), Stride(Strides),
Padding(Paddings), PaddingValue(PadValue) {
int N = KernelSize.size();
for (int I = -(N + 1); I < -1; I++) {
Axes.push_back(I);
Expand All @@ -144,29 +143,30 @@ mx::array Pool::forward(const mx::array &X) {

Pool2d::Pool2d(
const std::function<mx::array(const mx::array &, const std::vector<int> &)>
&PoolingFunction,
int PaddingValue, const std::vector<int> &KernelSize,
&PoolingFn,
int PadValue, const std::vector<int> &KernelSizes,
const std::optional<std::vector<int>> &StrideOpt,
const std::optional<std::vector<int>> &PaddingOpt)
: Pool(PoolingFunction,
KernelSize.size() == 1 ? valueOrList(KernelSize[0], 2) : KernelSize,
: Pool(PoolingFn,
KernelSizes.size() == 1 ? valueOrList(KernelSizes[0], 2)
: KernelSizes,
(StrideOpt.has_value()
? (StrideOpt.value().size() == 1
? valueOrList(StrideOpt.value()[0], 2)
: StrideOpt.value())
: (KernelSize.size() == 1 ? valueOrList(KernelSize[0], 2)
: KernelSize)),
: (KernelSizes.size() == 1 ? valueOrList(KernelSizes[0], 2)
: KernelSizes)),
makePaddingPairs(PaddingOpt.has_value() ? PaddingOpt.value()
: valueOrList(0, 2)),
PaddingValue) {}
PadValue) {}

AvgPool2d::AvgPool2d(const std::vector<int> &KernelSize,
const std::optional<std::vector<int>> &Stride,
const std::optional<std::vector<int>> &Padding)
AvgPool2d::AvgPool2d(const std::vector<int> &KernelSizes,
const std::optional<std::vector<int>> &StrideOpt,
const std::optional<std::vector<int>> &PaddingOpt)
: Pool2d(
[](const mx::array &A, const std::vector<int> &Axis) -> mx::array {
return mx::mean(A, Axis, false);
},
0, KernelSize, Stride, Padding) {}
0, KernelSizes, StrideOpt, PaddingOpt) {}

} // namespace mlx::core::nn
24 changes: 12 additions & 12 deletions mlx/pooling.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@
namespace mlx::core::nn {
class Pool : public nn::Module {
public:
Pool(const std::function<mx::array(
const mx::array &, const std::vector<int> &)> &PoolingFunction,
const std::vector<int> &KernelSize, const std::vector<int> &Stride,
const std::vector<std::pair<int, int>> &Padding, int PaddingValue);
Pool(const std::function<mx::array(const mx::array &,
const std::vector<int> &)> &PoolingFn,
const std::vector<int> &KernelSizes, const std::vector<int> &Strides,
const std::vector<std::pair<int, int>> &Paddings, int PadValue);
mx::array forward(const mx::array &X);

protected:
Expand All @@ -23,18 +23,18 @@ class Pool : public nn::Module {

class Pool2d : public Pool {
public:
Pool2d(const std::function<mx::array(
const mx::array &, const std::vector<int> &)> &PoolingFunction,
int PaddingValue, const std::vector<int> &KernelSize,
const std::optional<std::vector<int>> &Stride,
const std::optional<std::vector<int>> &Padding);
Pool2d(const std::function<mx::array(const mx::array &,
const std::vector<int> &)> &PoolingFn,
int PadValue, const std::vector<int> &KernelSizes,
const std::optional<std::vector<int>> &StrideOpt,
const std::optional<std::vector<int>> &PaddingOpt);
};

class AvgPool2d : public Pool2d {
public:
AvgPool2d(const std::vector<int> &KernelSize,
const std::optional<std::vector<int>> &Stride = std::nullopt,
const std::optional<std::vector<int>> &Padding = std::nullopt);
AvgPool2d(const std::vector<int> &KernelSizes,
const std::optional<std::vector<int>> &StrideOpt = std::nullopt,
const std::optional<std::vector<int>> &PaddingOpt = std::nullopt);
};

} // namespace mlx::core::nn
2 changes: 2 additions & 0 deletions model/gemma3/gemma3.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ class Model : public vlm::Module {
const mx::array &Mask,
const std::optional<std::vector<std::shared_ptr<vlm::BaseCache>>> &Cache =
std::nullopt) override;
// <eos> and <end_of_turn>.
std::vector<int> eosTokenIds() const override { return {1, 106}; }
static std::shared_ptr<Model> fromPretrained(const std::string &ModelPath);
std::unordered_map<std::string, mx::array>
sanitize(const std::unordered_map<std::string, mx::array> &Weights);
Expand Down
4 changes: 2 additions & 2 deletions model/utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,8 @@ void saveWeights(const mx::array &Weights, const std::string &Path) {
std::string loadBytesFromFile(const std::string &Path) {
std::ifstream Fs(Path, std::ios::in | std::ios::binary);
if (Fs.fail()) {
std::cerr << "Cannot open " << Path << std::endl;
exit(1);
spdlog::error("Cannot open {}.", Path);
return "";
}
std::string Data;
Fs.seekg(0, std::ios::end);
Expand Down
11 changes: 8 additions & 3 deletions model/vlm_base.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#include "base.h"
#include "spdlog/spdlog.h"
#include "vlm_sampling.h"
#include <algorithm>
#include <ctime>
#include <memory>
#include <mlx/array.h>
Expand Down Expand Up @@ -126,8 +127,8 @@ int KVCache::trim(int N) {
}

// RotatingKVCache implementation
RotatingKVCache::RotatingKVCache(int MaxSize, int Keep, int Step)
: KVCache(0, 0, Step), Keep(Keep), MaxSize(MaxSize), Idx(0) {}
RotatingKVCache::RotatingKVCache(int MaxSize, int Keep, int StepSize)
: KVCache(0, 0, StepSize), Keep(Keep), MaxSize(MaxSize), Idx(0) {}

mx::array RotatingKVCache::trim(int TrimSize, const mx::array &V,
std::optional<mx::array> Append) {
Expand Down Expand Up @@ -480,14 +481,14 @@ std::vector<int> Module::generate(
GenerationResult LastResponse;
auto Tic = std::chrono::system_clock::now().time_since_epoch();
std::vector<int32_t> TokenList;
const std::vector<int> EosTokenIds = eosTokenIds();
int N = 0;
while (true) {
int PromptTPS;
if (N == 0) {
PromptTPS = std::chrono::duration_cast<std::chrono::seconds>(Tic).count();
Tic = std::chrono::system_clock::now().time_since_epoch();
}
// TODO: if token = eos_token_id break
auto Response = GenerationResult();
N++;
if (N >= MaxTokens) {
Expand All @@ -497,6 +498,10 @@ std::vector<int> Module::generate(
mx::async_eval(NextY);
// TODO: handle decoder_input_ids
auto Token = Y.item<int>();
if (std::find(EosTokenIds.begin(), EosTokenIds.end(), Token) !=
EosTokenIds.end()) {
break;
}
TokenList.emplace_back(Token);
Y = NextY;
LogProbs = NextLogProbs;
Expand Down
5 changes: 4 additions & 1 deletion model/vlm_base.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ class RotatingKVCache : public KVCache {
int MaxSize;
int Idx;

RotatingKVCache(int MaxSize = -1, int Keep = 0, int Step = 256);
RotatingKVCache(int MaxSize = -1, int Keep = 0, int StepSize = 256);

std::tuple<mx::array, mx::array>
updateAndFetch(const mx::array &NewKeys, const mx::array &NewValues) override;
Expand Down Expand Up @@ -93,6 +93,9 @@ class Module : public mlx::core::nn::Module {
const mx::array &Mask,
const std::optional<std::vector<std::shared_ptr<vlm::BaseCache>>> &Cache =
std::nullopt) = 0;
// Token ids that end generation. Models override this with their EOS /
// end-of-turn ids; the default never stops early.
virtual std::vector<int> eosTokenIds() const { return {}; }
struct GenerationResult {
std::string Text;
int Token;
Expand Down
7 changes: 4 additions & 3 deletions model/whisper/decoding.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,8 +466,9 @@ DecodingTask::DecodingTask(const std::shared_ptr<Whisper> &Model,
}
}

DecodingOptions DecodingTask::verifyOptions(const DecodingOptions &Options) {
DecodingOptions Result = Options;
DecodingOptions
DecodingTask::verifyOptions(const DecodingOptions &InputOptions) {
DecodingOptions Result = InputOptions;

// Check beam_size and best_of conflicts
if (Result.BeamSize && Result.BestOf) {
Expand Down Expand Up @@ -740,7 +741,7 @@ std::vector<DecodingResult> DecodingTask::run(const mx::array &Mel) {
// tokens = mx.broadcast_to(tokens, [n_audio, self.n_group,
// len(self.initial_tokens)])
MlxShape NewShape = {NAudio, NGroup,
static_cast<int>(InitialTokens.size())};
static_cast<int>(InitialTokens.size())};
Tokens = mx::broadcast_to(Tokens, NewShape);

// tokens = tokens.reshape(n_audio * self.n_group, len(self.initial_tokens))
Expand Down
2 changes: 1 addition & 1 deletion model/whisper/decoding.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ class DecodingTask {
std::vector<DecodingResult> run(const mx::array &Mel);

private:
DecodingOptions verifyOptions(const DecodingOptions &Options);
DecodingOptions verifyOptions(const DecodingOptions &InputOptions);
std::vector<int> getInitialTokens();
std::vector<int> getSuppressTokens();
mx::array getAudioFeatures(const mx::array &Mel);
Expand Down
13 changes: 7 additions & 6 deletions model/whisper/tokenizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,8 +190,8 @@ Encoding::Encoding(const std::string &Name, int ExplicitNVocab,
const std::string &PatStr,
const std::unordered_map<std::string, int> &MergeableRanks,
const std::unordered_map<std::string, int> &SpecialTokens)
: Name(Name), PatStr(PatStr),
MergeableRanks(MergeableRanks), SpecialTokens(SpecialTokens) {
: Name(Name), PatStr(PatStr), MergeableRanks(MergeableRanks),
SpecialTokens(SpecialTokens) {

EotToken = 50257; // Default EOT token

Expand Down Expand Up @@ -447,13 +447,14 @@ int Tokenizer::languageToken() const {
return toLanguageToken(Language.value());
}

int Tokenizer::toLanguageToken(const std::string &Language) const {
std::string TokenName = "<|" + Language + "|>";
int Tokenizer::toLanguageToken(const std::string &LanguageCode) const {
std::string TokenName = "<|" + LanguageCode + "|>";
auto It = SpecialTokens.find(TokenName);
if (It != SpecialTokens.end()) {
return It->second;
}
throw std::runtime_error("Language " + Language + " not found in tokenizer.");
throw std::runtime_error("Language " + LanguageCode +
" not found in tokenizer.");
}

std::vector<int> Tokenizer::getAllLanguageTokens() const {
Expand Down Expand Up @@ -699,7 +700,7 @@ std::unique_ptr<Encoding> getEncoding(const std::string &Name,
std::vector<std::string> Specials = {"<|endoftext|>",
"<|startoftranscript|>"};

for (const auto &[Code, Name] : LANGUAGES) {
for (const auto &[Code, LanguageName] : LANGUAGES) {
Specials.push_back("<|" + Code + "|>");
if (static_cast<int>(Specials.size()) >= NumLanguages + 2)
break;
Expand Down
2 changes: 1 addition & 1 deletion model/whisper/tokenizer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ class Tokenizer {
int getNoTimestamps() const;
int getTimestampBegin() const;
int languageToken() const;
int toLanguageToken(const std::string &Language) const;
int toLanguageToken(const std::string &LanguageCode) const;
std::vector<int> getAllLanguageTokens() const;
std::vector<std::string> getAllLanguageCodes() const;
std::vector<int> getSotSequenceIncludingNotimestamps() const;
Expand Down
Loading
Loading