feat: expose embedTokens, lmHead, callCapturing for DFlash conformance on Llama/Qwen3/Qwen3MoE

Enables Sources/SwiftLM/{Llama,Qwen3,Qwen3MoE}+DFlash.swift to conform
to DFlashTargetModel by exposing:
  - LlamaModelInner.embedTokens (public)
  - LlamaModel.lmHead (public)
  - LlamaModel.callCapturing(_:cache:captureLayerIDs:)
  - Qwen3ModelInner.embedTokens + layers (public)
  - Qwen3Model.lmHead (public)
  - Qwen3Model.callCapturing(_:cache:captureLayerIDs:)
  - Qwen3MoEModelInner.embedTokens + layers (public)
  - Qwen3MoEModel.lmHead (public)
  - Qwen3MoEModel.callCapturing(_:cache:captureLayerIDs:)

Authored by: 0xClandestine (original design in unreachable commit b5762584)

Author: Aegis-AI

Libraries/MLXLLM/Models/Llama.swift

@@ -120,7 +120,7 @@ class LlamaTransformerBlock: Module {
 
 public class LlamaModelInner: Module, LayerPartitionable {
 
-    @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
+    @ModuleInfo(key: "embed_tokens") public var embedTokens: Embedding
 
     let layers: [LlamaTransformerBlock]
     let norm: RMSNorm
@@ -162,7 +162,7 @@ public class LlamaModel: Module, LLMModel, KVCacheDimensionProvider {
 
     public let model: LlamaModelInner
 
-    @ModuleInfo(key: "lm_head") var lmHead: Linear?
+    @ModuleInfo(key: "lm_head") public var lmHead: Linear?
 
     public init(_ args: LlamaConfiguration) {
         self.vocabularySize = args.vocabularySize
@@ -182,6 +182,21 @@ public class LlamaModel: Module, LLMModel, KVCacheDimensionProvider {
         }
     }
 
+    public func callCapturing(
+        _ 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)
+        var captured: [Int: MLXArray] = [:]
+        for (i, layer) in model.layers.enumerated() {
+            h = layer(h, mask: mask, cache: kvCache?[i])
+            if captureLayerIDs.contains(i) { captured[i] = h }
+        }
+        h = model.norm(h)
+        return (h, captured)
+    }
+
     public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
         // Remove unused precomputed rotary frequencies
         weights.filter {

Libraries/MLXLLM/Models/Qwen3.swift

@@ -145,9 +145,9 @@ public class Qwen3ModelInner: Module, LayerPartitionable {
     // LayerPartitionable
     public var gpuLayerCount: Int?
     public var totalLayerCount: Int { layers.count }
-    @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
+    @ModuleInfo(key: "embed_tokens") public var embedTokens: Embedding
 
-    fileprivate let layers: [Qwen3TransformerBlock]
+    public let layers: [Qwen3TransformerBlock]
     let norm: RMSNorm
 
     public init(_ args: Qwen3Configuration) {
@@ -185,7 +185,7 @@ public class Qwen3Model: Module, LLMModel, KVCacheDimensionProvider {
     public let model: Qwen3ModelInner
     let configuration: Qwen3Configuration
 
-    @ModuleInfo(key: "lm_head") var lmHead: Linear?
+    @ModuleInfo(key: "lm_head") public var lmHead: Linear?
 
     public init(_ args: Qwen3Configuration) {
         self.configuration = args
@@ -208,6 +208,21 @@ public class Qwen3Model: Module, LLMModel, KVCacheDimensionProvider {
         return out
     }
 
+    public func callCapturing(
+        _ 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)
+        var captured: [Int: MLXArray] = [:]
+        for (i, layer) in model.layers.enumerated() {
+            h = layer(h, mask: mask, cache: kvCache?[i])
+            if captureLayerIDs.contains(i) { captured[i] = h }
+        }
+        h = model.norm(h)
+        return (h, captured)
+    }
+
     public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
         var weights = weights
 

Libraries/MLXLLM/Models/Qwen3MoE.swift

@@ -196,9 +196,9 @@ public class Qwen3MoEModelInner: Module, LayerPartitionable, StreamableMoE {
     // StreamableMoE
     public var streamExperts: Bool = false
 
-    @ModuleInfo(key: "embed_tokens") var embedTokens: Embedding
+    @ModuleInfo(key: "embed_tokens") public var embedTokens: Embedding
 
-    fileprivate let layers: [Qwen3MoeDecoderLayer]
+    public let layers: [Qwen3MoeDecoderLayer]
     let norm: RMSNorm
     let args: Qwen3MoEConfiguration
 
@@ -238,7 +238,7 @@ public class Qwen3MoEModel: Module, LLMModel, KVCacheDimensionProvider {
     public let model: Qwen3MoEModelInner
     let configuration: Qwen3MoEConfiguration
 
-    @ModuleInfo(key: "lm_head") var lmHead: Linear?
+    @ModuleInfo(key: "lm_head") public var lmHead: Linear?
 
     public init(_ args: Qwen3MoEConfiguration) {
         self.configuration = args
@@ -261,6 +261,21 @@ public class Qwen3MoEModel: Module, LLMModel, KVCacheDimensionProvider {
         return out
     }
 
+    public func callCapturing(
+        _ 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)
+        var captured: [Int: MLXArray] = [:]
+        for (i, layer) in model.layers.enumerated() {
+            h = layer(h, mask: mask, cache: kvCache?[i])
+            if captureLayerIDs.contains(i) { captured[i] = h }
+        }
+        h = model.norm(h)
+        return (h, captured)
+    }
+
     public func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] {
         var sanitizedWeights = weights