From c5de3244cd85b1916451c0a70fc4b5dace9ab653 Mon Sep 17 00:00:00 2001 From: simba Date: Fri, 25 Sep 2026 13:00:11 -0700 Subject: [PATCH 1/3] fix(qwen35): SSD expert streaming crashes on prefill and decode Two regressions from the upstream sync (#62) broke --stream-experts on Qwen3.5/3.6 MoE: - The fork-only SSD block in SwitchGLU.projectExperts unsorted its output and still returned inverseOrder, so callers unsorted a second time. Prefill crashed with broadcast_shapes (N,8,8,D) vs (N,8,1). - Decode now runs through compiled traces, but SSD streaming evals mid-layer to load experts, which a trace cannot contain. Skip the compiled MoE block, decoder layer and decode segments while streaming. Co-Authored-By: Claude Opus 5.5 --- Libraries/MLXLLM/Models/Qwen35.swift | 9 ++++++--- Libraries/MLXLMCommon/SwitchLayers.swift | 4 +--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/Libraries/MLXLLM/Models/Qwen35.swift b/Libraries/MLXLLM/Models/Qwen35.swift index 0fa519324b..51a485f707 100644 --- a/Libraries/MLXLLM/Models/Qwen35.swift +++ b/Libraries/MLXLLM/Models/Qwen35.swift @@ -675,7 +675,8 @@ final class Qwen35SparseMoeBlock: Module, UnaryLayer { // Decode (S == 1) runs through a compiled trace: fusion merges the // elementwise chains into fewer kernels, bit-identically. Prefill // stays unfused — it is GEMM-bound and would pay a trace per shape. - if x.dim(1) != 1 { + // SSD expert streaming evals mid-layer, which a trace cannot contain. + if x.dim(1) != 1 || ExpertStreamingConfig.shared.isEnabled { return forward(x) } return compiledForward(self, x) @@ -788,7 +789,7 @@ final class Qwen35DecoderLayer: Module { // Single-token unmasked decode runs the layer as one traced function // (two for full attention, split at the KV write). Everything else // takes the general body below. - if x.dim(1) == 1, ssmMask == nil { + if x.dim(1) == 1, ssmMask == nil, !ExpertStreamingConfig.shared.isEnabled { if isLinear, let mambaCache = cache as? MambaCache { return decodeLinearLayer(x, cache: mambaCache) } @@ -1105,7 +1106,9 @@ public class Qwen35TextModelInner: Module, LayerPartitionable, StreamableMoE { /// everything else on a decode step is static-shaped, so the segments /// compile concretely. private func decodeStep(_ inputs: MLXArray, _ cache: [KVCache?]) -> MLXArray? { - guard cache.count == layers.count else { return nil } + guard cache.count == layers.count, !ExpertStreamingConfig.shared.isEnabled else { + return nil + } // The schedule is only valid when the masks the general path would // build both come out empty. if createSSMMask(h: inputs, cache: cache[ssmIdx] as? MambaCache) != nil { return nil } diff --git a/Libraries/MLXLMCommon/SwitchLayers.swift b/Libraries/MLXLMCommon/SwitchLayers.swift index 19d41fb533..9dae429ad8 100644 --- a/Libraries/MLXLMCommon/SwitchLayers.swift +++ b/Libraries/MLXLMCommon/SwitchLayers.swift @@ -1068,9 +1068,7 @@ open class SwitchGLU: Module, @unchecked Sendable { x = downProj.computeExperts(intermediate, buffers: downBuffers, ranges: ranges) } - if doSort { - x = scatterUnsort(x: x, invOrder: inverseOrder, shape: indices.shape) - } + // Callers unsort with `inverseOrder`; unsorting here too scrambled the rows. return (x, doSort ? inverseOrder : nil) } From e3083ccb885eaf715995fb7bb669ea0420cccb21 Mon Sep 17 00:00:00 2001 From: simba Date: Fri, 25 Sep 2026 13:14:11 -0700 Subject: [PATCH 2/3] fix(load): load lazily when SSD expert streaming is on The concurrent weight loader from the upstream sync evaluates every tensor, including the experts SSD streaming is meant to page in on demand. On a 32 GB Mac that held ~20 GB resident, grew swap, and cut Qwen3.6-35B-A3B SSD prefill from 315 to 77 tok/s. Keep the per-file lazy load while streaming. Co-Authored-By: Claude Opus 5.5 --- Libraries/MLXLMCommon/Load.swift | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/Libraries/MLXLMCommon/Load.swift b/Libraries/MLXLMCommon/Load.swift index bf966bb92f..57995adfee 100644 --- a/Libraries/MLXLMCommon/Load.swift +++ b/Libraries/MLXLMCommon/Load.swift @@ -380,7 +380,17 @@ public func loadWeights( in: modelDirectory, selection: weightFileSelection, additionalFiles: additionalFiles ?? []) - (weights, metadata) = try loadWeightArrays(urls: weightURLs) + if ExpertStreamingConfig.shared.isEnabled { + // Load lazily: the concurrent loader reads every tensor, including the experts + // that SSD streaming pages in on demand, which fills RAM and pushes into swap. + for url in weightURLs { + let (w, m) = try loadArraysAndMetadata(url: url) + weights.merge(w) { _, new in new } + if metadata.isEmpty { metadata = m } + } + } else { + (weights, metadata) = try loadWeightArrays(urls: weightURLs) + } // MTP add-ons (e.g. Qwen3.5/3.6-OptiQ's `optiq/mtp.safetensors`) live in a // subdirectory `safetensorWeightURLs` deliberately excludes by default (see its From ec91ae1190871596faa0e45717b4b58249369340 Mon Sep 17 00:00:00 2001 From: simba Date: Fri, 25 Sep 2026 13:26:09 -0700 Subject: [PATCH 3/3] fix(qwen3next): skip compiled decode segments while SSD streaming Same failure as the Qwen3.5 decode path: fp16 Qwen3-Next runs SwitchGLU inside compiled decode segments, and the SSD streaming path evals there. Found by code reading (SwiftLM review session); not reproduced, since no fp16 Qwen3-Next checkpoint is on the test machine. Co-Authored-By: Claude Opus 5.5 --- Libraries/MLXLLM/Models/Qwen3Next.swift | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/Libraries/MLXLLM/Models/Qwen3Next.swift b/Libraries/MLXLLM/Models/Qwen3Next.swift index 3192e5b567..41a07f56ea 100644 --- a/Libraries/MLXLLM/Models/Qwen3Next.swift +++ b/Libraries/MLXLLM/Models/Qwen3Next.swift @@ -684,7 +684,10 @@ public class Qwen3NextModelInner: Module { // this model. fp16 is pinned byte-identical against the general path. let embedded = embedTokens(inputs) guard embedded.dtype == .float16 else { return nil } - guard cache.count == layers.count else { return nil } + // SSD expert streaming evals mid-layer, which a trace cannot contain. + guard cache.count == layers.count, !ExpertStreamingConfig.shared.isEnabled else { + return nil + } if createSSMMask(h: inputs, cache: cache[ssmIdx] as? MambaCache) != nil { return nil }