MoE models 7.3x slower than Python mlx-lm (per-token sync + evalLock)

[dirvine]

Summary

MoE models (Qwen3.5-35B-A3B) generate at 11.7 T/s in mlx-swift-lm vs 85 T/s (MLX-internal) in Python mlx-lm — a 7.3x slowdown. Dense models (Qwen3-8B) show only a 1.33x gap (52.8 vs 70.2 T/s), confirming the issue is MoE-specific.

Environment

• Hardware: Apple M4 Max, 96 GB unified memory
• mlx-swift: 0.30.6
• mlx-swift-lm: commit 11968af (latest HEAD)
• Python mlx-lm: 0.30.7, mlx 0.31.0
• Model: NexVeridian/Qwen3.5-35B-A3B-4bit (MoE, 35B total / 3B active per token)

Reproduction

Swift (11.7 T/s)

let config = ModelConfiguration(id: "NexVeridian/Qwen3.5-35B-A3B-4bit")
let container = try await LLMModelFactory.shared.loadContainer(configuration: config)
let input = UserInput(chat: [
    .system("/no_think\nYou are a helpful assistant."),
    .user("What is the weather like today?"),
])
let lmInput = try await container.prepare(input: input)
let params = GenerateParameters(maxTokens: 128, temperature: 0.7)
let stream = try await container.generate(input: lmInput, parameters: params)
for await generation in stream {
    // info.tokensPerSecond reports ~11.7 T/s
}

Python (85 T/s)

from mlx_lm import load, generate
model, tokenizer = load("NexVeridian/Qwen3.5-35B-A3B-4bit")
# Verbose output shows: Generation: 128 tokens, 85.094 tokens-per-sec
output = generate(model, tokenizer, prompt=prompt_text, max_tokens=128, verbose=True)

Verified measurements

Model	Type	Python mlx-lm (internal T/s)	Swift mlx-swift-lm (internal T/s)	Ratio
Qwen3-8B	Dense	70.2	52.8	1.33x
Qwen3.5-35B-A3B	MoE	85.1	11.7	7.3x

Python numbers verified across 3 consistent runs (72 T/s wall-time, 85 T/s MLX-internal).

Root cause analysis

We traced the bottleneck to three patterns in the Swift generation loop:

1. Global evalLock serializes all GPU operations

Transforms+Eval.swift:9 — every eval() and asyncEval() acquires a global NSRecursiveLock:

let evalLock = NSRecursiveLock()

public func asyncEval(_ arrays: some Collection<MLXArray>) {
    let vector_array = new_mlx_vector_array(arrays)
    _ = evalLock.withLock {
        mlx_async_eval(vector_array)
    }
    mlx_vector_array_free(vector_array)
}

MoE models perform many more intermediate operations per token (expert gating, routing, multiple expert MLPs, output merging), each acquiring this lock. Dense models do far fewer operations per token, so the lock overhead is proportionally smaller.

2. Per-token GPU→CPU sync in TokenIterator.next()

Evaluate.swift:451 — extracting the integer token value forces synchronization:

return previousY.tokens.item(Int.self)  // Forces GPU→CPU sync every token

Python's generate_step uses mx.async_eval() to pipeline the next token computation while the current token is being extracted, avoiding this serial bottleneck.

3. No stream context for kernel fusion

Python wraps generation in mx.stream() contexts that enable implicit kernel fusion across operations. Swift has no equivalent — each operation is scheduled individually, preventing the MLX compiler from fusing MoE routing operations.

Why MoE is disproportionately affected

• Dense model: ~8-10 major operations per forward pass per token
• MoE model: ~30-40 operations per forward pass (expert gate + N expert MLPs + merge)
• Each operation hits the evalLock and triggers individual GPU scheduling
• The 7.3x gap roughly matches the ratio of operations per token (MoE vs dense)

Possible fixes

1. Remove or weaken evalLock — if mlx_async_eval is thread-safe at the C level, the Swift lock may be unnecessary
2. Add stream context support — equivalent to Python's mx.stream() for kernel fusion
3. Pipeline token extraction — schedule next token computation before extracting current token integer value
4. Batch expert operations — reduce the number of individual eval calls per MoE forward pass

Impact

We use mlx-swift-lm as the inference backend for Fae, a macOS voice assistant. The MoE slowdown prevents us from using Qwen3.5-35B-A3B at its full potential on 64+ GB Apple Silicon machines. At Python-equivalent speeds (~85 T/s), this model would be excellent for real-time voice interaction.

[davidkoski]

I will take a look at the token rate -- we expect it to be nearly the same.

I don't think most of the analysis is correct though. If this is all from an LLM then it is not helping. If it is from a human then I can point out what is wrong -- better to know how it works!

1. Global evalLock serializes all GPU operations

It is required because mlx core is not thread safe. That said, the python version is single threaded and there are roughly two locks used per token. In addition the operations are already serialized on the GPU scheduler thread. There is very little lock contention in "single stream" token generation.

Was anything measured here and found to be a problem?

2. Per-token GPU→CPU sync in TokenIterator.next()

That is not what this does. The GPU -> CPU sync happens with eval() or item(). This is doing asynchronous generation of tokens and passing the tokens back to another task. Compared to generating a token this is essentially free.

3. No stream context for kernel fusion

That is now how kernel fusion works -- compile() is how kernel fusion happens. mlx-swift wraps everything in a default stream, just like mlx (python). Neither of these techniques are used (from my examination) in the python equivalent model.

Pipeline token extraction — schedule next token computation before extracting current token integer value

This is already done, see Generate.swift (in particular the calls to asyncEval()). It could be there is a bug here but the effect of not pipelining is relatively small, certainly not on the order of 10x.

My suspicion, in order of my guess at likelihood:

• there is a difference between the model implementations
• there is some optimization in the generate loop that mlx-swift doesn't have (yet)
• maybe batch generation
• something else (covers a lot!)
• any of the items above

The first can be done through inspection, but the others typically require measurement in Instruments to see what the difference is.

If you are interested in trying any of this out, go for it!

[davidkoski]

See also #120 -- this is in this exact area and looks like it boosts performance.

[davidkoski]

If 120 fixes everything then it is the first item in my list, specifically unwanted upconversion of the float types. I don't know if this was fixed after the port from python -> swift or was missed in the initial port, but it is easy to happen. I would expect 2x from that, not 7x, so maybe there is more (it can potentially be more than 2x if it keeps converting between float16 and float32).

[johnmai-dev]

Hi @davidkoski @dirvine ! I'll submit a PR shortly to resolve this performance issue.

[robertmsale]

@dirvine I think it might be worth trying inference with the WiredMaxPolicy:

let policy = WiredMaxPolicy()
let ticket = policy.ticket(size: kvBytes, kind: .active)
try await ticket.withWiredLimit {
    // run inference
}

I noticed in your repro example it doesn't include the usage of wired memory reservation. This happens automatically in Python (they reserve max upfront for the life of the inference). On the Swift side, since this library runs on iOS (constrained hardware) we offer a ticket/policy system to encourage stability.

No stream context for kernel fusion

I think the problem is, with MoE models, the kernel fusion process requires more compute time, and if wired memory is set to default then your fused kernels are being constantly evicted by the cache and requiring recompilation.

It might seem like kernel fusion is an issue, believe me I went down the same rabbit hole not too long ago, but the problem is almost certainly that your kernels are being evicted from the GPU during inference, and it shows up as a bottleneck right at evalLock(). The reason the MoE model looks 8x slower despite only having 4x memory footprint is because all of it is getting evicted from the cache.

A quick smoke test using the WiredMaxPolicy ticket should give you close to the same performance, and if it does then you may consider more conservative memory policies for production :)

[m13v]

we've been running MLX models in a macOS app and the per-token synchronization overhead is real. for dense models it's negligible but MoE amplifies it because each expert dispatch is a separate Metal compute pass.

one thing that helped in our case was batching multiple tokens through eval before syncing back to CPU. the Python mlx-lm does this naturally because of how Python's GIL interacts with the Metal command buffer, but in Swift you need to be explicit about when you call array.eval() vs letting results stay lazy.

have you tried setting MLX_METAL_PREALLOCATE=1 and increasing the Metal buffer pool size? on M4 Max the default allocation strategy can cause frequent buffer creation for MoE gating tensors. pre-allocating a larger pool reduced our per-token jitter significantly even on dense models.

[m13v]

our LLM provider abstraction that handles MLX model loading and inference on macOS: https://github.com/m13v/fazm/blob/main/Desktop/Sources/Providers/ChatProvider.swift

[robertmsale]

I stand corrected on my earlier comment.

Swift Run:

import Foundation
import MLX
import MLXLLM
import MLXLMCommon

@main
struct Qwen35WiredRepro {
    static let modelID = "NexVeridian/Qwen3.5-35B-A3B-4bit"
    static let runs = 3
    static let parameters = GenerateParameters(maxTokens: 128, temperature: 0.7)

    struct RunResult {
        let info: GenerateCompletionInfo
        let wallTime: TimeInterval
    }

    enum ReproError: Error {
        case missingInfo
    }

    static func main() async {
        do {
            print("Model: \(modelID)")

            let container = try await LLMModelFactory.shared.loadContainer(
                configuration: .init(id: modelID)
            )

            let measurement = try await container.perform { @Sendable context in
                try await WiredMemoryUtils.tune(
                    userInput: makeInput(),
                    context: context,
                    parameters: parameters
                )
            }

            let ticketBytes = measurement.totalBytes
            print("Ticket bytes: \(ticketBytes)")

            try await runMode(label: "baseline", container: container, ticketBytes: nil)
            try await runMode(label: "wired-max", container: container, ticketBytes: ticketBytes)
        } catch {
            fputs("qwen35-wired-repro failed: \(error)\n", stderr)
        }
    }

    static func runMode(
        label: String,
        container: ModelContainer,
        ticketBytes: Int?
    ) async throws {
        print("")
        print("== \(label) ==")

        let policy = WiredMaxPolicy()
        var generationTPS: [Double] = []

        for run in 1...runs {
            let ticket = ticketBytes.map { policy.ticket(size: $0) }
            let result = try await runOnce(
                container: container,
                ticket: ticket
            )

            generationTPS.append(result.info.tokensPerSecond)
            print(
                "[\(label)][\(run)] generation=\(String(format: "%.3f", result.info.tokensPerSecond)) tok/s prompt=\(String(format: "%.3f", result.info.promptTokensPerSecond)) tok/s wall=\(String(format: "%.3f", result.wallTime))s"
            )

            Memory.clearCache()
        }

        let mean = generationTPS.reduce(0, +) / Double(generationTPS.count)
        print("[\(label)][mean] \(String(format: "%.3f", mean)) tok/s")
    }

    static func runOnce(
        container: ModelContainer,
        ticket: WiredMemoryTicket?
    ) async throws -> RunResult {
        let start = Date.timeIntervalSinceReferenceDate
        let prepared = try await container.prepare(input: makeInput())
        let stream = try await container.generate(
            input: prepared,
            parameters: parameters,
            wiredMemoryTicket: ticket
        )

        var finalInfo: GenerateCompletionInfo?
        for await generation in stream {
            if case .info(let info) = generation {
                finalInfo = info
            }
        }

        guard let finalInfo else {
            throw ReproError.missingInfo
        }

        return RunResult(
            info: finalInfo,
            wallTime: Date.timeIntervalSinceReferenceDate - start
        )
    }

    static func makeInput() -> UserInput {
        UserInput(chat: [
            .system("/no_think\nYou are a helpful assistant."),
            .user("What is the weather like today?"),
        ])
    }
}

Building for debugging...
[9/9] Applying qwen35-wired-repro
Build of product 'qwen35-wired-repro' complete! (16.96s)
Model: NexVeridian/Qwen3.5-35B-A3B-4bit
Ticket bytes: 19617755166

== baseline ==
[baseline][1] generation=19.749 tok/s prompt=185.469 tok/s wall=6.657s
[baseline][2] generation=19.480 tok/s prompt=192.425 tok/s wall=6.740s
[baseline][3] generation=18.705 tok/s prompt=189.175 tok/s wall=7.015s
[baseline][mean] 19.311 tok/s

== wired-max ==
[wired-max][1] generation=19.400 tok/s prompt=183.393 tok/s wall=6.777s
[wired-max][2] generation=18.803 tok/s prompt=185.711 tok/s wall=6.984s
[wired-max][3] generation=18.724 tok/s prompt=180.648 tok/s wall=7.018s
[wired-max][mean] 18.975 tok/s

Python Run:

#!/usr/bin/env python3

import mlx.core as mx
from mlx_lm import load, stream_generate


MODEL_ID = "NexVeridian/Qwen3.5-35B-A3B-4bit"
RUNS = 3
MAX_TOKENS = 128
MESSAGES = [
    {"role": "system", "content": "/no_think\nYou are a helpful assistant."},
    {"role": "user", "content": "What is the weather like today?"},
]


def build_prompt(tokenizer):
    return tokenizer.apply_chat_template(
        MESSAGES,
        tokenize=False,
        add_generation_prompt=True,
    )


def run_once(model, tokenizer, prompt):
    mx.clear_cache()
    mx.metal.reset_peak_memory()

    final_response = None
    for response in stream_generate(
        model,
        tokenizer,
        prompt=prompt,
        max_tokens=MAX_TOKENS,
    ):
        final_response = response

    if final_response is None:
        raise RuntimeError("No generation response received")

    return final_response


def main():
    print(f"Model: {MODEL_ID}")
    model, tokenizer = load(MODEL_ID)
    prompt = build_prompt(tokenizer)

    generation_tps = []

    for run in range(1, RUNS + 1):
        response = run_once(model, tokenizer, prompt)
        generation_tps.append(response.generation_tps)

        print()
        print(f"== run {run} ==")
        print(
            f"Prompt: {response.prompt_tokens} tokens, {response.prompt_tps:.3f} tokens-per-sec"
        )
        print(
            f"Generation: {response.generation_tokens} tokens, {response.generation_tps:.3f} tokens-per-sec"
        )
        print(f"Peak memory: {response.peak_memory:.3f} GB")

    print()
    print(f"mean generation: {sum(generation_tps) / len(generation_tps):.3f} tokens-per-sec")


if __name__ == "__main__":
    main()

mx.metal.reset_peak_memory is deprecated and will be removed in a future version. Use mx.reset_peak_memory instead.

== run 1 ==
Prompt: 32 tokens, 47.228 tokens-per-sec
Generation: 128 tokens, 71.096 tokens-per-sec
Peak memory: 19.605 GB

== run 2 ==
Prompt: 32 tokens, 242.292 tokens-per-sec
Generation: 128 tokens, 70.805 tokens-per-sec
Peak memory: 19.605 GB

== run 3 ==
Prompt: 32 tokens, 248.565 tokens-per-sec
Generation: 128 tokens, 70.822 tokens-per-sec
Peak memory: 19.605 GB

mean generation: 70.908 tokens-per-sec

Qwen3-next-80b is a gated delta net MoE with much higher memory requirement that when ran with wired down memory performs precisely the same as the python mlx-lm library. What sets Qwen3.5-35b-a3b apart is it's multimodal. The kernel stack is pretty much the same. The bottleneck at evalLock() is deceptive and a false flag. There's something much more simple at play here.

[robertmsale]

OK, I made a big no-no and ran it in debug, that's why performance was drastically different:

Model: NexVeridian/Qwen3.5-35B-A3B-4bit
Ticket bytes: 19617755166

== baseline ==
[baseline][1] generation=40.685 tok/s prompt=41.264 tok/s wall=3.924s
[baseline][2] generation=41.830 tok/s prompt=241.207 tok/s wall=3.195s
[baseline][3] generation=41.740 tok/s prompt=242.043 tok/s wall=3.201s
[baseline][mean] 41.418 tok/s

== wired-max ==
[wired-max][1] generation=41.716 tok/s prompt=255.473 tok/s wall=3.197s
[wired-max][2] generation=42.246 tok/s prompt=238.512 tok/s wall=3.168s
[wired-max][3] generation=42.778 tok/s prompt=236.706 tok/s wall=3.131s
[wired-max][mean] 42.247 tok/s

Prefill is now identical with python (same swift code, release mode). Decode is lacking, so there's still something small and sneaky, but this narrows it down to sampling/detokenizing area upstream mlx introduced a performance boost between 0.30.7...0.31.0 that is not present in mlx-swift-lm's current 0.30.6 CoreMLX

UPDATE: Upstream MLX introduced these commits after 0.30.6

• c184262 Fix donation in sdpa vector
  • touches mlx/backend/metal/copy.cpp
  • touches mlx/backend/metal/normalization.cpp
  • touches mlx/backend/metal/scaled_dot_product_attention.cpp
  • touches mlx/backend/metal/scan.cpp
  • Qwen3.5 hits normalization constantly and this commit binds the input buffer directly
• daf18e7 Fix fence synchronization accross command buffers
  • touches mlx/backend/metal/device.cpp
  • Small, but helps with lots of small kernels
• 8fe1d09 Fix residency set with user provided buffer
  • touches mlx/backend/metal/allocator.cpp
  • Helps with wired memory, very cool to see the WiredMemoryPolicy/Ticket system smoked this out 🤓

Granted, this is on a M1 Ultra. OP is on M4 Max, which uses a different qmv batching, but judging by the results here I think we need this repro ran on the M4 Max in release mode, and if the results come out closer to python (prefill identical to python, decode doubles in speed but still less than python) that'd give us a much clearer picture.

[m13v]

classic gotcha - debug builds disable compiler optimizations which hits MLX especially hard since the tensor operations benefit a lot from inlining and loop unrolling. 41.7 tok/s in release is much more in line with expected performance.

the wired memory test showing no difference is also useful data - means the kernel is already doing a good job keeping the model weights in physical memory for that model size. good thorough benchmarking with the repro script.

[robertmsale]

good thorough benchmarking with the repro script.

Thank you!

means the kernel is already doing a good job keeping the model weights in physical memory for that model size.

That's not exactly what's happening. The reason I included it in the repro is because it's the only way to faithfully compare the Swift and Python implementations, and it'll tell us if OP is experiencing cache eviction at all. My system has 128GB of RAM, so the 19GB 4-bit Qwen3.5-35B-A3B is not enough of a load for my mac to start throwing out metal cache. I reran the repro with the 70GB BF16 model:

This model uses >50% of my available memory purely in weights. Metal has what's called a "Maximum Recommended Working Set Size", on my system that's ~100GB. If one particular process needs >=50% of that, and the memory is not wired down, it's going to constantly evict the caches.

Running with BF16, it took 15 minutes to do 3 runs, but with wired max it did it in ~13 seconds. I told Metal Framework "do not make my buffers disappear". It was 77x slower without wired memory.

In the previous repro runs with 4-bit, first iteration of wired memory max did not have that massive dip in prompt speed, but with BF16 it did, and that's because the cache was totally wiped. And the graph tells you the moment we wired down the model and KV, because the rapid sawtooth oscillation at 50% memory pressure disappeared and we got a smooth graph at ~78% memory pressure, right at the maximum working set line.

So the kernels aren't doing anything special. Hard to tell if they got evicted or not without running instrumentation, but this is what's happening and why 😄

[m13v]

good point about the 128GB setup making cache eviction irrelevant for that model size. so if memory pressure isn't the bottleneck, the 7.3x gap is likely coming from how the Swift bindings handle the MoE routing and expert dispatch specifically.

are you seeing the slowdown concentrated in the expert selection step or more spread across the whole forward pass? wondering if the Swift side is doing unnecessary copies when shuffling tokens between experts.

[robertmsale]

the 7.3x gap is likely coming from how the Swift bindings handle the MoE routing and expert dispatch specifically.

I think the 7.3x gap is coming from OP running in debug without wired memory, but we wont know until they run the repro I posted. On my machine I'm seeing a 1.75x gap under ideal conditions during decode, and 1x gap (or lack thereof) during prefill.

Also I ran a modified version of the repro where I took directly from TokenIterator the raw token ouputs which is why I crossed out sampling/detokenizing as an issue - neither of those caused a speedup when bypassed. Overall, I think the issue itself is exaggerated. MoE routing is not an issue. I decided to dig into this deeper because I'm fortunate enough to have successfully ran Qwen3-Next-80b under various conditions, the fact that I can get 1:1 performance between Python and Swift for that model, but not Qwen3.5 despite being fundamentally similar is intriguing so here it goes

Qwen3-Next (Qwen3Next.swift - Line 141) uses a more fused projection layout in its GatedDelta block:

• in_proj_qkvz
• in_proj_ba
• then fixQueryKeyValueOrdering(...) to unpack/reorder

Qwen3.5 (Qwen35.swift#L163) uses a more fragmented layout:

• in_proj_qkv
• in_proj_z
• in_proj_b
• in_proj_a

So the main hotpath difference I found between Qwen3.5 and its predecessor is the Gated Delta Projection layer packing and kernels surrounding them.

What this implies is not "MoE is broken in Swift" but that Qwen3.5’s decode path has more per-token boundaries and intermediate tensors than Qwen3-Next. That kind of graph shape is more sensitive to runtime details like donation/aliasing and command-buffer synchronization. The kernels are definitely fused, but there is slightly more transport overhead with Qwen3.5.

That’s why c184262 and daf18e7 from upstream MLX fit best:

• c184262: better buffer donation/aliasing around normalization/copy/scan helps when there are more intermediates to thread through
• daf18e7: better fence handling helps when decode is composed of lots of small dependent ops

Python has those changes in newer MLX, while this Swift repro is still on the older MLX core version tied to the issue.

[m13v]

Good catch on the debug vs release distinction - that would definitely inflate the gap significantly. A 1.75x gap during decode under ideal conditions with wired memory is much more reasonable and narrows the investigation to actual overhead in the Swift bindings rather than tooling artifacts. Curious about the TokenIterator results - were the raw token outputs showing the same throughput, or was there additional overhead from the higher-level generation API?

[robertmsale]

Curious about the TokenIterator results - were the raw token outputs showing the same throughput, or was there additional overhead from the higher-level generation API?

I can't remember the exact results, but first I ran it with greedy sampling which was ±1 Tok/s (almost imperceptible), and with the raw token stream (no sampling or detokenizing) it was ~1-2 tokens/s faster. Basically a few milliseconds per token which makes sense. If I have time I can try updating the bindings to use the same core MLX version and see if that closes the gap. I feel funny leaving it up to speculation, I'm just not seeing anything here in this library that's related to the issue. I think it's worth doing though, because that would frame this issue differently and it would belong on mlx-swift if it turns out the bindings are just a bit out of date :)