From 362076c933c0a69584bf134bfa44d139404fa287 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sun, 9 Aug 2026 06:49:39 -0700 Subject: [PATCH 1/4] feat: make --mtp work for model families that ship MTP heads separately MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --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 --- Sources/SwiftLM/Server.swift | 85 +++++++++++++++++++++++++++++++++--- 1 file changed, 78 insertions(+), 7 deletions(-) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 2338fd0..626e4b8 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -286,6 +286,9 @@ struct MLXServer: AsyncParsableCommand { @Option(name: .long, help: "Number of MTP tokens to generate per speculation round (default: 3)") var numMtpTokens: Int = 3 + @Option(name: .long, help: "Assistant checkpoint providing MTP heads, for model families that ship them separately instead of in the main checkpoint (Gemma 4). Ignored when the main model carries its own MTP heads (Qwen3.5, DeepSeek V4).") + var mtpAssistantModel: String? + mutating func run() async throws { // Raise the open-file limit: large sharded models (e.g. Kimi K2.5, 182 safetensor // shards) + draft model + metallib + dylibs can exhaust the default macOS FD limit of 256. @@ -683,6 +686,47 @@ struct MLXServer: AsyncParsableCommand { draftModelRef = nil } + // ── Load the MTP assistant, for families that ship MTP heads separately ── + // Qwen3.5 and DeepSeek V4 carry their MTP heads inside the main checkpoint and + // conform to MTPLanguageModel directly, so generateMTP already works for them. + // 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. Loading it here and injecting that + // reference is what makes --mtp mean anything for Gemma 4; without it the flag + // is silently a no-op, because the main model fails the MTPLanguageModel test. + var mtpAssistantModelRef: (any DualModelMTP)? = nil + if self.mtp, let assistantPath = self.mtpAssistantModel { + print("[SwiftLM] Loading MTP assistant: \(assistantPath)") + var assistantConfig: ModelConfiguration + if FileManager.default.fileExists(atPath: assistantPath) { + assistantConfig = ModelConfiguration(directory: URL(filePath: assistantPath)) + } else if let local = ModelStorage.validatedContentDirectory(for: assistantPath) { + assistantConfig = ModelConfiguration(directory: local) + } else { + assistantConfig = ModelConfiguration(id: assistantPath) + } + if self.streamExperts { assistantConfig.lazyLoad = true } + let assistantDownloader = HubDownloader(hub: HubApi(downloadBase: cacheRoot)) + let assistantContainer = try await LLMModelFactory.shared.loadContainer( + from: assistantDownloader, + using: TransformersTokenizerLoader(), + configuration: assistantConfig + ) { _ in } + mtpAssistantModelRef = await assistantContainer.perform { assistantContext in + assistantContext.model as? (any DualModelMTP) + } + if mtpAssistantModelRef == nil { + print("[SwiftLM] ⚠️ \(assistantPath) does not provide MTP heads (not a DualModelMTP).") + print("[SwiftLM] Ignoring --mtp-assistant-model; generation will not use MTP.") + } else { + // The assistant drafts *for* this trunk, so it needs a reference to it. + await container.perform { mainContext in + mtpAssistantModelRef?.mainModelRef = mainContext.model + } + print("[SwiftLM] MTP assistant ready (\(self.numMtpTokens) tokens/round)") + } + } + // ── Load DFlash draft model for block-diffusion speculative decoding ── let dflashModel: DFlashDraftModel? let dflashBlockSizeConfig = self.dflashBlockSize @@ -810,7 +854,8 @@ struct MLXServer: AsyncParsableCommand { prefillSize: self.prefillSize, turboKV: self.turboKV, mtp: self.mtp, - numMtpTokens: self.numMtpTokens + numMtpTokens: self.numMtpTokens, + mtpAssistantModel: self.mtpAssistantModel ) let parallelSlots = self.parallel @@ -920,7 +965,8 @@ struct MLXServer: AsyncParsableCommand { request: request, bodyData: bodyData, config: config, container: container, semaphore: semaphore, stats: stats, promptCache: promptCache, draftModelRef: draftModelRef, numDraftTokens: numDraftTokensConfig, dflashModel: dflashModel, dflashBlockSize: dflashBlockSizeConfig, - dflashTargetModel: dflashTargetModel + dflashTargetModel: dflashTargetModel, + mtpAssistant: mtpAssistantModelRef ) } catch { let errMsg = String(describing: error).replacingOccurrences(of: "\"", with: "'") @@ -1091,6 +1137,7 @@ struct ServerConfig: Sendable { let turboKV: Bool let mtp: Bool let numMtpTokens: Int + let mtpAssistantModel: String? } // ── SSD Memory Budget ──────────────────────────────────────────────────────── @@ -1352,7 +1399,8 @@ func handleChatCompletion( numDraftTokens: Int = 4, dflashModel: DFlashDraftModel? = nil, dflashBlockSize: Int? = nil, - dflashTargetModel: (any DFlashTargetModel)? = nil + dflashTargetModel: (any DFlashTargetModel)? = nil, + mtpAssistant: (any DualModelMTP)? = nil ) async throws -> Response { let chatReq = try JSONDecoder().decode(ChatCompletionRequest.self, from: bodyData) let isStream = chatReq.stream ?? false @@ -1626,9 +1674,9 @@ func handleChatCompletion( } let remainingTokens = lmInput.text.tokens[startIndex...] let trimmedInput = LMInput(tokens: remainingTokens) - if config.mtp, context.model is any MTPLanguageModel { + if config.mtp, let mtpCtx = mtpContext(main: context, assistant: mtpAssistant) { stream = try MLXLMCommon.generateMTP( - input: trimmedInput, cache: cache, parameters: params, context: context, numMTPTokens: config.numMtpTokens + input: trimmedInput, cache: cache, parameters: params, context: mtpCtx, numMTPTokens: config.numMtpTokens ) } else { stream = try MLXLMCommon.generate( @@ -1637,9 +1685,9 @@ func handleChatCompletion( } } else { // Cache miss: process the full prompt. - if config.mtp, context.model is any MTPLanguageModel { + if config.mtp, let mtpCtx = mtpContext(main: context, assistant: mtpAssistant) { stream = try MLXLMCommon.generateMTP( - input: lmInput, cache: cache, parameters: params, context: context, numMTPTokens: config.numMtpTokens + input: lmInput, cache: cache, parameters: params, context: mtpCtx, numMTPTokens: config.numMtpTokens ) } else { stream = try MLXLMCommon.generate( @@ -2714,6 +2762,29 @@ func pendingStopPrefixLength(_ text: String, stopSequences: [String]) -> Int { return longest } +/// The context `generateMTP` should run against, and whether MTP applies at all. +/// +/// Two shapes exist. Qwen3.5 and DeepSeek V4 carry MTP heads inside the main checkpoint, +/// so the main context is already an `MTPLanguageModel` and is used as-is. Gemma 4 ships +/// the heads as a separate assistant checkpoint: there the *assistant* is the +/// `MTPLanguageModel`, so it is swapped into a derived context while the tokenizer, +/// processor and configuration — and the KV cache passed alongside — stay the trunk's. +/// That mirrors the reference usage in Gemma4MTPBench, which passes the assistant as the +/// model and the main model's cache. +/// +/// Returns nil when MTP does not apply, so callers fall through to plain generation. +func mtpContext(main: ModelContext, assistant: (any DualModelMTP)?) -> ModelContext? { + if let assistant, let assistantModel = assistant as? (any LanguageModel) { + return ModelContext( + configuration: main.configuration, + model: assistantModel, + processor: main.processor, + tokenizer: main.tokenizer + ) + } + return main.model is any MTPLanguageModel ? main : nil +} + /// Trims `text` at the earliest stop sequence it contains. /// /// Earliest *in the text*, not first in the caller's list: returning whichever entry From 04f7dcdb0109421750f6e164c71745754362d922 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sun, 9 Aug 2026 13:17:02 -0700 Subject: [PATCH 2/4] fix: bump mlx-swift-lm to pick up the dual-model MTP prefill fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- mlx-swift-lm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx-swift-lm b/mlx-swift-lm index b320bc4..47f2cc5 160000 --- a/mlx-swift-lm +++ b/mlx-swift-lm @@ -1 +1 @@ -Subproject commit b320bc485c68b009f5d6fc4eef717e34b658e6f4 +Subproject commit 47f2cc5876be28fd7afda947db4d207f77867295 From 869109f9fd442397d37628bc9b24aa8030e79949 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sun, 9 Aug 2026 13:27:47 -0700 Subject: [PATCH 3/4] test: exercise chunked prefill with a prompt over the 512-token boundary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- tests/test-contract.sh | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test-contract.sh b/tests/test-contract.sh index d31235d..8b8cc55 100755 --- a/tests/test-contract.sh +++ b/tests/test-contract.sh @@ -210,6 +210,37 @@ else skip "node not available on this machine" fi +# ── 8. A prompt long enough to be prefilled in chunks ──────────────────────── +# Below prefillStepSize (512 tokens) the generator's prepare() returns the prompt +# tokens without ever forwarding them, so a whole code path — chunked prefill — +# goes unexercised. Every other prompt in this repo's tests is under 80 characters, +# which is how a dual-model MTP crash on any real-sized prompt reached a merge +# queue with green CI (SharpAI/mlx-swift-lm#46). This closes the gap for the +# ordinary generate path only — CI runs no --mtp job, so the speculative variant +# of the same path stays uncovered until one exists (#128). +log "Test 8: prompt exceeding the prefill chunk size" +LONG_PROMPT=$(python3 -c ' +# ~3000 tokens: comfortably past 512 even with an efficient tokenizer. +print(("The quick brown fox jumps over the lazy dog near the river bank. " * 190) + "Reply with the single word: done.")') +LONG_BODY=$(python3 -c ' +import json, sys +print(json.dumps({"messages": [{"role": "user", "content": sys.argv[1]}], + "max_tokens": 16, "stream": False}))' "$LONG_PROMPT") +LONG_RESP=$(curl -sf --max-time 300 "$URL/v1/chat/completions" \ + -H 'Content-Type: application/json' -d "$LONG_BODY" 2>/dev/null || true) +if [ -z "$LONG_RESP" ]; then + # A crash in the prefill path kills the connection rather than returning an error + # body, so an empty response is the signal we are actually looking for here. + fail "no response to a chunk-prefilled prompt — server may have died; see /tmp/SwiftLM-test-contract.log" +elif echo "$LONG_RESP" | python3 -c ' +import json, sys +d = json.load(sys.stdin) +sys.exit(0 if d["choices"][0]["message"]["content"].strip() else 1)' 2>/dev/null; then + pass "chunk-prefilled prompt produced content" +else + fail "chunk-prefilled prompt returned no content: $(echo "$LONG_RESP" | head -c 120)" +fi + log "═══════════════════════════════════════" log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total" log "═══════════════════════════════════════" From 03784a2a6f0e27e60f70c4162defc798e64ea752 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sun, 9 Aug 2026 20:35:09 -0700 Subject: [PATCH 4/4] chore: re-point mlx-swift-lm at the merged prefill fix on main MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- mlx-swift-lm | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlx-swift-lm b/mlx-swift-lm index 47f2cc5..6a2c179 160000 --- a/mlx-swift-lm +++ b/mlx-swift-lm @@ -1 +1 @@ -Subproject commit 47f2cc5876be28fd7afda947db4d207f77867295 +Subproject commit 6a2c179998723107bb3d271963eeb9061056bee5