Merge pull request #30 from SharpAI/feat/qwen3next-dflash-support

solderzzc · web-flow · commit ef3318e4dacf · 2026-04-23T14:38:19.000-07:00
feat: Qwen3Next and Qwen35 DFlash Speculative Decoding Support
diff --git a/Libraries/MLXLLM/Models/Qwen35.swift b/Libraries/MLXLLM/Models/Qwen35.swift
@@ -574,7 +574,7 @@ final class Qwen35DecoderLayer: Module {
 // MARK: - Text Model
 
 public class Qwen35TextModelInner: Module, LayerPartitionable, StreamableMoE {
-    @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
+    @ModuleInfo(key: "embed_tokens") public var embedTokens: Embedding
 
     fileprivate let layers: [Qwen35DecoderLayer]
     let norm: MathRMSNorm
@@ -613,7 +613,7 @@ public class Qwen35TextModelInner: Module, LayerPartitionable, StreamableMoE {
         var hiddenStates = embedTokens(inputs)
 
         var cacheArray = cache
-        if cacheArray == nil {
+        if cacheArray == nil || cacheArray?.count != layers.count {
             cacheArray = Array(repeating: nil as KVCache?, count: layers.count)
         }
 
@@ -633,6 +633,37 @@ public class Qwen35TextModelInner: Module, LayerPartitionable, StreamableMoE {
 
         return norm(hiddenStates)
     }
+
+    public func callCapturing(_ inputs: MLXArray, cache: [KVCache?]? = nil, captureLayerIDs: Set<Int>) -> (MLXArray, [Int: MLXArray]) {
+        var hiddenStates = embedTokens(inputs)
+
+        var cacheArray = cache
+        if cacheArray == nil || cacheArray?.count != layers.count {
+            cacheArray = Array(repeating: nil as KVCache?, count: layers.count)
+        }
+
+        let faMask = createAttentionMask(h: hiddenStates, cache: cacheArray?[faIdx])
+        let ssmMask = createSSMMask(h: hiddenStates, cache: cacheArray?[ssmIdx] as? MambaCache)
+
+        var captured = [Int: MLXArray]()
+
+        for (i, layer) in layers.enumerated() {
+            let mask = layer.isLinear ? ssmMask : nil
+            let attnMask =
+                layer.isLinear
+                ? MLXFast.ScaledDotProductAttentionMaskMode.none : faMask
+            hiddenStates = partitionedLayerCall(index: i, gpuLayerCount: gpuLayerCount, stream: streamExperts, cacheToEval: cacheArray?[i]) {
+                layer(
+                    hiddenStates, attentionMask: attnMask, ssmMask: mask, cache: cacheArray?[i])
+            }
+            
+            if captureLayerIDs.contains(i) {
+                captured[i] = hiddenStates
+            }
+        }
+
+        return (norm(hiddenStates), captured)
+    }
 }
 
 public class Qwen35TextModel: Module, LLMModel, KVCacheDimensionProvider {
@@ -642,7 +673,7 @@ public class Qwen35TextModel: Module, LLMModel, KVCacheDimensionProvider {
     public let model: Qwen35TextModelInner
     let configuration: Qwen35TextConfiguration
 
-    @ModuleInfo(key: "lm_head") var lmHead: Linear?
+    @ModuleInfo(key: "lm_head") public var lmHead: Linear?
 
     public init(_ args: Qwen35TextConfiguration) {
         self.configuration = args
@@ -725,7 +756,7 @@ public class Qwen35Model: Module, LLMModel, KVCacheDimensionProvider {
     public let vocabularySize: Int
     public let kvHeads: [Int]
 
-    @ModuleInfo(key: "language_model") var languageModel: Qwen35TextModel
+    @ModuleInfo(key: "language_model") public var languageModel: Qwen35TextModel
 
     public init(_ args: Qwen35Configuration) {
         let textModel = Qwen35TextModel(args.textConfig)
diff --git a/Libraries/MLXLLM/Models/Qwen3Next.swift b/Libraries/MLXLLM/Models/Qwen3Next.swift
@@ -419,7 +419,7 @@ final class Qwen3NextDecoderLayer: Module {
 }
 
 public class Qwen3NextModelInner: Module {
-    @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
+    @ModuleInfo(key: "embed_tokens") public var embedTokens: Embedding
 
     fileprivate let layers: [Qwen3NextDecoderLayer]
     let norm: RMSNorm
@@ -451,7 +451,7 @@ public class Qwen3NextModelInner: Module {
         var hiddenStates = embedTokens(inputs)
 
         var cacheArray = cache
-        if cacheArray == nil {
+        if cacheArray == nil || cacheArray?.count != layers.count {
             cacheArray = Array(repeating: nil as KVCache?, count: layers.count)
         }
 
@@ -467,6 +467,33 @@ public class Qwen3NextModelInner: Module {
 
         return norm(hiddenStates)
     }
+
+    public func callCapturing(_ inputs: MLXArray, cache: [KVCache?]? = nil, captureLayerIDs: Set<Int>) -> (MLXArray, [Int: MLXArray]) {
+        var hiddenStates = embedTokens(inputs)
+
+        var cacheArray = cache
+        if cacheArray == nil || cacheArray?.count != layers.count {
+            cacheArray = Array(repeating: nil as KVCache?, count: layers.count)
+        }
+
+        let faMask = createAttentionMask(h: hiddenStates, cache: cacheArray?[faIdx])
+        let ssmMask = createSSMMask(h: hiddenStates, cache: cacheArray?[ssmIdx] as? MambaCache)
+
+        var captured = [Int: MLXArray]()
+
+        for (i, layer) in layers.enumerated() {
+            let mask = layer.isLinear ? ssmMask : nil
+            let attnMask = layer.isLinear ? MLXFast.ScaledDotProductAttentionMaskMode.none : faMask
+            hiddenStates = layer(
+                hiddenStates, attentionMask: attnMask, ssmMask: mask, cache: cacheArray?[i])
+            
+            if captureLayerIDs.contains(i) {
+                captured[i] = hiddenStates
+            }
+        }
+
+        return (norm(hiddenStates), captured)
+    }
 }
 
 public class Qwen3NextModel: Module, LLMModel, KVCacheDimensionProvider {
@@ -476,7 +503,7 @@ public class Qwen3NextModel: Module, LLMModel, KVCacheDimensionProvider {
     public let model: Qwen3NextModelInner
     let configuration: Qwen3NextConfiguration
 
-    @ModuleInfo(key: "lm_head") var lmHead: Linear?
+    @ModuleInfo(key: "lm_head") public var lmHead: Linear?
 
     public init(_ args: Qwen3NextConfiguration) {
         self.configuration = args
diff --git a/Libraries/MLXLMCommon/KVCache.swift b/Libraries/MLXLMCommon/KVCache.swift
@@ -1251,7 +1251,7 @@ public class ChunkedKVCache: KVCacheSimple {
 }
 
 /// Base cache for array-based state storage
-public class ArraysCache: BaseKVCache {
+open class ArraysCache: BaseKVCache {
     private var cache: [MLXArray?]
     internal var leftPadding: MLXArray?
 
@@ -1265,12 +1265,12 @@ public class ArraysCache: BaseKVCache {
         cache.compactMap { $0 }
     }
 
-    public subscript(index: Int) -> MLXArray? {
+    open subscript(index: Int) -> MLXArray? {
         get { cache[index] }
         set { cache[index] = newValue }
     }
 
-    public override var state: [MLXArray] {
+    open override var state: [MLXArray] {
         get {
             return cache.compactMap { $0 }
         }
@@ -1279,7 +1279,7 @@ public class ArraysCache: BaseKVCache {
         }
     }
 
-    public override func copy() -> any KVCache {
+    open override func copy() -> any KVCache {
         let new = ArraysCache(size: cache.count)
         let s = self.state
         if !s.isEmpty {
@@ -1322,7 +1322,7 @@ public class ArraysCache: BaseKVCache {
 
     /// metaState format: [slotCount, presentSlots (comma-separated), leftPadding (comma-separated, optional)]
     /// Legacy format (BaseKVCache default): [""]
-    public override var metaState: [String] {
+    open override var metaState: [String] {
         get {
             var result = [
                 "\(cache.count)",
@@ -1379,7 +1379,7 @@ public class ArraysCache: BaseKVCache {
 }
 
 /// Simple cache for Mamba-style state space models
-public class MambaCache: ArraysCache {
+open class MambaCache: ArraysCache {
     /// Saved state for speculative decoding rollback.
     /// Mamba state is recurrent and cannot be partially "trimmed" like attention KV caches.
     /// Instead, we checkpoint before speculation and restore on rollback.
@@ -1390,10 +1390,10 @@ public class MambaCache: ArraysCache {
     }
 
     /// Mark as trimmable to enable speculative decoding on hybrid Attention+Mamba models.
-    public override var isTrimmable: Bool { true }
+    open override var isTrimmable: Bool { true }
 
     /// Save a checkpoint of the current Mamba state (call before speculative draft round).
-    public func checkpoint() {
+    open func checkpoint() {
         let s = self.state
         if !s.isEmpty {
             savedState = s.map { $0[.ellipsis] }  // deep copy
@@ -1404,7 +1404,7 @@ public class MambaCache: ArraysCache {
     /// When n > 0, rejected draft tokens have polluted the state — restore checkpoint.
     /// When n == 0, all drafts accepted — keep current state and clear checkpoint.
     @discardableResult
-    public override func trim(_ n: Int) -> Int {
+    open override func trim(_ n: Int) -> Int {
         if n > 0, let saved = savedState {
             self.state = saved
             savedState = nil
@@ -1414,7 +1414,7 @@ public class MambaCache: ArraysCache {
         return 0
     }
 
-    public override func copy() -> any KVCache {
+    open override func copy() -> any KVCache {
         let new = MambaCache()
         let s = self.state
         if !s.isEmpty {
diff --git a/Libraries/MLXLMCommon/SwitchLayers.swift b/Libraries/MLXLMCommon/SwitchLayers.swift
@@ -144,6 +144,7 @@ public class SwitchGLU: Module, @unchecked Sendable {
                     _persistentUp = qUp.allocateExpertBuffers(maxBuffers)
                     _persistentDown = qDown.allocateExpertBuffers(maxBuffers)
 
+
                     // Merged eval: idx + buffer allocations (same as ssd-opt-v1)
                     var toEval: [MLXArray] = [idx]
                     toEval.append(contentsOf: _persistentGate!)
diff --git a/Tests/MLXLMTests/Qwen35Tests.swift b/Tests/MLXLMTests/Qwen35Tests.swift
@@ -0,0 +1,48 @@
+import Foundation
+import MLX
+import MLXLLM
+import MLXLMCommon
+import MLXNN
+import Testing
+
+extension MLXTestingSuite {
+    @Suite
+    struct Qwen35Tests {
+
+        private func makeTinyConfigData() -> Data {
+            let json = """
+            {
+                "model_type": "qwen3_5",
+                "hidden_size": 64,
+                "num_hidden_layers": 4,
+                "intermediate_size": 128,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "rms_norm_eps": 1e-6,
+                "vocab_size": 100,
+                "rope_theta": 10000.0,
+                "max_position_embeddings": 512
+            }
+            """
+            return json.data(using: .utf8)!
+        }
+
+        @Test("Qwen35 callCapturing returns captured layers")
+        func testQwen35CallCapturing() throws {
+
+            let data = makeTinyConfigData()
+            let config = try JSONDecoder().decode(Qwen35Configuration.self, from: data)
+            let model = Qwen35Model(config)
+
+            let input = MLXArray(0..<8).reshaped(1, 8)
+            let captureLayerIDs: Set<Int> = [0, 1]
+            
+            let (hiddenStates, captured) = model.languageModel.model.callCapturing(input, captureLayerIDs: captureLayerIDs)
+
+            #expect(hiddenStates.shape == [1, 8, 64])
+            #expect(captured.count == 2)
+            #expect(captured[0]?.shape == [1, 8, 64])
+            #expect(captured[1]?.shape == [1, 8, 64])
+        }
+    }
+}
diff --git a/Tests/MLXLMTests/Qwen3NextTests.swift b/Tests/MLXLMTests/Qwen3NextTests.swift
@@ -0,0 +1,60 @@
+import Foundation
+import MLX
+import MLXLLM
+import MLXLMCommon
+import MLXNN
+import Testing
+
+extension MLXTestingSuite {
+    @Suite
+    struct Qwen3NextTests {
+
+        private func makeTinyConfigData() -> Data {
+            let json = """
+            {
+                "model_type": "qwen3_next",
+                "hidden_size": 64,
+                "num_hidden_layers": 4,
+                "intermediate_size": 128,
+                "num_attention_heads": 4,
+                "linear_num_value_heads": 4,
+                "linear_num_key_heads": 2,
+                "linear_key_head_dim": 16,
+                "linear_value_head_dim": 16,
+                "linear_conv_kernel_dim": 4,
+                "num_experts": 2,
+                "num_experts_per_tok": 1,
+                "decoder_sparse_step": 2,
+                "shared_expert_intermediate_size": 64,
+                "mlp_only_layers": [],
+                "moe_intermediate_size": 64,
+                "rms_norm_eps": 1e-6,
+                "vocab_size": 100,
+                "num_key_value_heads": 2,
+                "rope_theta": 10000.0,
+                "max_position_embeddings": 512,
+                "full_attention_interval": 2
+            }
+            """
+            return json.data(using: .utf8)!
+        }
+
+        @Test("Qwen3Next callCapturing returns captured layers")
+        func testQwen3NextCallCapturing() throws {
+
+            let data = makeTinyConfigData()
+            let config = try JSONDecoder().decode(Qwen3NextConfiguration.self, from: data)
+            let model = Qwen3NextModel(config)
+
+            let input = MLXArray(0..<8).reshaped(1, 8)
+            let captureLayerIDs: Set<Int> = [0, 2]
+            
+            let (hiddenStates, captured) = model.model.callCapturing(input, captureLayerIDs: captureLayerIDs)
+
+            #expect(hiddenStates.shape == [1, 8, 64])
+            #expect(captured.count == 2)
+            #expect(captured[0]?.shape == [1, 8, 64])
+            #expect(captured[2]?.shape == [1, 8, 64])
+        }
+    }
+}
diff --git a/Tests/MLXLMTests/ToolRegressionTests.swift b/Tests/MLXLMTests/ToolRegressionTests.swift
