perf(mtp): Gemma4 MTP sliding-window fix + 8-bit benchmark results - #109
perf(mtp): Gemma4 MTP sliding-window fix + 8-bit benchmark results#109solderzzc wants to merge 1 commit into
Conversation
…rmup ## Changes ### mlx-swift-lm (submodule bump → c552b4d) - Cap Gemma4 MTP shared-KV cross-attention to last 16 backbone positions (O(T) → O(16)), eliminating throughput regression at 40K-100K context - MTPPartialRollback protocol: callMTPHeadOnly re-seeds MTP draft from cached backbone state without re-running the main model - numMTPDraftTokens=2 per pass (depth=4 empirically worse on Metal) ### Benchmark results (gemma-4-26b-a4b-it-8bit, M5 Pro 64GB) 8-bit is bandwidth-bound (2× heavier weights). KV reads amortize across the verification batch → MTP provides real throughput gains at 8-bit: 40K ctx: 38.8 tok/s MTP vs 32.4 vanilla (+20%) 100K ctx: 22.5 tok/s MTP vs 14.9 vanilla (+51%) 4-bit MoE is compute-bound (MoE FFN dominates); MTP neutral/overhead. TQ+MTP counterproductive at both precisions (TQ removes bandwidth bottleneck, making MTP's batch cost proportional again). ### scripts/profiling/mtp_bench.py - Add Metal shader warmup request before first timed run per config (fixes inflated TTFT on first 512-token measurement — 1.77s → ~0.3s) ### README.md - Split 4-bit / 8-bit benchmark tables with accurate current numbers - Add precision-specific guidance: --mtp for 8-bit at 40K+; --turbo-kv alone for max throughput; don't combine TQ+MTP ### Server.swift / CLICommandBuilder / GenerationConfig - --mtp / --num-mtp-tokens wired through CLI → GenerationConfig - Architectural note in MTP dispatch path explaining compute-bound behaviour on 4-bit MoE (documents why TQ+MTP underperforms TQ)
There was a problem hiding this comment.
Pull request overview
This PR wires Gemma4 MTP assistant-model speculative decoding into the SwiftLM server/CLI flow and adds benchmarking/documentation for precision-dependent long-context performance.
Changes:
- Adds Gemma4 MTP assistant auto-resolution/loading and MTP acceptance logging in server responses.
- Persists/exports MTP generation settings via
GenerationConfigand CLI command builder. - Adds a new Test 13 profiling script and updates benchmark guidance in README/run script.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
Sources/SwiftLM/Server.swift |
Loads Gemma4 MTP assistant models, passes them into generation, and logs MTP acceptance. |
Sources/MLXInferenceCore/GenerationConfig.swift |
Adds explicit decoding defaults for persisted MTP-related generation settings. |
Sources/MLXInferenceCore/CLICommandBuilder.swift |
Emits --mtp and --num-mtp-tokens in generated CLI commands. |
scripts/profiling/mtp_bench.py |
Adds HTTP-based MTP/TurboQuant benchmark runner. |
run_benchmark.sh |
Routes Test 13 to the new Python benchmark script. |
README.md |
Updates Gemma 4 benchmark tables and MTP/TurboQuant guidance. |
Comments suppressed due to low confidence (1)
Sources/MLXInferenceCore/CLICommandBuilder.swift:60
- This omits
--num-mtp-tokenswhen the UI config is1, but the SwiftLM server flag defaults to3(Server.swift:286-287). A copied command forenableMTP=true, numMTPTokens=1would therefore run with 3 draft tokens instead of the configured value.
if config.numMTPTokens != 1 {
parts.append("--num-mtp-tokens \(config.numMTPTokens)")
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| let asstConfig = ModelConfiguration(id: asstId) | ||
| let asstDownloader = HubDownloader(hub: HubApi(downloadBase: cacheRoot)) | ||
| let asstContainer = try await LLMModelFactory.shared.loadContainer( | ||
| from: asstDownloader, | ||
| using: TransformersTokenizerLoader(), | ||
| configuration: asstConfig | ||
| ) { _ in } |
| self.enableMTP = try container.decodeIfPresent(Bool.self, forKey: .enableMTP) ?? false | ||
| self.numMTPTokens = try container.decodeIfPresent(Int.self, forKey: .numMTPTokens) ?? 1 |
| if config.enableMTP { | ||
| parts.append("--mtp") | ||
| if config.numMTPTokens != 1 { | ||
| parts.append("--num-mtp-tokens \(config.numMTPTokens)") | ||
| } |
| if let asstId = asstModelId, resolveModelDirectory(modelId: asstId) != nil { | ||
| print("[SwiftLM] Loading Gemma4 MTP assistant model: \(asstId)") | ||
| let asstConfig = ModelConfiguration(id: asstId) | ||
| let asstDownloader = HubDownloader(hub: HubApi(downloadBase: cacheRoot)) | ||
| let asstContainer = try await LLMModelFactory.shared.loadContainer( | ||
| from: asstDownloader, | ||
| using: TransformersTokenizerLoader(), | ||
| configuration: asstConfig | ||
| ) { _ in } | ||
| mtpAsstRef = await asstContainer.extractDraftModel() | ||
| print("[SwiftLM] MTP assistant loaded (\(self.numMtpTokens) draft tokens/round). Using DualModelMTP speculative path.") | ||
| } else { | ||
| if let asstId = asstModelId { | ||
| print("[SwiftLM] ⚠️ Gemma4 MTP: assistant model '\(asstId)' not found in HF cache. Run: python -m mlx_lm.convert --hf-path \(asstId) to download it.") | ||
| } else { | ||
| print("[SwiftLM] ⚠️ --mtp: model '\(modelId)' is not a known Gemma4 MTP model. MTP requires a Gemma4 assistant checkpoint.") | ||
| } |
| **Key takeaways (4-bit):** | ||
| - 🚀 **TurboQuant is the headline win**: At 100K context, `Vanilla + TurboQuant` delivers **66.9 tok/s** vs **27.5 tok/s** Vanilla — a **2.43× speedup**. | ||
| - 💾 **Massive memory savings**: OS RAM at 40K context drops from **48.7 GB → 18.2 GB** with TurboQuant (63% reduction). | ||
| - ⚡ **MTP neutral on 4-bit MoE**: The 4-bit model is compute-bound (MoE expert dispatch). Batch verification scales linearly with token count, so MTP provides no net throughput gain over vanilla at 4-bit. |
|
Reviewed against current main. The textual conflict is small — one submodule pointer — but there is a semantic collision that needs a decision before any of this lands, so recording the analysis rather than merging. The decisionMain grew its own MTP path while this sat: This PR adds a competing branch ( So: native What is genuinely unlanded and worth having
Suggested split
Also: the branch adds |
|
Followed up on the "native How MTP is actually structured
Those checkpoints carry Gemma 4 has no in-model MTP heads. What that means for main today
Meanwhile the dual-model half is present but unreachable: nothing in So this PR is not a competing implementation. It is the missing wiring for a family the native path cannot cover. What I would change about itThe objections from my earlier comment stand, but they are placement and hygiene, not architecture:
Suggested order, unchanged otherwise
Happy to take any of these if useful. |
…ly (#137) * feat: make --mtp work for model families that ship MTP heads separately --mtp has been silently a no-op for Gemma 4. The gate is `context.model is any MTPLanguageModel`, and only Qwen35Model, Qwen35TextModel and DeepseekV4Model conform — those carry their MTP heads inside the main checkpoint. Gemma 4 does not: Google ships the heads as a separate assistant checkpoint, and Gemma4AssistantModel conforms to DualModelMTP (MTPLanguageModel plus a back-reference to the trunk it drafts for). Nothing in Sources/ ever set that reference except Gemma4MTPBench, which is not a target in Package.swift and so cannot build — leaving the whole path unreachable. --mtp-assistant-model loads the assistant, injects mainModelRef, and routes through the existing generateMTP call. Rather than adding a second generation branch, mtpContext() picks which context generateMTP should run against: the main context for in-checkpoint MTP, or a derived context whose model is the assistant while tokenizer, processor and configuration — and the KV cache passed alongside — stay the trunk's. That mirrors the reference usage in Gemma4MTPBench and keeps one code path, so the prompt cache is unaffected. An explicit flag rather than an id table: the table in #109 maps gemma-4-e4b-it to the E2B assistant and gemma-4-31b-it to the 26B one, which look like slips, and a wrong guess here silently drafts from the wrong model. Measured, and the result is not favourable yet. Output is correct — identical prefixes to baseline — but throughput is worse on both pairs available here: gemma-4-e2b-it-4bit + E2B assistant: 136.8 → 117.2 tok/s gemma-4-26b-a4b-4bit + 26B assistant: 74.1 → 63.6 tok/s and flat across --num-mtp-tokens 1/2/3 (63.3 / 64.1 / 63.6 on the 26B pair). Invariance to draft depth points at a fixed per-round cost rather than draft token cost, which is what the unlanded maxSharedKV=16 cap in #109 targets. Both assistants also ship bf16 against 4-bit trunks, so each drafted token costs more than the trunk token it replaces. So this makes the flag mean something and gives the perf work something to be measured against; it is not a speedup on its own. MTP stays opt-in and off by default, and with no --mtp-assistant-model the behaviour is byte-identical to before. 259 tests pass. Refs #109. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: bump mlx-swift-lm to pick up the dual-model MTP prefill fix Points at SharpAI/mlx-swift-lm#46, which makes the Gemma 4 assistant's callAsFunction delegate to the trunk. Without it this PR's feature aborts on any prompt over prefillStepSize (512 tokens) with Fatal error: Layer 0 is a KV-shared layer but received no sharedKV because MTPTokenIterator.prepare() prefills through context.model, which this PR makes the assistant — and an assistant checkpoint is entirely KV-shared layers that cannot run without sharedKV from the trunk. To be re-pointed at main once #46 lands, since the squash rewrites the SHA. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * test: exercise chunked prefill with a prompt over the 512-token boundary Every prompt in this repo's test suite is under 80 characters, so prepare() always returned prompt tokens without forwarding them and chunked prefill was never run. That gap is how a dual-model MTP crash on any real-sized prompt reached a green CI (SharpAI/mlx-swift-lm#46) — the failure needed only a prompt past prefillStepSize to appear, and nothing in CI supplied one. Adds one ~2700-token request to the contract suite. An empty response is treated as a failure, not an error case: a crash in prefill drops the connection rather than returning an error body, which is precisely the signature being watched for. This covers the ordinary generate path only. CI runs no --mtp job, so the speculative variant of the same code path remains uncovered (#128). Verified locally: server logs prompt=2697t for the new case, suite 10 passed 0 failed 2 skipped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * chore: re-point mlx-swift-lm at the merged prefill fix on main SharpAI/mlx-swift-lm#46 landed as squash commit 6a2c179, which replaces the branch SHA the previous bump pointed at. The tree is byte-identical to the interim pointer, so the CI already run against this PR still applies — only the commit identity changes, from a now-deleted branch to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
#109 bundled a sliding-window KV cap, a Gemma4-specific MTP-assistant auto-resolution path, and this benchmark tooling into one PR. The first two are not part of this extraction: - the KV cap was independently ported and evaluated earlier this session — no measurable benefit at 150 tokens or ~9k context, parked - the auto-resolution wiring (Gemma4MTPRegistry, mtpAsstRef reusing the draft-model path) duplicates what #137 already shipped more generally as the explicit --mtp-assistant-model flag; rebasing it forward would reintroduce a second, narrower implementation of a feature that already exists on main This PR is only the tooling, and it rebases clean because it is genuinely orthogonal: mtp_bench.py drives the server with --mtp --num-mtp-tokens N --turbo-kv, the single-checkpoint MTP path, not the Gemma4 dual-model wiring in conflict. Verified none of these three files reference Gemma4MTPRegistry, mtp-assistant-model, or anything else from the withheld part of #109. run_benchmark.sh's Test 13 previously shelled out to `swift run Gemma4MTPBench`, a product that no longer exists in Package.swift — that path was already broken on main before this PR. It now drives mtp_bench.py against the regular SwiftLM binary instead. README's benchmark numbers (Gemma4 26B, 4-bit and 8-bit) are carried over from the original PR's measurements, not reproduced in this extraction — a 40K/100K context benchmark run is multi-hour. The 8-bit table matches data already used as reference in mlx-swift-lm#46 and the #137 comment thread earlier this session, so it is not new to this repo's history, just newly landing in the README. Verified: run_benchmark.sh syntax checked, mtp_bench.py compiles and its --help output parses correctly; confirmed --mtp, --num-mtp-tokens and --turbo-kv all already exist on main independent of --mtp-assistant-model. Refs #109 Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
|
Closing as superseded. This PR bundled three things, and each has been resolved independently: MTP-assistant wiring — superseded by #137 (merged), which ships the same capability more generally: an explicit Sliding-window KV cap — independently ported and evaluated on this session's hardware. No measurable benefit at either 150 tokens (62.6 vs 63.6 tok/s) or ~9k context (18.5 vs 18.6 tok/s) on the 4-bit model available for testing. Parked rather than merged, since the numbers didn't support it in the regime I could reach. The 8-bit long-context claim in this PR's benchmark tables (2–4× at 40K–100K) may still hold — I couldn't test 8-bit at that context length — so this isn't a refutation, just an unreproduced-in-this-environment result. If anyone revisits it, the ported branch is Benchmark tooling — landed via #151. Thanks for the original investigation — the 8-bit MTP numbers here ( |
…e it exposed (#152) * feat: auto-detect VLM checkpoints on the CLI, and fix a false positive it exposed Extracted from fix/pr57-speculative-ci rather than a rebase (see #109's resolution and the discussion around it). Its ModelArchitectureProbe.swift and tests were already shipped verbatim; only the CLI wiring and its test were still missing. **The gap.** Server.swift already probed the architecture at load time but discarded the result for the CLI entry point: `let isVision = self.vision`. Confirmed the user-facing effect directly: `SwiftLM --model LiquidAI/LFM2.5-VL-450M-MLX-4bit` (no --vision) printed "Loading LLM" and the subsequent image request failed. InferenceEngine.swift — SwiftBuddy's own loader — already auto-detects unconditionally; the CLI now mirrors that, keeping --vision/--audio as explicit overrides: let isVision = self.vision || (!self.audio && architecture.supportsVision) Verified both directions: LFM2.5-VL-450M-MLX-4bit now loads as a VLM and answers an image request without --vision; a plain LLM (Qwen2.5-0.5B) still loads as LLM with no flags; --vision still forces VLM loading explicitly. **What wiring this up exposed.** The fixtures suite broke: `moe-nested` started failing with "Key vision_tower.patch_embedder... not found", because that fixture's `model_type: "gemma4"` — carried since #145, for the earlier, narrower reason that Gemma4Configuration decodes text_config only, giving the MoE-nesting test a real nested config to exercise — was now read by the CLI's newly-active auto-detection and routed to VLMModelFactory. That is not a fixture bug on its own. `MLXLLM/Gemma4.swift`'s Gemma4Configuration (the LLMModelFactory "gemma4" entry) has no vision_config field at all — a real, intentionally-supported text-only checkpoint shape. ModelArchitectureProbe.knownVisionModelTypes listed bare "gemma4" as vision-triggering regardless, which is a genuine false positive for any real text-only Gemma4 checkpoint using that shape, not just this fixture — auto- detection would have broken loading one in production. Removed "gemma4" from that list; the separate `vision_config != nil` check already distinguishes the two correctly, since only VLMModelFactory's Gemma4Configuration requires that field. Also swapped the fixture's decoy container from vision_config to audio_config — unrelated to the bug above, but vision_config as a decoy was already the wrong choice: a real text-only Gemma4 checkpoint would never carry that key, so the fixture is more honest this way regardless of the probe fix. audio_config is excluded from the expert-count walk by ModelProfiler.nonLanguageContainers the same way vision_config was, so the decoy still proves what it proved before. Regression test added: testVLM_AutoDetectsLFM25WithoutVisionFlag in tests/SwiftBuddyTests/VLMTests.swift, adapted from fix/pr57-speculative-ci to match main's already-refactored captureStartupOutput helper (main had the refactor; only this second test case was missing). Red-green verified: fails with "Unsupported model type: lfm2-vl" against the pre-fix isVision line, passes once restored. Verified together: fixtures 6/0, contract 10/0/2, both VLMTests cases, LFM2.5-VL-450M-MLX-4bit auto-detects, Qwen2.5-0.5B does not false-positive, explicit --vision still works. Refs #109 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * fix: don't auto-detect vision when speculative/MTP decoding is requested CI's speculative-decoding suite started crashing (Trace/BPT trap on the first real generation request) after auto-detection began routing mlx-community/Qwen3.5-*-4bit through VLMModelFactory instead of LLMModelFactory. That checkpoint genuinely ships a vision_config even when used purely as a text draft/main pair in this test, and qwen3_5 is registered in both factories — the same ambiguity as gemma4, just surfaced through a different flag combination. --draft-model/--dflash/--mtp only wire up BaseLanguageModel from LLMModelFactory, so auto-detection now backs off whenever speculative decoding is requested, matching the pre-auto-detect behavior for that path. --vision remains a valid explicit override. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Summary
Resolves throughput regression in Gemma4 MTP speculative decoding at long-context lengths, and documents precision-dependent MTP behaviour from empirical benchmarking.
mlx-swift-lm submodule (→ fix/compiler-warnings-mtp-optim @ c552b4d)
Benchmark Results (M5 Pro 64 GB, gemma-4-26b-a4b-it-8bit)
8-bit is bandwidth-bound → KV reads amortize across the batch → MTP provides real gains:
4-bit MoE is compute-bound; MTP neutral. TQ+MTP counterproductive at both precisions.
Other
User Guidance