fix(provider): v0.7.2 — engine_v2 text routing on VLM-loaded Gemma 4 slots (weight-shared extraction, per-request routing)#505
Conversation
Every prod Gemma 4 checkpoint ships a vision tower, so the provider loads it via VLMModelFactory and the slot is isVLM=true. The per-slot v2 gate (guard !isVLM else return nil) excluded 100% of Gemma traffic from engine v2 while GPT-OSS on v2 shows +78% p50 TPS. For an allowlisted VLM slot, the slot factory now extracts the CBv2-adapted MLXLLM Gemma4TextModel over the SAME weight arrays the MLXVLM wrapper holds (zero extra weight memory) and serves TEXT through engine v2; image/video requests keep the legacy VLM prepare/generate path (per-request routing, media peeled off before the bridge branch). Extraction re-applies the checkpoint's quantization structure the way loadWeights does (4bit-QAT and 8bit, incl. mixed per-layer precision), verifies with update(verify:[.all]) so any drift throws into engine_v2_fallback WARN + legacy, and runs a load-time forward parity gate (top-k containment + bounded max|Δlogit|). Telemetry logs vlm_text_routing=true.
Unit (live-isolated, scripted CBv2 stub): allowlisted VLM slot builds + registers a bridge; extraction failure -> legacy + engine_v2_fallback WARN; non-allowlisted VLM slot -> silent legacy; text request routes through the bridge; image-bearing request never reaches the bridge (legacy vision path). Weight-gated live (DARKBLOOM_LIVE_MLX_TESTS + _GEMMA) on the real prod qat-4bit and 8bit checkpoints: weight-share growth ~10KB (zero copy), load-time parity gate passes, v2 == legacy greedy token-exact on the same extracted model, and a real interleaved vision request does not perturb the v2 text output.
ProviderCore.version and coordinator LatestProviderVersion -> 0.7.2 for the engine_v2 VLM text-routing capability.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 1 finding(s)
- 🔵 [INFO]
provider-swift/Sources/ProviderCore/Inference/EngineV2VLMTextExtraction.swift:1-344— Large new file with complex extraction logic may introduce memory allocations without clear bounds- Suggestion: Review memory allocation patterns in extraction logic, especially in parameter tree operations and forward pass validation
Type_diligence — 1 finding(s) (1 blocking)
- 🟡 [MEDIUM]
provider-swift/Sources/ProviderCore/Inference/EngineV2VLMTextExtraction.swift:183— Missing type assertion check when casting JSONSerialization result to [String: Any]- Suggestion: Add proper error handling for the type cast: guard let parsed = try JSONSerialization.jsonObject(with: configData) as? [String: Any] else { throw EngineV2VLMTextExtractionError.invalidConfig("config.json is not an object") }
Additive_complexity — ✅ No issues found
2 finding(s) total, 1 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
| // Copyright © 2026 Eigen Labs. | ||
| // | ||
| // ContinuousBatchingV2 — weight-sharing text-model extraction for VLM slots. | ||
| // | ||
| // Every production Gemma 4 checkpoint ships a vision tower, so the provider | ||
| // loads it through `VLMModelFactory` and the resident module is MLXVLM's | ||
| // `Gemma4` wrapper — whose language model is a PRIVATE inline duplicate of | ||
| // the text architecture ("MLXVLM can't import MLXLLM") with none of the | ||
| // CBv2 hooks (`cbv2LayerKinds` / `newCacheV2`). The v2 engine therefore | ||
| // could never serve a VLM slot directly, and before v0.7.2 the per-slot | ||
| // `isVLM` gate excluded 100% of prod Gemma traffic from engine v2. | ||
| // | ||
| // This file builds the CBv2-adapted MLXLLM `Gemma4TextModel` OVER THE SAME | ||
| // WEIGHT ARRAYS the wrapper already holds: | ||
| // | ||
| // 1. decode the checkpoint's `config.json` `text_config` with MLXLLM's | ||
| // `Gemma4TextConfiguration` decoder and construct a lazy skeleton | ||
| // (nothing is materialized — MLXArray init is lazy until eval); | ||
| // 2. re-apply the checkpoint's quantization STRUCTURE to the skeleton the | ||
| // exact way `loadWeights` did for the wrapper (scales-presence gate + | ||
| // the same per-layer table from `BaseConfiguration`, whose keys live in | ||
| // the checkpoint's `language_model.`-prefixed key space); | ||
| // 3. re-key the wrapper's live parameter tree (`language_model.model.X` → | ||
| // `model.X`, `language_model.lm_head.X` → `lm_head.X`), run it through | ||
| // `Gemma4TextModel.sanitize` (drops shared-KV k/v duplicates the | ||
| // wrapper allocates but the MLXLLM model does not), and | ||
| // `update(parameters:verify:[.all])` — missing/extra/mis-shaped keys | ||
| // all THROW, which the engine factory catches as the standard | ||
| // `engine_v2_fallback` WARN + legacy path (never silent wrongness); | ||
| // 4. run a tiny forward through BOTH the wrapper's text path and the | ||
| // extracted model and require cross-containment of each side's greedy | ||
| // argmax in the other side's top-5, plus a bounded max |Δlogit| — a | ||
| // load-time backstop against catastrophic extraction bugs (see | ||
| // `assertForwardParity` for why bit-parity is structurally | ||
| // unattainable; env `DARKBLOOM_ENGINE_V2_VLM_PARITY_CHECK=0` skips). | ||
| // | ||
| // The result is a SEPARATE module instance (own norms/rope/layer objects) | ||
| // sharing only the immutable parameter arrays — zero extra weight memory, | ||
| // and no module-level mutable state shared with the wrapper, which is what | ||
| // makes the Qwen3.5-mrope class of cross-path state corruption structurally | ||
| // impossible here. Concurrency: the legacy vision forward (wrapper) and v2 | ||
| // text forward (extracted model) may interleave on the same arrays; both | ||
| // are read-only over the parameters (forward passes never mutate weights) | ||
| // and each path owns its private KV caches. | ||
|
|
||
| import Foundation | ||
| import MLX | ||
| import MLXLLM | ||
| import MLXLMCommon | ||
| import MLXNN | ||
| import MLXVLM | ||
|
|
||
| /// Failure modes of VLM text-model extraction. Every case lands in | ||
| /// `EngineV2Factory.makeBridgeIfSelected`'s catch → WARN `engine_v2_fallback` | ||
| /// telemetry + legacy serving. The messages are operator-facing (they ride | ||
| /// the telemetry `error` field), so they say exactly what to look at. | ||
| enum EngineV2VLMTextExtractionError: Error, CustomStringConvertible { | ||
| /// The loaded module is not a VLM wrapper this extraction understands. | ||
| case unsupportedWrapper(String) | ||
| /// The slot factory could not hand us the model directory (needed for | ||
| /// `config.json`). | ||
| case missingModelDirectory | ||
| /// `config.json` was unreadable or had no decodable `text_config`. | ||
| case invalidConfig(String) | ||
| /// The checkpoint has quantized weights but no `quantization` block in | ||
| /// `config.json` to derive the skeleton's quantization structure from. | ||
| case missingQuantizationConfig | ||
| /// The load-time forward parity gate failed: the extracted text model | ||
| /// disagrees with the wrapper's own text path on the same weights. | ||
| case parityMismatch(String) | ||
|
|
||
| var description: String { | ||
| switch self { | ||
| case .unsupportedWrapper(let type): | ||
| return "engine_v2 vlm extraction: unsupported VLM wrapper \(type)" | ||
| case .missingModelDirectory: | ||
| return "engine_v2 vlm extraction: model directory unavailable for config.json" | ||
| case .invalidConfig(let detail): | ||
| return "engine_v2 vlm extraction: config.json unusable (\(detail))" | ||
| case .missingQuantizationConfig: | ||
| return "engine_v2 vlm extraction: quantized weights but no quantization block in config.json" | ||
| case .parityMismatch(let detail): | ||
| return "engine_v2 vlm extraction: wrapper/extracted forward parity failed (\(detail))" | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// Weight-sharing extraction of the CBv2-adapted MLXLLM text model from a | ||
| /// loaded MLXVLM wrapper. Pure functions; no state. | ||
| enum EngineV2VLMTextExtraction { | ||
|
|
||
| /// Env kill switch for the load-time forward parity gate ("0"/"false"/ | ||
| /// "no"/"off" disables). Default ON — the check is one tiny prefill at | ||
| /// model-load time and is the backstop against silent architecture | ||
| /// drift between MLXVLM's inline text model and MLXLLM's. | ||
| static let parityCheckFlag = "DARKBLOOM_ENGINE_V2_VLM_PARITY_CHECK" | ||
|
|
||
| /// Checkpoint key-space prefix of the wrapper's language model. Both the | ||
| /// parameter tree and the per-layer quantization table use it. | ||
| private static let languageModelPrefix = "language_model." | ||
|
|
||
| /// Result of one extraction: the CBv2-adapted text model (sharing the | ||
| /// wrapper's weight arrays) plus the parity probe's max |Δlogit| for the | ||
| /// slot factory's log line (nil when the parity gate was disabled). | ||
| struct Extraction { | ||
| let model: Gemma4TextModel | ||
| let parityMaxAbsLogitDiff: Float? | ||
| } | ||
|
|
||
| /// Build an MLXLLM `Gemma4TextModel` over the weight arrays of a loaded | ||
| /// MLXVLM `Gemma4` wrapper. See the file header for the full mechanism. | ||
| /// | ||
| /// - Parameters: | ||
| /// - model: the slot's loaded module (must be `MLXVLM.Gemma4`). | ||
| /// - modelDirectory: the checkpoint directory (for `config.json`). | ||
| /// - environment: env snapshot (parity-gate kill switch). | ||
| static func extractTextModel( | ||
| from model: any LanguageModel, | ||
| modelDirectory: URL, | ||
| environment: [String: String] = ProcessInfo.processInfo.environment | ||
| ) throws -> Extraction { | ||
| guard let wrapper = model as? MLXVLM.Gemma4 else { | ||
| throw EngineV2VLMTextExtractionError.unsupportedWrapper( | ||
| String(describing: type(of: model))) | ||
| } | ||
|
|
||
| // 1. Text config: decode the checkpoint's `text_config` with the SAME | ||
| // decoder `LLMModelFactory` would use for a text-only checkpoint. | ||
| let configURL = modelDirectory.appendingPathComponent("config.json") | ||
| let configData: Data | ||
| do { | ||
| configData = try Data(contentsOf: configURL) | ||
| } catch { | ||
| throw EngineV2VLMTextExtractionError.invalidConfig( | ||
| "read \(configURL.path): \(error)") | ||
| } | ||
| let textConfig = try decodeTextConfiguration(configData: configData) | ||
|
|
||
| // Quantization table: `BaseConfiguration` holds the checkpoint-wide | ||
| // default plus the per-layer overrides, keyed in the checkpoint's | ||
| // `language_model.`-prefixed key space. | ||
| let baseConfig = try? JSONDecoder.json5().decode(BaseConfiguration.self, from: configData) | ||
|
|
||
| // 2-3. Lazy skeleton → quantization structure → weight-sharing update. | ||
| let skeleton = Gemma4TextModel(textConfig) | ||
| let textWeights = reKeyedTextWeights(wrapper: wrapper, sanitizer: skeleton) | ||
| try applyQuantizationStructure( | ||
| skeleton: skeleton, weights: textWeights, | ||
| perLayerQuantization: baseConfig?.perLayerQuantization) | ||
| // verify: [.all] — a missing model key, an unused weight key, or a | ||
| // shape mismatch all throw here. That is the design: any drift | ||
| // between the wrapper's parameter tree and the MLXLLM architecture | ||
| // must fail LOUDLY at load (→ engine_v2_fallback WARN + legacy), | ||
| // never produce a silently wrong serving model. | ||
| try skeleton.update( | ||
| parameters: ModuleParameters.unflattened(textWeights), verify: [.all]) | ||
|
|
||
| // 4. Load-time forward parity gate (env-gated, default on). | ||
| var parityDiff: Float? = nil | ||
| if parityCheckEnabled(environment: environment) { | ||
| parityDiff = try assertForwardParity( | ||
| wrapper: wrapper, extracted: skeleton, vocabSize: textConfig.vocabSize) | ||
| } | ||
|
|
||
| return Extraction(model: skeleton, parityMaxAbsLogitDiff: parityDiff) | ||
| } | ||
|
|
||
| // MARK: - Steps | ||
|
|
||
| /// Decode MLXLLM's `Gemma4TextConfiguration` from the VLM checkpoint's | ||
| /// `text_config` block. The top-level `quantization` block is merged in | ||
| /// so the configuration's informational `quantizationBits`/`GroupSize` | ||
| /// fields reflect the checkpoint (they do not drive the skeleton's | ||
| /// quantization — `applyQuantizationStructure` does). | ||
| static func decodeTextConfiguration(configData: Data) throws -> Gemma4TextConfiguration { | ||
| let root: [String: Any] | ||
| do { | ||
| guard | ||
| let parsed = try JSONSerialization.jsonObject(with: configData) | ||
| as? [String: Any] | ||
| else { | ||
| throw EngineV2VLMTextExtractionError.invalidConfig("config.json is not an object") | ||
| } | ||
| root = parsed | ||
| } catch let error as EngineV2VLMTextExtractionError { | ||
| throw error | ||
| } catch { | ||
| throw EngineV2VLMTextExtractionError.invalidConfig("parse config.json: \(error)") | ||
| } | ||
| guard var textConfigJSON = root["text_config"] as? [String: Any] else { | ||
| throw EngineV2VLMTextExtractionError.invalidConfig("no text_config object") | ||
| } | ||
| if textConfigJSON["quantization"] == nil, let quantization = root["quantization"] { | ||
| textConfigJSON["quantization"] = quantization | ||
| } | ||
| do { | ||
| let textConfigData = try JSONSerialization.data(withJSONObject: textConfigJSON) | ||
| return try JSONDecoder.json5().decode(Gemma4TextConfiguration.self, from: textConfigData) | ||
| } catch { | ||
| throw EngineV2VLMTextExtractionError.invalidConfig("decode text_config: \(error)") | ||
| } | ||
| } | ||
|
|
||
| /// Re-key the wrapper's live parameter tree into the text model's key | ||
| /// space and drop everything outside the language model (vision tower, | ||
| /// multimodal embedder). The result then passes through the text | ||
| /// model's own `sanitize`, which drops the k/v-projection duplicates | ||
| /// the wrapper allocates for KV-shared layers (MLXLLM does not allocate | ||
| /// those modules) and leaves everything else untouched. | ||
| private static func reKeyedTextWeights( | ||
| wrapper: MLXVLM.Gemma4, sanitizer: Gemma4TextModel | ||
| ) -> [String: MLXArray] { | ||
| var textWeights: [String: MLXArray] = [:] | ||
| for (key, value) in wrapper.parameters().flattened() { | ||
| guard key.hasPrefix(languageModelPrefix) else { continue } | ||
| textWeights[String(key.dropFirst(languageModelPrefix.count))] = value | ||
| } | ||
| return sanitizer.sanitize(weights: textWeights) | ||
| } | ||
|
|
||
| /// Mirror `loadWeights`' quantization pass onto the skeleton: a module is | ||
| /// quantized iff the (re-keyed, live) weights carry `<path>.scales`, with | ||
| /// (groupSize, bits, mode) resolved from the checkpoint's per-layer table | ||
| /// under the module's `language_model.`-prefixed checkpoint key — the | ||
| /// exact lookup the wrapper's own load performed, so the two module trees | ||
| /// can never disagree on quantization structure. All arrays involved stay | ||
| /// lazy; `update(parameters:)` replaces them before anything evaluates. | ||
| private static func applyQuantizationStructure( | ||
| skeleton: Gemma4TextModel, | ||
| weights: [String: MLXArray], | ||
| perLayerQuantization: BaseConfiguration.PerLayerQuantization? | ||
| ) throws { | ||
| let hasQuantizedWeights = weights.keys.contains { $0.hasSuffix(".scales") } | ||
| guard hasQuantizedWeights else { return } | ||
| guard let perLayerQuantization else { | ||
| throw EngineV2VLMTextExtractionError.missingQuantizationConfig | ||
| } | ||
| quantize(model: skeleton) { path, _ in | ||
| guard weights["\(path).scales"] != nil else { return nil } | ||
| return perLayerQuantization | ||
| .quantization(layer: languageModelPrefix + path)?.asTuple | ||
| } | ||
| } | ||
|
|
||
| // MARK: - Parity gate | ||
|
|
||
| static func parityCheckEnabled(environment: [String: String]) -> Bool { | ||
| guard | ||
| let raw = environment[parityCheckFlag]? | ||
| .trimmingCharacters(in: .whitespaces).lowercased(), !raw.isEmpty | ||
| else { return true } | ||
| return !["0", "false", "no", "off"].contains(raw) | ||
| } | ||
|
|
||
| /// Top-k window for the bidirectional argmax-containment check. | ||
| static let parityTopK = 5 | ||
| /// Ceiling on max |Δlogit| across the probe. Final logits are | ||
| /// softcapped to ±`final_logit_softcapping` (30 on every Gemma 4 | ||
| /// checkpoint), so unrelated distributions differ by up to ~60; | ||
| /// same-weights implementation noise measured ≤ ~8 (see below). | ||
| static let parityMaxAbsLogitDiff: Float = 20 | ||
|
|
||
| /// One tiny forward through the wrapper's text path and the extracted | ||
| /// model on identical tokens — a CATASTROPHIC-EXTRACTION detector, not | ||
| /// a bit-parity check. | ||
| /// | ||
| /// Token-exact parity between the two is structurally unattainable | ||
| /// (measured on gemma-4-26B-A4B qat-4bit, 2026-07): | ||
| /// | ||
| /// * the two implementations have different bf16 kernel/fusion | ||
| /// orderings, giving ~0.5 max |Δlogit| even at position 0 (where | ||
| /// RoPE is the identity), which flips near-tie top-8-of-128 MoE | ||
| /// expert selections at later positions (~8 max |Δlogit|); | ||
| /// * MLXVLM's wrapper mis-implements the checkpoint's declared | ||
| /// `rope_type: "proportional"` for full-attention layers as an | ||
| /// HF-style truncated-dims partial RoPE (freqs /128, partner +64), | ||
| /// while MLXLLM's `ProportionalRoPE` implements the declared type | ||
| /// (freqs /512 over the full head, partner +256 — verified against | ||
| /// the mlx_lm Python reference). The extracted model keeps the | ||
| /// CORRECT scheme; the wrapper discrepancy is flagged upstream. | ||
| /// | ||
| /// What a correct extraction guarantees — and what this gate enforces — | ||
| /// is that both models compute the same function up to implementation | ||
| /// noise: each side's greedy argmax must sit inside the other side's | ||
| /// top-`parityTopK` at EVERY position, and max |Δlogit| must stay under | ||
| /// `parityMaxAbsLogitDiff`. A mis-extracted model (wrong keys, wrong | ||
| /// quantization structure, wrong config) produces unrelated | ||
| /// distributions and fails both with overwhelming probability. The | ||
| /// probe is fixed and both models are deterministic, so for a given | ||
| /// (checkpoint, binary) the gate either always passes or always fails — | ||
| /// no per-load flakiness. | ||
| private static func assertForwardParity( | ||
| wrapper: MLXVLM.Gemma4, extracted: Gemma4TextModel, vocabSize: Int | ||
| ) throws -> Float { | ||
| // Fixed probe: small ids well inside every Gemma vocab; length stays | ||
| // far under the sliding window so the check exercises both layer | ||
| // types without materializing meaningful KV. | ||
| let probeTokens = [2, 651, 6134, 1024, 578, 108, 2364].map { | ||
| min($0, max(0, vocabSize - 1)) | ||
| } | ||
| let inputs = MLXArray(probeTokens.map(Int32.init)).expandedDimensions(axis: 0) | ||
|
|
||
| let wrapperLogits = wrapper(inputs, cache: nil).asType(.float32) | ||
| let extractedLogits = extracted(inputs, cache: nil).asType(.float32) | ||
|
|
||
| let k = parityTopK | ||
| // Top-k token ids per position, [1, L, k] (unordered within k). | ||
| let wrapperTopK = argPartition(-wrapperLogits, kth: k - 1, axis: -1)[.ellipsis, ..<k] | ||
| let extractedTopK = argPartition(-extractedLogits, kth: k - 1, axis: -1)[.ellipsis, ..<k] | ||
| let wrapperArgmax = argMax(wrapperLogits, axis: -1) | ||
| let extractedArgmax = argMax(extractedLogits, axis: -1) | ||
| let maxAbsDiff = MLX.abs(wrapperLogits - extractedLogits).max() | ||
| eval(wrapperTopK, extractedTopK, wrapperArgmax, extractedArgmax, maxAbsDiff) | ||
|
|
||
| let positions = probeTokens.count | ||
| let wrapperIds = wrapperArgmax.asArray(Int32.self) | ||
| let extractedIds = extractedArgmax.asArray(Int32.self) | ||
| let wrapperTop = wrapperTopK.asArray(Int32.self) // row-major [L × k] | ||
| let extractedTop = extractedTopK.asArray(Int32.self) | ||
| let diff = maxAbsDiff.item(Float.self) | ||
|
|
||
| for position in 0 ..< positions { | ||
| let window = position * k ..< (position + 1) * k | ||
| let wrapperWindow = Set(wrapperTop[window]) | ||
| let extractedWindow = Set(extractedTop[window]) | ||
| guard extractedWindow.contains(wrapperIds[position]), | ||
| wrapperWindow.contains(extractedIds[position]) | ||
| else { | ||
| throw EngineV2VLMTextExtractionError.parityMismatch( | ||
| "top-\(k) containment failed at position \(position): " | ||
| + "wrapper argmax \(wrapperIds[position]) vs extracted top-\(k) " | ||
| + "\(extractedWindow.sorted()); extracted argmax " | ||
| + "\(extractedIds[position]) vs wrapper top-\(k) " | ||
| + "\(wrapperWindow.sorted()); maxAbsLogitDiff=\(diff)") | ||
| } | ||
| } | ||
| guard diff <= parityMaxAbsLogitDiff else { | ||
| throw EngineV2VLMTextExtractionError.parityMismatch( | ||
| "max |Δlogit| \(diff) exceeds \(parityMaxAbsLogitDiff) " | ||
| + "(argmax containment held — distributions drifted)") | ||
| } | ||
| return diff | ||
| } | ||
| } |
There was a problem hiding this comment.
🔵 [INFO] ⚡ Large new file with complex extraction logic may introduce memory allocations without clear bounds
💡 Suggestion: Review memory allocation patterns in extraction logic, especially in parameter tree operations and forward pass validation
📊 Score: 2×3 = 6 · Category: unbounded_allocations
| as? [String: Any] | ||
| else { | ||
| throw EngineV2VLMTextExtractionError.invalidConfig("config.json is not an object") | ||
| } |
There was a problem hiding this comment.
🟡 [MEDIUM] 🏷️ Missing type assertion check when casting JSONSerialization result to [String: Any]
💡 Suggestion: Add proper error handling for the type cast: guard let parsed = try JSONSerialization.jsonObject(with: configData) as? [String: Any] else { throw EngineV2VLMTextExtractionError.invalidConfig("config.json is not an object") }
📊 Score: 4×3 = 12 · Category: missing_type_assertion
|
This PR bumps the provider version floor to 0.7.2 and adds VLM text-model extraction for CBv2 routing — the security-relevant changes are narrow and largely neutral, with one positive clarification and one new attack surface worth tracking. Trust boundaries touched
Per-threat assessmentT-007 / T-012 / T-027 — Weight hash and model integrity
This does not make the weight-hash situation worse than before (SEC-007 / fail-open enforcement is the pre-existing weakness), but it is worth flagging explicitly: the parity gate must not be documented or treated as a tamper-detection control — it is a development safety net, not a security boundary. ℹ️ Neutral against T-007/T-027. T-007 — VLM routing ordering contract ( The T-004 / T-024 —
New attack surface not covered by an existing threatOperator-controllable parity kill switch expanding the VLM extraction path
This is subsumed under SEC-007 / T-027 in spirit, but the new env var creates an additional explicit, documented bypass path that is not currently listed in the threat model. Recommend adding it to the T-027 detection hint or as a note under SEC-007. Recommendation: add an operator-facing warning in the binary startup log whenever Open findings resolved by this PRNone. SEC-007 remains open; this PR does not change coordinator-side weight-hash enforcement. 🔐 Threat model: |
|
Found 1 test failure on Blacksmith runners: Failure
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b6735f7cdf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| parityDiff = try assertForwardParity( | ||
| wrapper: wrapper, extracted: skeleton, vocabSize: textConfig.vocabSize) |
There was a problem hiding this comment.
When an allowlisted VLM slot loads with the default parity gate enabled, this assertForwardParity call runs wrapper and extracted-model forwards and MLX retains their transient activation buffers in GPU.cacheMemory; the load path then installs the slot and updates capacity, whose load/KV gates count cache memory as used. Without clearing MLX cache after this load-time probe, Gemma VLM loads can falsely advertise much less headroom or reject co-resident loads/requests until a later reclaimer happens to run.
Useful? React with 👍 / 👎.
| let slot = try await loadVLMSlot( | ||
| modelID: Self.qat4bitModelID, budgetBytes: 48 * 1024 * 1024 * 1024) | ||
| let scheduler = slot.scheduler | ||
| defer { Task { await scheduler.unloadModel() } } |
There was a problem hiding this comment.
When the gated Gemma live tests are run, this defer only launches an unstructured unload task and lets the test return immediately, so the next serialized live test can start loading the 8-bit checkpoint while the 4-bit model is still resident. On the memory-sized machines these tests target, that can create avoidable OOM/cap failures and flaky results; the teardown needs to be awaited before the test completes.
Useful? React with 👍 / 👎.
|
Both post-merge Codex findings triaged: (1) parity-probe MLX cache inflating capacity sizing — real, non-critical (single-digit % of the construction-time KV grant, self-healing, no correctness impact); fix (cache clear after probe, before sizing) + regression test ships in the v0.7.3 batch. (2) live-test teardown not awaited — test-only flakiness on weight-gated dev runs; awaited teardown ships in the same batch. Neither blocks the v0.7.2 rollout; monitoring adds a watch for unexpected co-resident load rejections on Gemma boxes as the field check for (1). |
2 participants
v0.7.2 — Engine V2 (continuous batching) for text on VLM-loaded Gemma 4 slots
Root cause
Engine V2 (ContinuousBatchingV2) never engaged for any production Gemma 4
model, so 100% of Gemma traffic — overwhelmingly text — ran on the legacy
engine, while GPT-OSS on v2 shows +78% p50 TPS in prod.
Why: every production Gemma 4 checkpoint ships a vision tower
(
gemma-4-26b-8bit,gemma-4-26b-qat-4bitboth havevision_configinconfig.json, soProviderLoop.modelIsVLM→ true). The provider loads themvia
VLMModelFactory, and the loaded module is MLXVLM'sGemma4wrapper.ProviderLoop+EngineV2.swiftthen gated v2 out per slot:The wrapper's language model is a private inline duplicate of the text
architecture ("MLXVLM can't import MLXLLM") with no CBv2 hooks
(
cbv2LayerKinds/newCacheV2), so the v2 engine genuinely could not serveit directly — and the blanket per-slot gate meant text requests (the vast
majority of Gemma load) silently missed the +78% win GPT-OSS already gets.
Fix — weight-sharing text-model extraction + per-request routing
For an allowlisted VLM slot, the slot factory now extracts the
CBv2-adapted MLXLLM
Gemma4TextModelover the same weight arrays thewrapper already holds (zero extra weight memory), and serves text through
v2. Image/video requests keep the legacy VLM prepare/generate path.
Extraction (
EngineV2VLMTextExtraction.swift):text_configwith MLXLLM'sGemma4TextConfigurationdecoder; construct a lazy
Gemma4TextModelskeleton.loadWeightsdoes —<path>.scales-presence gate + the per-layer(groupSize, bits, mode)table fromBaseConfiguration, keyed in thecheckpoint's
language_model.-prefixed key space. Handles 4bit (QAT) and8bit identically (quantization-agnostic).
language_model.model.X → model.X,language_model.lm_head.X → lm_head.X), run it throughGemma4TextModel.sanitize(drops the shared-KV k/v duplicates the wrapper allocates but MLXLLM does not),
then
update(parameters:verify:[.all])— a missing/extra/mis-shaped keythrows → the existing
engine_v2_fallbackWARN + legacy (never silentwrongness).
text path and the extracted model; each side's greedy argmax must sit in the
other's top-5 at every probe position, with bounded max |Δlogit|
(
DARKBLOOM_ENGINE_V2_VLM_PARITY_CHECK=0to skip).The extracted model is a separate module instance sharing only immutable
weight arrays — no module-level mutable state is shared with the wrapper, which
makes the Qwen3.5-mrope class of cross-path state corruption structurally
impossible. Legacy vision forward and v2 text forward may interleave on the
same arrays; both are read-only over the parameters and each owns its private
KV caches.
Per-request routing (
MultiModelBatchSchedulerEngine.streamChatCompletion):the media check (
VLMRequestInference.hasMedia) already peels image/videorequests off to the legacy vision path before the
engineV2Bridgebranch —so text-only requests on a bridge-carrying VLM slot route to v2, and
image-bearing requests stay legacy. Made the ordering an explicit contract.
Telemetry: the slot-build log line reports
vlm_text_routing=trueso the modeis visible in prod logs.
Before / After
flowchart TB subgraph Before A1[Gemma 4 request<br/>text or image] --> B1[VLM slot<br/>isVLM=true] B1 --> C1{makeEngineV2BridgeForSlot} C1 -->|guard !isVLM else return nil| D1[legacy BatchScheduler<br/>ALL Gemma traffic] D1 --> E1[legacy engine<br/>no +78% TPS] end subgraph After A2[Gemma 4 request] --> B2[allowlisted VLM slot] B2 --> C2{makeEngineV2BridgeForSlot<br/>isVLM + selection==v2} C2 --> X2[extract MLXLLM Gemma4TextModel<br/>over SAME weight arrays<br/>verify .all + parity gate] X2 -->|extraction fails| D2[engine_v2_fallback WARN → legacy] X2 -->|ok| F2[EngineV2 bridge<br/>vlm_text_routing=true] F2 --> G2{streamChatCompletion<br/>hasMedia?} G2 -->|text| H2[Engine V2<br/>+78% p50 TPS class] G2 -->|image/video| I2[legacy VLM prepare/generate] endProd evidence
is_vision=trueon both prod Gemma builds (config.jsoncarriesvision_config) → every one loads viaVLMModelFactory→ old per-slot gateexcluded them all.
was missing.
Engine-repo changes
None.
libs/mlx-swift-lmis untouched. The extraction codes purely againstexisting public API (
Gemma4TextModel,Gemma4TextConfiguration,BaseConfiguration,Module.parameters()/update(),quantize(),MLXVLM.Gemma4). One upstream discrepancy was found and documented (notfixed here): MLXVLM's inline Gemma4 text model implements the checkpoint's
declared full-attention
rope_type: "proportional"as an HF-styletruncated-dims partial RoPE (freqs /128, partner +64), while MLXLLM's
ProportionalRoPEimplements the declared type (freqs /512 over the full head,partner +256 — matches the mlx_lm Python reference). The extracted model keeps
the correct MLXLLM scheme; this (plus bf16 kernel-order noise flipping
near-tie top-8-of-128 MoE expert picks) is why token-exact wrapper parity is
structurally unattainable and the gate uses top-k containment instead. Flagged
for upstream resolution.
Tests
Unit (
EngineV2ProductionWiringTests, live-isolated, scripted CBv2 stub):engine_v2_fallback;vision path), scripted engine sees zero submissions.
Weight-gated live (
GemmaVLMEngineV2LiveTests,DARKBLOOM_LIVE_MLX_TESTS+DARKBLOOM_LIVE_MLX_GEMMA), on the exact prod checkpoints — both 4bit and 8bit:v2 == legacy-extracted == wrapper-legacy token-exact ("1, 2, 3, 4, 5"),
vision interleave hygiene holds (v2 output unchanged after an interleaved real
vision request);
token-exact.
Full suite:
swift build -j 8 && swift test -j 6→ 1333 tests / 114 suitespass (0 failures). Coordinator
go build ./...OK,gofmtclean.Version bumps
ProviderCore.version→0.7.2coordinator/api/server.goLatestProviderVersion→0.7.2Need help on this PR? Tag
/codesmithwith what you need. Autofix is disabled.