Per-request MTP/DFlash selection + DFlash streaming - #23
Open
dusterbloom wants to merge 2 commits into
Open
Conversation
DFlash speculative decoding was bound at model-load and was absent from the
streaming code path, so a streamed request silently fell back to MTP even when
a DFlash drafter was loaded. There was also no way to choose the method per
request.
- Refactor the ~350-line DFlash draft-verify loop into a shared
`dflash_decode<S: DflashSink>` driver fed by two sinks: `DflashBufferedSink`
(non-streaming GenerationOutput, unchanged behavior) and `DflashStreamSink`
(per-round StreamingOutput chunks). The gate/EMA/thermal logic is written once.
- Add `Speculation { Auto, DFlash, Mtp, None }` to SamplingParams and honor it at
all four dispatch sites; streaming branches before prefill (mirrors the
non-streaming site) to avoid a double prefill. Both heads already coexist in a
drafter-loaded engine, so this is dispatch-only — no extra weights.
- Parse a per-request `speculation` field on the OpenAI chat + Anthropic routes
(400 on invalid). Default `auto` = DFlash when a drafter is loaded (incl.
streaming), else MTP.
- Verify: dflash_streaming_matches_nonstreaming loads the real 9B + drafter and
asserts streaming == non-streaming DFlash byte-for-byte (deterministic gate-off
run), plus MTP routing via last_dflash_accepts. Speculation::parse unit test.
- Docs: README + docs/configuration.md document the field and draft_model.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
`trim_drafter_cache` and `rollback_kv_only` (both pub fn in dflash.rs) had zero callers anywhere in the repo — superseded by `crop_drafter_cache` (the one the draft-verify loop uses) and the tape-replay rollback path. Remove them. Co-Authored-By: Claude Opus 4.8 (1M context) <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.
What
Builds on the DFlash engine in panbanda#204 (base branch
feat/dflash-modal). Two fixes:DFlash now works while streaming. The streaming path had no DFlash branch, so a streamed request — what interactive clients send — silently fell back to MTP even with a drafter loaded. The ~350-line draft-verify loop is refactored into one shared
dflash_decode<S: DflashSink>driver feeding two sinks:DflashBufferedSink(non-streamingGenerationOutput, behavior unchanged) andDflashStreamSink(per-roundStreamingOutputchunks). The gate/EMA/thermal logic is written once; streaming dispatches before prefill so DFlash runs its own tap-prefill with no double prefill.Per-request method selection. New
Speculation { Auto, DFlash, Mtp, None }onSamplingParams, exposed as a"speculation"request field on the OpenAI chat + Anthropic routes (400 on invalid value). Defaultauto= DFlash when a drafter is loaded (incl. streaming), else MTP;mtpforces MTP even with a drafter;none= plain AR. Both heads already coexist in a drafter-loaded engine, so this is dispatch-only — no extra weights.Verification (Apple Silicon + MLX)
Asserts streaming DFlash == non-streaming DFlash byte-for-byte (deterministic gate-off run), and that
speculation=mtproutes to MTP even with a drafter loaded (last_dflash_acceptsstays empty). Plus aSpeculation::parseunit test. clippy (nursery) + fmt clean.Base / stacking
Targets
feat/dflash-modalso the diff here is just this change (~+570), not the whole engine. Retarget tomainonce panbanda#204 merges.🤖 Generated with Claude Code