Conversation
…ne-hot DFlash v1 draft positions are independent per-position marginals, so sampling each block position separately produced an incoherent block sequence the target rejected almost everything (0.45x, 1.16 tok/step measured on the Meta head at T=1). Propose argmax instead and rejection-sample against a one-hot proposal distribution, same as vLLM's DFlash scheme and standard greedy-draft rejection sampling. Target distribution is unchanged; only the proposal changes.
shoemoney
added a commit
to shoemoney/dflash
that referenced
this pull request
Sep 8, 2026
This branch has not been deployed
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.
Problem
In
_stream_generate, a v1 draft attemperature > 0samples every block position independently from its own distribution (_sampling_probsthen_sample_probs). A v1 head's positions are marginals, not a joint sequence distribution, so the sampled block is incoherent and the target rejects nearly all of it.Measured with
meta-models/Muse-Glimmer-30B-assistant(4-bit) onmlx-community/Muse-Glimmer-30B-4bit, M3 Ultra, T=1 / top-p 0.95 / top-k 64, 6 prompts x 256 tokens:Greedy acceptance for the same head is 2.93 at block 5, so the fix restores the head's full acceptance.
DFlash2DraftModelis unaffected either way (its candidate selector already proposes a coherent block; 3.35 accepted/step at the same settings).Fix
For non-DFlash2 drafts at
temperature > 0, proposeargmax(draft_logits)per position and pass a one-hotdraft_probsinto the existing_rejection_sample(draft_indices=None). Withqone-hot the accept testu * q < preduces tou < p(token)and the residualmax(p - q, 0)renormalized is the target distribution with the rejected token removed, i.e. the standard greedy-draft rejection scheme. The output distribution is exactly the target's; only the proposal changed. This is also what vLLM's DFlash proposer does.temperature == 0path and the DFlash 2 path are untouched. 10 insertions, 3 deletions.Note
This changes sampled-mode behavior for every v1 head on the MLX backend (Qwen3 / Qwen3.5 / Qwen3.6 / Gemma 4 drafts), not just Muse. I only measured Muse; if you have a Qwen3 v1 head bench handy it is worth a look, but the argument above is not model-specific.
Independent of the Muse-Glimmer adapter PR (#168); either can land first.