perf(mlx-lm): Gemma 4 v0.8.2 — shared tower, layer-18 submission, coupled weighted-unsort+R1, packed multimodal - #102
Conversation
…n, coupled weighted-unsort+R1, packed multimodal Retained slices (each source-reviewed; retention per 2026-08-09 production matrix against A/B default brackets): - Phase 1 shared VLM text tower: MLXVLM.Gemma4 owns exactly one Gemma4TextModel exposed as textModel; sanitizer normalizes only language keys, quantization/quantization_config alias overlays enter before strict load; nested VLM decode defaults and unsupported-flag rejections preserved; LoRA decoder roots; 'all' bidirectional attention; maxKVSize-bounded direct caches. - Phase 2 layer-interval lazy prefill submission: post-layer asyncEval at the configured interval (default 18, DARKBLOOM_GEMMA4_PREFILL_CHUNK_EVAL), CBv2 scheduled L>1 only, ordinary/decode/MTP untouched; effective interval provenance + counted StepProfiler engagement, fail-closed benchmark check. - Phase 4+7 coupled expert optimizations: one contract (gemma4SupportsCoupledExpertOptimizations = exact production topology AND safe-R1 quantization) drives the fused weighted-unsort direct reduction in SwitchGLU and the safe expert-QMM request together; lock-free-armed WeightedExpertUnsortProbe provenance, decode/MTP/small shapes always on the legacy scatter+weighted-sum route. - Phase 5 packed multimodal prefill: explicit model+cache capability claims, row-local image splice with row-aligned span contexts threaded through rectangular CBv2 groups; vision rows keep per-block span intersection inside the existing q=128 query blocks; text rows unchanged; paged/custom providers fail closed; MTP preconditioned nil. Dropped per matrix and removed: expert gate/up packing (own decode regression), dense gate/up packing (never engaged), and standalone weighted/R1 profiles (prefill/decode regressions); GateUp/Dense focused suites are superseded by Gemma4ExpertOptimizationEligibilityTests and BenchCBv2ProvenanceTests.
|
@codex review |
![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: 53e8ccd4e1
ℹ️ 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".
…ower Four cutover regressions found in review of #102: 1. fp16 SDPA promotion (vmlx #52): the deleted inline text tower promoted Q/K/V to float32 around attention when activations are fp16 and cast back before o_proj; the canonical tower passed fp16 straight into the fused/composed SDPA shapes, allowing non-finite scores. Restored in Gemma4Attention's direct forward (fused and fallback paths); bf16 (production) is unchanged. 2. Rank-1 token acceptance: both public Gemma4TextModel entry points (callAsFunction, callCapturingPreNorm) now expand [N] to [1, N] before any dim(1) read, matching the deleted twin's cache-reuse contract (e.g. deprecated TokenIterator callers). 3. Full-layer global KV heads: Gemma4Attention now honors num_global_key_value_heads on full layers independent of attention_k_eq_v (k_eq_v only elides v_proj), and the model-level kvHeads map uses the same per-layer rule. 4. Quantization round-trip: explicit Gemma4TextConfiguration.encode(to:) re-emits every keyed property plus the nested quantization block, so decode-encode-decode no longer drops text_config-local quantization (previously skipped module quantization at strict load). New suite Tests/MLXLMTests/Gemma4SharedTowerContractTests.swift pins all four (bit-exact rank-1 parity, projection-shape head rule, forged fp16-overflow stability, idempotent round trips incl. root precedence). Full package suites: 594 swift-testing tests + 471 XCTest pass (CBv2CompiledDecodeTests.testSinksParity flaked once compiled-vs-eager at token 2 on a path this change cannot reach; passed on rerun and is unrelated).
|
Codex review addressed in
New suite @codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1d46abc2df
ℹ️ 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".
… normalize layer_types Deep-review follow-ups on the shared-tower fixes: - CBv2 layer-kind derivation and the three MTP validation sites carried the old k_eq_v-gated global-KV-head rule while the model constructor moved past it — KV storage would mis-size for k_eq_v=false + global- heads configs (prefill trap on shape precondition). gemma4LayerKinds now derives kvHeads = sliding ? numKeyValueHeads : (global ?? numKV) and loses the dead attentionKeqV parameter (wrapper + both test call sites updated); MTP effectiveFullKVHeads/bind checks use the same rule and the stale numGlobal-equality site is removed. - The fp16 overflow test was vacuous: RMSNorm divides projection scale out before attention, so forging q/k_proj reached nothing. Scores are now forged via q_norm/k_norm gains (|q·k| ≈ 600²×16 ≫ 65504); a non-promoting build fails it. - layer_types decode: short explicit lists pad with sliding_attention and empty lists fall back to all-sliding instead of trapping at model construction — restoring the deleted towers' out-of-range tolerance. - Rank-1 acceptance now also pinned through callCapturingPreNorm (MTP verify surface); fp16 promotion comment records the CBv2 scope boundary (CBv2 serves the bf16 contract; fp16 CBv2 is not admitted). - New pins: kinds↔model kvHeads agreement across k_eq_v, and the canonical full-layer ProportionalRoPE construction (over full head dim, /dims frequencies, +inf pass-through) — disclosing that the deleted tower's truncated-rope pattern was the anomaly corrected by the cutover, not a semantic we restore. - Old-rule expectations in Gemma4AssistantDraftModelTests replaced with matched/mismatched effective-KV-head coverage. Suites: contract 10/10; CBv2Model/MTP/Assistant/BaseConfiguration 72/72.
|
Deep-review round complete — fixes landed as
New findings fixed this wave: short/empty Suites: contract 10/10; full package 599 swift-testing + 471 XCTest, 0 failures, twice ( @codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f00c9bdcfa
ℹ️ 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".
…TP attention-products check validateAttentionProducts still derived fullKVHeads via the gated rule — provably unfailable as a rejection check (per-field caps bound the product far below the tensor limit), but the old rule alive in code and inconsistent with the neighboring divides check. Full layers now honor num_global_key_value_heads whenever present there too. Final LM-verifier sign-off item; the k_eq_v rule is now gone from every KV-head derivation in the tree.
|
Final review wave (#2) complete. Every prior finding is closed; the final verifier pass signed off each repo:
Perf-gate status unchanged: attribution rerun deferred to a cooled machine (see the prior comment — the box measured ~1.7× slow during the late-night window). The remaining open perf item is the nested PR's per-hit @codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
…arch-v0.8.2 # Conflicts: # Libraries/MLXLMCommon/ContinuousBatchingV2/AttentionV1.swift # Libraries/MLXLMCommon/ContinuousBatchingV2/EngineLoopV2.swift # Package.swift # Sources/BenchCBv2Core/BenchCBv2RealModel.swift
|
Final merge-readiness pass complete:
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 29ff95bc84
ℹ️ 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".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 072da0662a
ℹ️ 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".
|
Review follow-up complete at |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1ddb070f74
ℹ️ 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".
|
Final follow-up at |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b677d81f43
ℹ️ 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".
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
1 participant
Status
Ready for review at
c5c5338. The lower mlx #4 and mlx-swift #11 dependencies are merged. Currentmain(802398d) is integrated, GitHub reports this PR clean and mergeable, every review thread is resolved, and final Codex review found no major issues.Stack order: merged mlx #4 -> merged mlx-swift #11 -> this PR -> d-inference ml-explore#607.
What ships
Gemma4TextModel.SwitchGLUinitializers retainfuseGateUp, fused checkpoints registergate_up_proj, and the Gemma-specific optimized reduction stays on its supported separate-projection topology.attentionKeqV:layer-kind overload remains source-compatible, and fully bidirectional legacy prompts run as one whole-prompt forward instead of invalid chunked prefill.Before / After: behavior
flowchart LR subgraph Before A1[Gemma VLM load] --> B1[second reconstructed text tower] A2[CBv2 prefill] --> B2[one terminal submission] A3[fused MoE checkpoint] --> B3[public fused initializer unavailable] end subgraph After C1[Gemma VLM load] --> D1[one owned shared text tower] C2[scheduled contiguous CBv2 prefill] --> D2[layer-interval submission plus safe expert path] C3[fused MoE checkpoint] --> D3[gate_up_proj topology loads and forwards] endBefore / After: code
flowchart TD subgraph Before V1[MLXVLM.Gemma4] --> X1[extraction and re-key] S1[SwitchGLU] --> U1[separate gate_proj and up_proj only] P1[EngineV2 prefix cache] --> L1[lookup and donation] end subgraph After V2[MLXVLM.Gemma4.textModel] --> C2[CBv2 / MTP / LoRA / sizing] S2[SwitchGLU.callAndWeightedReduce] --> G2[shared eligibility gate] S2 --> F2[optional fused gate_up_proj] P2[cbv2LayerKindsAllowPrefixReuse] --> L2[lookup and donation gates] endVerification
c5c5338: 460 XCTest passed, 1 skipped; 781 Swift Testing tests in 110 suites passed.SwitchGLUinitializer overloads.Performance evidence
The clean attribution epoch measured the production CBv2 path at 2K TTFT 1715 -> 1390 ms, arrival TTFT 1.41-1.76x, arrival E2E +18-25%, and decode +1.3-2.3% on an M4 Max. A final same-binary contiguous bracket passed correctness and provenance, but host throughput drifted materially, so its late timing is retained only as diagnostic evidence.
The final scope change removes the direct-prefill regression by allowing weighted reduction only through
Gemma4TextModel.cbv2Prefill; it preserves the serving path that produced the TTFT signal.Release posture
The stock provider release remains contiguous KV. Paged support and explicit paged gates remain available, but paged is not the default release posture.
No release or production deployment is part of this PR.