perf(mlx-lm): Gemma 4 v0.8.2 — shared tower, layer-18 submission, coupled weighted-unsort+R1, packed multimodal

Libraries/MLXLLM/Models/Gemma4MTPTarget.swift

@@ -7,14 +7,10 @@ import MLXLMCommon
 /// Abstraction over a Gemma 4 text tower that can drive MTP speculative
 /// decoding.
 ///
-/// Both the text-only ``Gemma4TextModel`` (MLXLLM) and the vision-language
-/// `Gemma4` tower (MLXVLM) conform, so the single-stream MTP round loop, the
-/// token iterator, and the drafter binding all work against either tower. The
-/// MTP drafter (``Gemma4AssistantDraftModel``) is trained against the Gemma 4
-/// text architecture; because the VLM tower implements the *same* text
-/// architecture and loads the *same* text weights, a drafter bound to a VLM
-/// tower produces the same speculative tokens it would against the text-only
-/// tower (validated by the parity spike).
+/// ``Gemma4TextModel`` is the canonical target for text-only loads and for
+/// MLXVLM Gemma 4: the VLM owns and exposes this exact object as `textModel`.
+/// Single-stream and CBv2 MTP therefore bind to the same text architecture,
+/// weights, hidden capture, and cache identity used by direct VLM forwards.
 public protocol Gemma4MTPTarget: AnyObject {
 
     /// The resolved text configuration, used for drafter-compatibility
@@ -40,7 +36,7 @@ public protocol Gemma4MTPTarget: AnyObject {
         _ caches: [KVCache], accepted: Gemma4AcceptCount, blockSize: Int)
 }
 
-// MARK: - Text-only tower conformance
+// MARK: - Shared text-tower conformance
 
 extension Gemma4TextModel: Gemma4MTPTarget {
     public var mtpConfiguration: Gemma4TextConfiguration { configuration }

Libraries/MLXLMCommon/BaseConfiguration.swift

@@ -205,6 +205,26 @@ public struct BaseConfiguration: Codable, Sendable {
     enum CodingKeys: String, CodingKey {
         case modelType = "model_type"
         case quantizationContainer = "quantization"
+        case quantizationConfiguration = "quantization_config"
         case eosTokenIds = "eos_token_id"
     }
+
+    public init(from decoder: any Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        modelType = try container.decode(String.self, forKey: .modelType)
+        quantizationContainer =
+            try container.decodeIfPresent(
+                QuantizationContainer.self, forKey: .quantizationContainer)
+            ?? container.decodeIfPresent(
+                QuantizationContainer.self, forKey: .quantizationConfiguration)
+        eosTokenIds = try container.decodeIfPresent(IntOrIntArray.self, forKey: .eosTokenIds)
+    }
+
+    public func encode(to encoder: any Encoder) throws {
+        var container = encoder.container(keyedBy: CodingKeys.self)
+        try container.encode(modelType, forKey: .modelType)
+        try container.encodeIfPresent(
+            quantizationContainer, forKey: .quantizationContainer)
+        try container.encodeIfPresent(eosTokenIds, forKey: .eosTokenIds)
+    }
 }

Libraries/MLXLMCommon/ContinuousBatchingV2/EngineLoopV2.swift

@@ -55,11 +55,16 @@ public protocol CBv2LayerCacheProvider: AnyObject {
     /// equal-length text chunks may be coalesced into one layer-major
     /// forward. Fail-safe default false (paged/custom providers).
     var supportsPackedPrefill: Bool { get }
+    /// True only when every layer cache can bind one optional span context
+    /// per rectangular row. This is stricter than single-row multimodal
+    /// support and fails closed for paged/custom providers.
+    var supportsPackedMultimodalSpans: Bool { get }
 }
 
 extension CBv2LayerCacheProvider {
     public var supportsMultimodalSpans: Bool { false }
     public var supportsPackedPrefill: Bool { false }
+    public var supportsPackedMultimodalSpans: Bool { false }
 }
 
 // MARK: - Sampler interface (WS-E's CBv2DefaultSampler is the production impl)
@@ -1368,10 +1373,10 @@ public final class EngineLoopV2: @unchecked Sendable {
 
         // Prompt chunks. Default shape is per-request [1, chunk]; when the
         // model AND the cache provider both prove rectangular per-row
-        // semantics, equal-length TEXT chunks are coalesced into one
-        // layer-major [B, chunk] forward so each layer's weights are read
-        // once for the whole cohort. Span-bearing multimodal chunks always
-        // stay per-request.
+        // semantics, equal-length chunks are coalesced into one layer-major
+        // [B, chunk] forward so each layer's weights are read once for the
+        // whole cohort. Span-bearing rows require the stronger model and
+        // cache capabilities for row-local embeddings and attention masks.
         var prefillSampled: [CBv2RequestID: MLXArray] = [:]
         var evalTargets: [MLXArray] = []
         var packedIDs = Set<CBv2RequestID>()
@@ -1380,20 +1385,23 @@ public final class EngineLoopV2: @unchecked Sendable {
             let packedModel = model as? CBv2PackedPrefillSteppableModel,
             packedModel.supportsPackedPrefill
         {
+            let canPackMultimodal =
+                packedModel.supportsPackedMultimodalPrefill
+                && cacheProvider.supportsPackedMultimodalSpans
             struct PackedGroup {
                 let count: Int
                 let samples: Bool
                 var rows: [RowWork]
             }
             var groups: [PackedGroup] = []
             for row in work where !row.isDecode {
-                // Pure function of has-spans, exactly like the singleton
-                // path: a multimodal REQUEST whose current chunk carries no
-                // span is still packable.
+                // A multimodal request's text-only chunks remain packable.
+                // A span-bearing chunk needs explicit rectangular embedding
+                // and row-mask capability from both model and cache provider.
                 let hasSpan =
                     multimodalByID[row.rec.id]?.chunkContext(
                         start: row.start, count: row.count) != nil
-                if hasSpan { continue }
+                if hasSpan && !canPackMultimodal { continue }
                 if let index = groups.firstIndex(where: {
                     $0.count == row.count && $0.samples == row.samples
                 }) {
@@ -1416,9 +1424,24 @@ public final class EngineLoopV2: @unchecked Sendable {
                 let caches = eagerCaches(rowStates: group.rows.map { kvStates[$0.rec.id]! })
                 let requirement: CBv2PrefillRequirement =
                     group.samples ? .lastPositionLogits : .evaluationOnly
-                let output = prefillOutput(
-                    tokens: inputs, inputEmbeddings: nil, caches: caches,
-                    requirement: requirement)
+                let spanContexts = group.rows.map {
+                    multimodalByID[$0.rec.id]?.chunkContext(
+                        start: $0.start, count: $0.count)
+                }
+                let output: MLXArray
+                if spanContexts.contains(where: { $0 != nil }) {
+                    output = packedMultimodalChunksForward(
+                        tokens: inputs,
+                        starts: group.rows.map(\.start),
+                        multimodal: group.rows.map { multimodalByID[$0.rec.id] },
+                        spanContexts: spanContexts,
+                        caches: caches,
+                        requirement: requirement)
+                } else {
+                    output = prefillOutput(
+                        tokens: inputs, inputEmbeddings: nil, caches: caches,
+                        requirement: requirement)
+                }
                 cacheInnerState.append(contentsOf: eagerCacheInnerState(caches))
 
                 if group.samples {
@@ -1569,6 +1592,58 @@ public final class EngineLoopV2: @unchecked Sendable {
             requirement: requirement)
     }
 
+    /// Rectangular counterpart of `multimodalChunkForward`. Each row keeps
+    /// its own token embeddings, image splice coordinates, KV state, and
+    /// optional span mask while the model traverses the cohort layer-major.
+    func packedMultimodalChunksForward(
+        tokens: MLXArray,
+        starts: [Int],
+        multimodal: [CBv2ResolvedMultimodal?],
+        spanContexts: [CBv2SpanChunkContext?],
+        caches: [CBv2AttendingLayerCache],
+        requirement: CBv2PrefillRequirement
+    ) -> MLXArray {
+        guard let mmModel = model as? CBv2MultimodalSteppableModel else {
+            preconditionFailure(
+                "CBv2 packed multimodal chunk reached a model without embedding-forward support")
+        }
+        let batch = tokens.dim(0)
+        let count = tokens.dim(1)
+        precondition(
+            starts.count == batch && multimodal.count == batch
+                && spanContexts.count == batch,
+            "CBv2 packed multimodal metadata must match batch \(batch)")
+
+        let textEmbeddings = mmModel.embedPromptTokens(tokens)
+        var embeddingRows: [MLXArray] = []
+        embeddingRows.reserveCapacity(batch)
+        for index in 0 ..< batch {
+            let textRow = textEmbeddings[index ..< index + 1]
+            if spanContexts[index] != nil, let rowMultimodal = multimodal[index] {
+                embeddingRows.append(
+                    CBv2MultimodalPlan.spliceEmbeddings(
+                        textEmbeddings: textRow,
+                        chunkStart: starts[index],
+                        spans: rowMultimodal.spansInChunk(
+                            start: starts[index], count: count)))
+            } else {
+                embeddingRows.append(textRow)
+            }
+        }
+        let spliced = concatenated(embeddingRows, axis: 0)
+
+        let bindables =
+            count > 1 ? caches.compactMap { $0 as? CBv2PackedSpanMaskBinding } : []
+        precondition(
+            count == 1 || bindables.count == caches.count,
+            "CBv2 packed multimodal prefill requires per-row span binding on every cache")
+        for bindable in bindables { bindable.bindSpanContexts(spanContexts) }
+        defer { for bindable in bindables { bindable.bindSpanContexts(nil) } }
+        return prefillOutput(
+            tokens: tokens, inputEmbeddings: spliced, caches: caches,
+            requirement: requirement)
+    }
+
     // MARK: Finalization (deferred stop detection)
 
     private func finalize(_ step: CBv2InFlightStep, now: ContinuousClock.Instant) {

Libraries/MLXLMCommon/ContinuousBatchingV2/LayerCacheBankV2.swift

@@ -122,6 +122,12 @@ public final class CBv2LayerCacheBank: CBv2LayerCacheProvider, CBv2CompositionIn
         caches.allSatisfy { $0 is CBv2LayerCache }
     }
 
+    /// Packed vision rows additionally require one independently bound
+    /// optional span context per row on every owning and borrowing layer.
+    public var supportsPackedMultimodalSpans: Bool {
+        caches.allSatisfy { $0 is CBv2PackedSpanMaskBinding }
+    }
+
     public func layerCaches(rowStates: [[CBv2SequenceKV?]]) -> [CBv2AttendingLayerCache] {
         let identity = rowStates.map { row -> ObjectIdentifier in
             guard let anchor = row.compactMap({ $0 }).first else {

Libraries/MLXLMCommon/ContinuousBatchingV2/LayerCacheV2.swift

@@ -58,15 +58,11 @@ public final class CBv2LayerCache: CBv2AttendingLayerCache {
     /// the same parameter); never part of the per-call contract surface.
     public let attentionSoftcap: Float?
 
-    /// Vision span context for the CURRENT span-containing prefill chunk,
-    /// bound by the engine immediately before that chunk's graph build and
-    /// unbound immediately after (`CBv2SpanMaskBinding`). ALWAYS nil outside
-    /// that window, so decode, text chunks, and other requests' chunks take
-    /// the untouched pinned paths. The engine only ever binds it while the
-    /// cache's rows are exactly the one vision request's row (prefill is
-    /// per-request [1, chunk]), so the span mask can never leak onto a
-    /// batchmate's computation.
-    private(set) var boundSpanContext: CBv2SpanChunkContext?
+    /// Optional vision span context for each CURRENT prefill row. The engine
+    /// binds this array immediately before graph construction and clears it
+    /// immediately after. nil outside that window; nil entries are ordinary
+    /// text rows sharing a rectangular call.
+    private(set) var boundSpanContexts: [CBv2SpanChunkContext?]?
 
     public init(
         layerIndex: Int, kind: CBv2LayerKind, rows: [CBv2SequenceKV] = [],
@@ -120,10 +116,10 @@ public final class CBv2LayerCache: CBv2AttendingLayerCache {
             rows: rows, kind: kind,
             queries: queries, keys: keys, values: values,
             scale: scale, sinks: sinks, softcap: attentionSoftcap,
-            spanContext: boundSpanContext,
+            spanContexts: boundSpanContexts,
             serializeQueries: mtpSerializesRectangularAttention)
-        // Advance offsets ON-DEVICE (uniform: decode is [B,1], prefill is
-        // [1,chunk] — L is the same for every row in the call).
+        // Advance offsets ON-DEVICE. Decode and packed prefill are
+        // rectangular, so L is uniform across every bound row.
         cachedPositionOffsets = cachedPositionOffsets + Int32(queries.dim(2))
         return output
     }
@@ -164,7 +160,7 @@ public final class CBv2LayerCache: CBv2AttendingLayerCache {
         return CBv2AttentionV1.attendBorrowing(
             sourceRows: source.rows, sourceKind: source.kind, kind: kind,
             queries: queries, scale: scale, sinks: sinks, softcap: attentionSoftcap,
-            spanContext: boundSpanContext,
+            spanContexts: boundSpanContexts,
             serializeQueries: mtpSerializesRectangularAttention)
     }
 
@@ -187,9 +183,13 @@ extension CBv2LayerCache: CBv2LastQueryPrefillLayerCache {}
 
 // MARK: - Vision span-mask binding
 
-extension CBv2LayerCache: CBv2SpanMaskBinding {
+extension CBv2LayerCache: CBv2PackedSpanMaskBinding {
     public func bindSpanContext(_ context: CBv2SpanChunkContext?) {
-        boundSpanContext = context
+        boundSpanContexts = context.map { [$0] }
+    }
+
+    public func bindSpanContexts(_ contexts: [CBv2SpanChunkContext?]?) {
+        boundSpanContexts = contexts
     }
 }
 

Libraries/MLXLMCommon/ContinuousBatchingV2/MultimodalV2.swift

@@ -13,20 +13,20 @@
 //  2. SCHEDULES prefill chunks that never split a block (SchedulerV2 snaps
 //     chunk boundaries to block edges; a block longer than
 //     `maxBatchedTokensPerStep` is rejected at submit).
-//  3. EXECUTES span-containing chunks on a NEW pinned attention path:
-//     input embeddings = scaled text-token embeddings with the span
-//     embeddings spliced verbatim (`spliceEmbeddings`), and a per-chunk
+//  3. EXECUTES span-containing chunks through the embedding-forward path:
+//     input embeddings = scaled text-token embeddings with each row's span
+//     embeddings spliced verbatim (`spliceEmbeddings`), and a row-local
 //     boolean mask implementing causal-plus-bidirectional-within-blocks —
 //     the exact semantics of MLXVLM Gemma4's
-//     `gemma4BidirectionalVisionMask` overlay (`_apply_blockwise_
-//     bidirectional_overlay` in the Python reference): tokens of the same
-//     image block attend each other in BOTH directions, overriding both
-//     causality and the sliding window; everything else stays causal
-//     (∧ window). Text-only chunks — including the text chunks of a vision
-//     request — keep the existing maskless/causal path untouched, and
-//     decode after prefill is UNCHANGED (image tokens are history in KV by
-//     then). The path choice is a pure function of has-spans-in-chunk;
-//     never data-dependent.
+//     `gemma4BidirectionalVisionMask` overlay. Tokens of the same image block
+//     attend each other in BOTH directions, overriding causality and the
+//     sliding window; everything else stays causal (∧ window). Long rows
+//     retain q=128 query blocking, expanding only touched blocks' K/V slices
+//     enough to include the complete image span. Explicit model + cache
+//     capabilities allow independently spliced vision rows and nil-context
+//     text neighbors to share a rectangular call. Decode after prefill is
+//     unchanged (image tokens are history in KV by then). The path choice is
+//     a pure function of has-spans-in-chunk; never data-dependent.
 //  4. EXCLUDES vision requests from prefix-cache lookup AND donation
 //     (v1 policy): token-id chain hashes cannot see image content, so a
 //     hit would silently reuse the wrong KV. Image-digest extra keys in the
@@ -67,6 +67,15 @@ public protocol CBv2SpanMaskBinding: AnyObject {
     func bindSpanContext(_ context: CBv2SpanChunkContext?)
 }
 
+/// Stronger binding contract for rectangular multimodal prefill. The array
+/// is row-aligned with the cache's current batch: non-nil entries carry that
+/// row's vision overlay and nil entries retain ordinary text attention.
+/// Custom and paged providers must opt in explicitly; structural
+/// `CBv2SpanMaskBinding` conformance alone is insufficient.
+public protocol CBv2PackedSpanMaskBinding: CBv2SpanMaskBinding {
+    func bindSpanContexts(_ contexts: [CBv2SpanChunkContext?]?)
+}
+
 // MARK: - Model surfaces
 
 /// Steppable models that can prefill from spliced input embeddings.

Libraries/MLXLMCommon/ContinuousBatchingV2/PrefillOutputV2.swift

@@ -57,17 +57,23 @@ public protocol CBv2PrefillSteppableModel: CBv2SteppableModel {
 }
 
 /// Opt-in refinement for models whose prompt forward keeps per-row
-/// semantics when several EQUAL-LENGTH text chunks are executed as one
+/// semantics when several EQUAL-LENGTH chunks are executed as one
 /// rectangular `[B, L]` pass. This makes the transformer traversal
 /// layer-major across those rows: each layer's weights are read once for
 /// the whole cohort instead of once per row.
 ///
 /// This is a claim about the MODEL only. The engine additionally requires
 /// the cache provider to vouch that its layer caches keep independent rows
-/// (`CBv2LayerCacheProvider.supportsPackedPrefill`), and it never packs
-/// span-bearing multimodal chunks, which stay on the per-request path.
+/// (`CBv2LayerCacheProvider.supportsPackedPrefill`). Packing rows that splice
+/// image embeddings and carry row-local span masks requires the stronger,
+/// separately fail-closed `supportsPackedMultimodalPrefill` claim.
 public protocol CBv2PackedPrefillSteppableModel: CBv2PrefillSteppableModel {
     var supportsPackedPrefill: Bool { get }
+    var supportsPackedMultimodalPrefill: Bool { get }
+}
+
+extension CBv2PackedPrefillSteppableModel {
+    public var supportsPackedMultimodalPrefill: Bool { false }
 }
 
 /// Model-level (KVCache-shaped) twin of `CBv2PrefillSteppableModel`, for
@@ -85,8 +91,13 @@ public protocol CBv2LanguageModelPrefillForwardable {
     /// Whether this model's prompt forward is safe to run as a rectangular
     /// `[B, L]` cohort of independent rows. Fail-closed default: false.
     var cbv2SupportsPackedPrefill: Bool { get }
+
+    /// Stronger claim for rectangular embedding-forward rows with one
+    /// optional vision span-mask context per row. Fail-closed default: false.
+    var cbv2SupportsPackedMultimodalPrefill: Bool { get }
 }
 
 extension CBv2LanguageModelPrefillForwardable {
     public var cbv2SupportsPackedPrefill: Bool { false }
+    public var cbv2SupportsPackedMultimodalPrefill: Bool { false }
 }