[solderzzc]

Proposed changes

Follow-ups from the review of #72 (the 12 lows). They are grouped by what changes.

Fixes

1. Hybrid draft finalize. SpeculativeTokenIterator.finalizeGeneration rewound the draft cache only with a per-layer trim. That can't take unemitted drafts out of a recurrent (Mamba) layer. The draft's recurrent state and draftY are now snapshotted at round start; finalize restores them and re-feeds the emitted tokens, as the main cache already does.
2. Small hybrid draft window. A hybrid draft rewinds up to numDraftTokens + 1 single-token writes. A wrapped RotatingKVCache can only give back maxSize - keep rows, so a small maxKVSize used to trap at the "Speculative draft rewind was partial" precondition mid-generation. Init now throws instead. The bound is tight: with numDraftTokens 3, maxKVSize 7 traps when the check is loosened by one, and 8 runs.
3. MTP round reset. MTPTokenIterator now resets committedPendingTokenCount and roundStart before each round. Before, a verify round that returned early (empty callMTP) left the previous round's count, and finalize then rewound tokens that had already been emitted.
4. FalconH1 num_logits_to_keep. Adds a TrailingLogitsLanguageModel protocol in MLXLMCommon, which FalconH1 conforms to. SpeculativeTokenIterator rejects a main model that returns fewer logit rows than a verify pass reads (numDraftTokens + 1). FalconH1 defaults num_logits_to_keep to 1.
5. CI. Removed the Gemma4StabilizationTests gate. No such suite has ever existed, so the gate selected 0 tests. Gemma4Tests still gates the PLE math.

Tests (MTPSpeculativeDecodingTests, Gemma4Tests, FalconH1Tests)

• 2.7i: MTP finalize after a partly accepted round and a fallback round. The last head is always wrong.
• 2.7j: MTP finalize drops two drafts at once, and the MTP caches are trimmed with the main cache. Uses a CountingDraftingModel that exposes MTP caches.
• 2.7k: non-hybrid MTP finalize (full_attention_interval 1).
• 2.7l: a round without output commits nothing.
• 2.7m: generateMTP stops on an EOS token mid-round; the loop's finalize leaves the caller's cache holding only the emitted tokens.
• 2.7g: the draft cache is compared with a reference; a second finalize must leave draft offsets and both processedTokenCounts unchanged; main lookahead ≥ 2 (accepted ≥ 2) must occur.
• 2.7n: the small hybrid draft window throws; one row more runs.
• Gemma4: the draft perturbation is seeded with withRandomState, and the test asserts draft offsets and the draft timeline.
• FalconH1: keeping 1 or 2 logit rows throws; keeping 3 or all runs.

Verification

• swift test --no-parallel: all new tests pass.
• Each fix was mutation-checked, and each mutation makes its test fail:
  • draft hybrid finalize off
  • window bound loosened by one (traps)
  • MTP reset removed
  • MTP-cache trim off, in finalize and in the round
  • loop finalize off
  • non-hybrid trims off, for MTP and for the draft
  • FalconH1 check off
  • draft count not cleared after finalize
• Full suite: 2 failures, both in AllowedToolOutputRouterTests (ordinaryTagLikeTextRemainsAResponse, nearProtocolMarkersSuppressWithoutStrippingOrdinaryTags). They fail the same way on main 9d5041a and are unrelated to this PR.
• Formatting: only the lines this PR adds were formatted with swift-format. The touched files aren't swift-format clean on main, so formatting whole files would change many unrelated lines.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

AI usage

• I have read this PR description in full and approve it as my own, and it
  accurately describes the code changes.
• AI usage disclosure: Claude Code (Claude Opus 5.5) wrote the code, the tests and this description, and ran the tests and the mutation checks. The repository owner approved opening this PR.

🤖 Generated with Claude Code

[solderzzc]

Pushed 06e5c07 addressing the review: doc comment kept on SpeculativeTokenIterator (error type moved above it); hybrid-draft window bound only enforced for numDraftTokens >= 2 (with a test that numDraftTokens == 1 and a one-row window does not throw); exact draft-position assertions after finalize in 2.7g and the Gemma4 test (a re-feed-only-start.y mutant now fails); and a deterministic partly accepted round test (2.7o). Full suite: only the 2 pre-existing AllowedToolOutputRouterTests failures.

🤖 Generated with Claude Code