Fix Qwen3.5/3.6 SSD expert streaming after the upstream sync (crashes + full weight load)

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 }

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
         }

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

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)
         }