fix: address Copilot review — safe cache normalization + partitionedLayerCall

- Replace compactMap{} cache handling with explicit normalization that
  preserves layer↔cache index alignment (nil-padding to layers.count)
- Wrap per-layer calls in partitionedLayerCall() to respect gpuLayerCount
  CPU/GPU routing (matches callAsFunction behavior)
- Qwen3MoE additionally passes stream:model.streamExperts to match
  its Qwen3MoEModelInner.callAsFunction routing

Author: Aegis-AI

Libraries/MLXLLM/Models/Llama.swift

@@ -186,11 +186,19 @@ public class LlamaModel: Module, LLMModel, KVCacheDimensionProvider {
         _ inputs: MLXArray, cache: [KVCache?]? = nil, captureLayerIDs: Set<Int>
     ) -> (MLXArray, [Int: MLXArray]) {
         var h = model.embedTokens(inputs)
-        let kvCache = cache?.compactMap { $0 }
-        let mask = createAttentionMask(h: h, cache: kvCache?.first)
+        let layerCount = model.layers.count
+        let kvCache: [KVCache?] = {
+            guard let c = cache else { return Array(repeating: nil, count: layerCount) }
+            var normalized: [KVCache?] = Array(repeating: nil, count: layerCount)
+            for (i, v) in c.prefix(layerCount).enumerated() { normalized[i] = v }
+            return normalized
+        }()
+        let mask = createAttentionMask(h: h, cache: kvCache.first ?? nil)
         var captured: [Int: MLXArray] = [:]
         for (i, layer) in model.layers.enumerated() {
-            h = layer(h, mask: mask, cache: kvCache?[i])
+            h = partitionedLayerCall(index: i, gpuLayerCount: model.gpuLayerCount) {
+                layer(h, mask: mask, cache: kvCache[i])
+            }
             if captureLayerIDs.contains(i) { captured[i] = h }
         }
         h = model.norm(h)

Libraries/MLXLLM/Models/Qwen3.swift

@@ -212,11 +212,19 @@ public class Qwen3Model: Module, LLMModel, KVCacheDimensionProvider {
         _ inputs: MLXArray, cache: [KVCache?]? = nil, captureLayerIDs: Set<Int>
     ) -> (MLXArray, [Int: MLXArray]) {
         var h = model.embedTokens(inputs)
-        let kvCache = cache?.compactMap { $0 }
-        let mask = createAttentionMask(h: h, cache: kvCache?.first)
+        let layerCount = model.layers.count
+        let kvCache: [KVCache?] = {
+            guard let c = cache else { return Array(repeating: nil, count: layerCount) }
+            var normalized: [KVCache?] = Array(repeating: nil, count: layerCount)
+            for (i, v) in c.prefix(layerCount).enumerated() { normalized[i] = v }
+            return normalized
+        }()
+        let mask = createAttentionMask(h: h, cache: kvCache.first ?? nil)
         var captured: [Int: MLXArray] = [:]
         for (i, layer) in model.layers.enumerated() {
-            h = layer(h, mask: mask, cache: kvCache?[i])
+            h = partitionedLayerCall(index: i, gpuLayerCount: model.gpuLayerCount) {
+                layer(h, mask: mask, cache: kvCache[i])
+            }
             if captureLayerIDs.contains(i) { captured[i] = h }
         }
         h = model.norm(h)

Libraries/MLXLLM/Models/Qwen3MoE.swift

@@ -265,11 +265,21 @@ public class Qwen3MoEModel: Module, LLMModel, KVCacheDimensionProvider {
         _ inputs: MLXArray, cache: [KVCache?]? = nil, captureLayerIDs: Set<Int>
     ) -> (MLXArray, [Int: MLXArray]) {
         var h = model.embedTokens(inputs)
-        let kvCache = cache?.compactMap { $0 }
-        let mask = createAttentionMask(h: h, cache: kvCache?.first)
+        let layerCount = model.layers.count
+        let kvCache: [KVCache?] = {
+            guard let c = cache else { return Array(repeating: nil, count: layerCount) }
+            var normalized: [KVCache?] = Array(repeating: nil, count: layerCount)
+            for (i, v) in c.prefix(layerCount).enumerated() { normalized[i] = v }
+            return normalized
+        }()
+        let mask = createAttentionMask(h: h, cache: kvCache.first ?? nil)
         var captured: [Int: MLXArray] = [:]
         for (i, layer) in model.layers.enumerated() {
-            h = layer(h, mask: mask, cache: kvCache?[i])
+            h = partitionedLayerCall(
+                index: i, gpuLayerCount: model.gpuLayerCount, stream: model.streamExperts
+            ) {
+                layer(h, mask: mask, cache: kvCache[i])
+            }
             if captureLayerIDs.contains(i) { captured[i] = h }
         }
         h = model.norm(h)