[HamedShams]

Closes #261.

Two commits, reviewable apart: the first is the rollback fix from #261 (Sources/DFlash), the second makes tool calls work under --dflash (the bridge in Server.swift), which the degenerate output had been hiding.

--dflash on Qwen3.5 produced text that collapsed after the first drafted block on every prompt, at temperature 0 (Paris,\n,\n,\n,… for five capitals, Sunlight appears blue. for why the sky is blue), while the DFlash summaries reported 70-90 % acceptance. Block size 1 was byte-identical to the plain server, hence the target forward was fine; what was missing was the rollback after a partially accepted block.

Why. Qwen3.5's gated delta net layers keep their state in a plain MambaCache. The verify pass runs the whole block through them and when only part of it is accepted that state has to go back. Every DFlashTargetModel conformance says dflashIsHybridGDN = false, therefore, those layers never got a rollback-capable cache and restoreTargetCacheAfterAcceptance only moved their offset ("we can't undo them"). The rejected drafts stayed in the state and the next cycle verified against a prefix the model never produced. The hybrid path would not have helped as it stands: its tape is filled by recordTape, which nothing in mlx-swift-lm's layers calls.

What changes (first commit)

• armTargetRollback checkpoints plain MambaCache layers before the verify pass, with the checkpoint() the library added for speculative decoding.
• On a partial acceptance with such layers, the new replayCommittedTokens restores them (trim(1) restores the checkpoint), trims the attention layers to the pre-verify prefix, and runs the committed tokens, the staged first token plus the accepted drafts, through the target once more. This is what mlx-swift-lm's own draft-model loop does on a rejection with recurrent layers (Evaluate.swift, "re-feed the kept tokens"); its MTP path never needs it because Qwen3.5 advertises a native rewind depth of one draft.
• Full acceptance and the no-draft case are unchanged, apart from dropping the checkpoint. The next draft still conditions on the verify pass's hidden states, which are exact at the accepted positions.

Second commit: tool calls under DFlash. The bridge turned tokens into text with a bare detokenizer, so a tool call reached the client as content with finish_reason: "stop" and no tool_calls. Tokens now go through the TextToolTokenLoopHandler the standard path builds (the model's tool-call format, the request's tools, the policy, the configured stop strings); the handler is flushed before .info, and a stop it reports ends the stream and cancels the DFlash producer - as #224 requires. Reasoning stays with the server's own <think> tracker - as on the standard path.

Verified (M1 Max, mlx-community/Qwen3.5-4B-4bit + z-lab/Qwen3.5-4B-DFlash, the CI pair, greedy)

• The 8 prompts, 200 tokens, plain server vs --dflash: before the fix, every one degenerated within the first block; but after, they match to the last character, with 2 exceptions worth stating. DFlash emits the model's trailing newline before <|im_end|> where the plain path stops at the end token, so completion_tokens reads 2 higher. And on one prompt a close call between 2 plausible tokens went the other way; the same prompt also diverges between the plain path and DFlash at block size 1, which drafts nothing, so that is the multi-token kernel path, not the rollback. Thinking (reasoning_content) and a tool-call round trip now match the plain server.
• tests/test-dflash.sh: 12 passed, 0 failed, with 2 new blocks. 4b: a tool call under DFlash arrives as tool_calls with finish_reason: "tool_calls" (skips when the model does not call). 4c: the five-capitals answer still names the capitals with no repeated line; on main it is Paris,\n,\n,\n,….
• DFlashRecurrentRollbackTests (new, on the real MambaCache and KVCacheSimple): arming checkpoints the recurrent state, a partial acceptance restores it, trims attention to the prefix and replays exactly the committed ids, a full acceptance keeps the state and drops the checkpoint and only plain recurrent caches ask for a replay. DFlashStreamBridgeTests (rewritten, 7): tokens reach the handler the way the standard loop shows them, a stop token goes to onStopToken, a handler stop tears the producer down, and a real TextToolTokenLoopHandler over a byte tokenizer turns <tool_call>{…}</tool_call> into a .toolCall event.
• swift test --filter SwiftLMTests: 345 tests, 0 failures.

Cost, stated plainly. A partially accepted cycle now pays one extra target forward of the committed tokens. On my machine, with the 4-bit 4B model, that puts --dflash below plain decoding: about 66 tok/sec plain, 34 at --dflash-block-size 4, 30 at 8, 23 at the default 16. Before the fix the summaries showed 110 tok/sec, on garbage. A correct answer slower than the plain path seemed the right trade until the library can roll a GDN state back per position (the checkpointAfter machinery its MTP path uses), at which point the replay can go. I think the README's note that DFlash on Qwen3.6-35B-A3B "locks into low-entropy attractors" is this same corruption seen through a bigger model; though, I couldn't test that one here.