Perf: cap Gemma4 MTP head cross-attention at long context (remaining piece of mlx-swift-lm#42) #237

Tracks what is still worth taking from SharpAI/mlx-swift-lm#42 (2026-05, 2 commits, conflicting with main). Issues are disabled on mlx-swift-lm, so this lives here.

What #42 did

• Compiler-warning cleanup and an output2D scatter refinement in the MTP path.
• perf(mtp): cap shared-KV cross-attention in runMTPHead to the last 16 backbone positions (maxSharedKV = 16), so the cost is O(16) instead of O(T). The commit reports removing a throughput regression at 40K-100K context (gemma-4-26b-a4b-it-8bit, M5 Pro).
• A MTPPartialRollback protocol on Gemma4AssistantModel (keeps lastBackboneHiddenStateAll), callMTPHeadOnly, and numMTPDraftTokens = 2, to roll back a partial draft rejection without re-running the main model.

State on main (SharpAI/mlx-swift-lm main 4d76a2e; ml-explore/mlx-swift-lm main 48114ee)

• None of maxSharedKV, MTPPartialRollback, callMTPHeadOnly, numMTPDraftTokens, lastBackboneHiddenStateAll exist on either branch. Upstream's MTP is the generic MTPDrafterModel / MTPSpeculativeTokenIterator (plus Qwen3.5 MTP); the Gemma4 assistant head (Gemma4AssistantModel) is fork-only.
• The speculative loop changed underneath ci: switch to SWIFTLM_PR_TOKEN to trigger regression suite on dependency bumps #42: main now uses upstream's staged rounds (beginRound/commit/rollback/rewindLastRound) and finalizeGeneration feeds the last draft. ci: switch to SWIFTLM_PR_TOKEN to trigger regression suite on dependency bumps #42's own rollback overlaps with that and cannot be applied as is.

Still worth doing

• Long-context cap on the Gemma4 MTP head's cross-attention. Re-measure first: run the MTP acceptance/throughput benchmark at 40K and 100K context with the cap off and on. If the regression still shows, add the window to Gemma4AssistantModel on current main.
• Partial-rejection rollback: only revisit if profiling shows the staged-round rewind is a cost for the Gemma4 assistant. Prefer extending the staged-round API over a second protocol.

Notes

• Do it on current main, not by rebasing the old branch. The branch also carries stray test_array_init.swift and test_scatter.swift.
• Keep MTPSpeculativeDecodingTests and Gemma4Tests green; their finalize expectations were adjusted in mlx-swift-lm#83.