fix: prefill the trunk, not the assistant, in dual-model MTP - #46
Merged
Conversation
A Gemma 4 assistant checkpoint is all KV-shared layers — it ships no K/V
projections of its own and can only run with sharedKV handed in from the
trunk's cache, which is what callMTP does. But the dual-model path installs
the assistant as the iterator's `model`, so MTPTokenIterator.prepare() ran
its prefill through the assistant's plain callAsFunction, standalone.
That only aborts when prepare() actually forwards a chunk, i.e. when the
prompt exceeds prefillStepSize (512). Shorter prompts return .tokens without
touching the model, which is why every test so far passed and a ~9k-token
prompt died with:
Fatal error: Layer 0 is a KV-shared layer but received no sharedKV
Delegate callAsFunction to mainModelRef when it is set. Prefill belongs to
the trunk regardless: it is the trunk's cache being primed and the trunk's
logits the iterator samples the first token from.
Verified on gemma-4-26b-a4b-it-4bit + gemma-4-26B-A4B-it-assistant-bf16 with
a cold server and a ~9k-token prompt: previously a hard abort, now 120 tok
@ 19.1 tok/s with coherent output; short prompts and streaming unchanged.
Also clamps the KVCacheSimple shared-KV slice to the key buffer, matching
the rotating branch — `offset` counts positions seen and can outrun it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
solderzzc
added a commit
to SharpAI/SwiftLM
that referenced
this pull request
Aug 9, 2026
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>
solderzzc
added a commit
to SharpAI/SwiftLM
that referenced
this pull request
Aug 9, 2026
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>
solderzzc
added a commit
to SharpAI/SwiftLM
that referenced
this pull request
Aug 10, 2026
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>
solderzzc
added a commit
to SharpAI/SwiftLM
that referenced
this pull request
Aug 10, 2026
…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>
solderzzc
added a commit
that referenced
this pull request
Aug 11, 2026
The MTP unit tests built their assistant from a config with num_kv_shared_layers: 0, which no shipped assistant checkpoint uses — real ones are entirely KV-shared and carry no K/V projections at all. So the existing tests exercised a model that cannot occur in practice, and dual-model MTP aborting on any chunk-prefilled prompt (#46) passed straight through them. Adds a matching config (every layer shared) and two tests: - a plain forward must reach the trunk, asserted by comparing logits against the trunk directly rather than merely surviving the call - MTPTokenIterator with prefillStepSize squeezed to 2, which forces the chunked prefill path with a 5-token prompt The second is the regression proper. prepare() only forwards a chunk when the prompt exceeds prefillStepSize; below it the model is never called, which is why a 512-token threshold hid the bug behind every short test prompt. Lowering the threshold reproduces it without a multi-gigabyte model pair — the smallest real one is 7.1 GB, more than GitHub's 10 GB per-repo cache budget can host alongside the existing modalities. Verified red-green: with the delegation in Gemma4AssistantModel.callAsFunction removed, both abort with the production error Layer 0 is a KV-shared layer but received no sharedKV Note that a regression here surfaces as a crashed test run rather than a failed assertion, since the guard is a fatalError. Full suite: 98 tests, 16 suites, all passing. Co-authored-by: Aegis AI Assistant <simba@aegis-ai.dev> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
solderzzc
added a commit
that referenced
this pull request
Aug 12, 2026
…#48) * test: cover the all-KV-shared assistant and its chunked prefill The MTP unit tests built their assistant from a config with num_kv_shared_layers: 0, which no shipped assistant checkpoint uses — real ones are entirely KV-shared and carry no K/V projections at all. So the existing tests exercised a model that cannot occur in practice, and dual-model MTP aborting on any chunk-prefilled prompt (#46) passed straight through them. Adds a matching config (every layer shared) and two tests: - a plain forward must reach the trunk, asserted by comparing logits against the trunk directly rather than merely surviving the call - MTPTokenIterator with prefillStepSize squeezed to 2, which forces the chunked prefill path with a 5-token prompt The second is the regression proper. prepare() only forwards a chunk when the prompt exceeds prefillStepSize; below it the model is never called, which is why a 512-token threshold hid the bug behind every short test prompt. Lowering the threshold reproduces it without a multi-gigabyte model pair — the smallest real one is 7.1 GB, more than GitHub's 10 GB per-repo cache budget can host alongside the existing modalities. Verified red-green: with the delegation in Gemma4AssistantModel.callAsFunction removed, both abort with the production error Layer 0 is a KV-shared layer but received no sharedKV Note that a regression here surfaces as a crashed test run rather than a failed assertion, since the guard is a fatalError. Full suite: 98 tests, 16 suites, all passing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * feat: load glm_moe_dsa / deepseek_v3_2 with dense attention (stage 1) GLM-5.2 uses `glm_moe_dsa`, which in mlx-lm is 54 lines whose model class is a bare `class Model(DSV32Model)` — it is DeepSeek V3.2, not its own architecture. V3.2 in turn is V3 plus DeepSeek Sparse Attention: a lightning indexer scores cached keys and attention is restricted to the top `index_topk`. The V3 half of that already exists in this repo, so both model types map to one new Swift model built on the V3 stack. The indexer is not implemented here. That is not an approximation below `index_topk`, because the reference returns no selection at all until the cache is longer than it: if k.shape[2] <= self.index_topk: return None GLM-5.2 ships index_topk 2048, so output is exact for the first 2048 positions and diverges past them — dense instead of top-2048 sparse, a long-context quality question rather than a failure. Implementing the indexer is stage 2 and also closes the DeepseekV4Compressor/Indexer TODO. Config decoding differs from V3 in one way worth noting: GLM-5.2 carries no top-level rope_theta, nesting it in `rope_parameters` (mlx-lm unpacks this in __post_init__). Both layouts are accepted. `rope_type: "default"` is deliberately not forwarded as rope_scaling — a non-nil dict there sends the V3 stack looking for yarn fields the config does not have. Also generalises a trap in DeepseekV3.sanitize, which dropped `model.layers.61` as a string literal. That number is just numHiddenLayers — V3 has 61 layers and its MTP block sits at index 61 — but GLM-5.2 has 78, so the literal deleted a real layer while keeping the actual MTP block. It now filters on `layerIdx >= numHiddenLayers`, identical for V3 and correct for any depth. Covered both ways: a 78-layer model keeps layer 61 and drops 78, a 61-layer model still drops 61. Verified red-green — under the old literal the GLM case fails on both counts. Tests are tiny random-init configs in the style of the Gemma 4 suite; they establish shape, flow and weight handling, not numerics against a reference. Nothing here is validated against real weights: the smallest glm_moe_dsa checkpoint is 308 GB. Full suite: 104 tests, 17 suites, passing. Refs SharpAI/SwiftLM#111 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Aegis AI Assistant <simba@aegis-ai.dev> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The bug
A Gemma 4 assistant checkpoint is entirely KV-shared layers — it ships no K/V projections and can only run with
sharedKVhanded in from the trunk's cache, which is whatcallMTPdoes.But the dual-model MTP path installs the assistant as the iterator's
model, soMTPTokenIterator.prepare()ran its prefill through the assistant's plaincallAsFunction, standalone:Why nothing caught it
prepare()only forwards a chunk when the prompt exceedsprefillStepSize(512). Below that it returns.tokenswithout touching the model at all — so every short-prompt test passed and the failure only appeared on a cold server given a long prompt. I found it running a ~9k-token benchmark against SharpAI/SwiftLM#137.The fix
Delegate
callAsFunctiontomainModelRefwhen set. Prefill belongs to the trunk in any case: it is the trunk's cache being primed and the trunk's logits the iterator samples the first token from. Generation and verification already route throughcallMTP, which runs the trunk internally — this was the one path that didn't.Also clamps the
KVCacheSimpleshared-KV slice to the key buffer, matching what the rotating branch already does (offsetcounts positions seen and can outrun the allocation).Verification
gemma-4-26b-a4b-it-4bit+gemma-4-26B-A4B-it-assistant-bf16, cold server, ~9k-token prompt:'Tokyo')Also worth recording, from the same runs: with the crash fixed, dual-model MTP is faster than no-MTP at long context (18.5 vs 17.4 tok/s, +6%) — the inverse of the ~14% loss measured at 150 tokens in ml-explore#137.
Blocks SharpAI/SwiftLM#137, which needs the submodule pointer bumped to this.
🤖 Generated with Claude Code