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

Libraries/MLXLLM/Models/Gemma4.swift

@@ -33,6 +33,7 @@ public struct Gemma4Configuration: Codable, Sendable {
     public init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
         let quantizationContainer = try decoder.container(keyedBy: QuantizationCodingKeys.self)
+        let rootPerLayerQuantization = (try? BaseConfiguration(from: decoder))?.perLayerQuantization
         self.modelType = try container.decodeIfPresent(String.self, forKey: .modelType) ?? "gemma4"
         self.vocabSize = try container.decodeIfPresent(Int.self, forKey: .vocabSize) ?? 262144
 
@@ -52,9 +53,12 @@ public struct Gemma4Configuration: Codable, Sendable {
                 Gemma4WeightQuantizationMetadata.self, forKey: .quantization)
             ?? quantizationContainer.decodeIfPresent(
                 Gemma4WeightQuantizationMetadata.self, forKey: .quantizationConfig)
-        if let quantization {
+        if let rootPerLayerQuantization {
+            self.textConfig.mergeQuantization(rootPerLayerQuantization)
+        } else if let quantization {
             self.textConfig.quantizationBits = quantization.bits
             self.textConfig.quantizationGroupSize = quantization.groupSize
+            self.textConfig.quantizationMode = quantization.mode ?? .affine
         }
     }
 }

Libraries/MLXLLM/Models/Gemma4MTPConfigurationValidation.swift

@@ -217,9 +217,10 @@ enum Gemma4AssistantConfigurationValidator {
                     "textConfig.numGlobalKeyValueHeads",
                     "is required when full attention uses K=V")
             }
-            let fullKVHeads = text.attentionKeqV
-                ? text.numGlobalKeyValueHeads!
-                : text.numKeyValueHeads
+            // Mirrors Gemma4Attention.init: full layers honor
+            // num_global_key_value_heads whenever present, independent of
+            // attention_k_eq_v (k_eq_v only elides v_proj).
+            let fullKVHeads = text.numGlobalKeyValueHeads ?? text.numKeyValueHeads
             try divides(
                 fullKVHeads,
                 into: text.numAttentionHeads,
@@ -266,9 +267,9 @@ enum Gemma4AssistantConfigurationValidator {
                 field: "textConfig.slidingAttention")
         }
         if hasFull {
-            let fullKVHeads = text.attentionKeqV
-                ? text.numGlobalKeyValueHeads!
-                : text.numKeyValueHeads
+            // Full layers honor num_global_key_value_heads whenever present,
+            // independent of attention_k_eq_v (mirrors Gemma4Attention.init).
+            let fullKVHeads = text.numGlobalKeyValueHeads ?? text.numKeyValueHeads
             try validateAttentionProducts(
                 heads: text.numAttentionHeads,
                 kvHeads: fullKVHeads,
@@ -584,12 +585,9 @@ enum Gemma4MTPCompatibilityValidator {
                     drafterText.attentionKeqV,
                     target.attentionKeqV,
                     field: "fullAttention.attentionKeqV")
-                if drafterText.attentionKeqV {
-                    try equalOptional(
-                        drafterText.numGlobalKeyValueHeads,
-                        target.numGlobalKeyValueHeads,
-                        field: "fullAttention.numGlobalKeyValueHeads")
-                }
+                // The gated `numGlobalKeyValueHeads` equality check was
+                // removed with the k_eq_v-gated head rule: full-layer KV
+                // geometry is now compared unconditionally right below.
                 try equal(
                     effectiveFullKVHeads(drafterText),
                     effectiveFullKVHeads(target),
@@ -611,10 +609,10 @@ enum Gemma4MTPCompatibilityValidator {
     }
 
     private static func effectiveFullKVHeads(_ config: Gemma4TextConfiguration) -> Int {
-        if config.attentionKeqV, let global = config.numGlobalKeyValueHeads {
-            return global
-        }
-        return config.numKeyValueHeads
+        // Mirrors Gemma4Attention.init: full layers honor
+        // num_global_key_value_heads whenever present, independent of
+        // attention_k_eq_v (k_eq_v only elides v_proj).
+        config.numGlobalKeyValueHeads ?? config.numKeyValueHeads
     }
 
     private static func equal<T: Equatable & CustomStringConvertible>(

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 }