speculative : fix DFlash 0% draft acceptance on Metal (sanitize non-finite target features) - #2
Merged
joerowell merged 1 commit intoJul 21, 2026
Conversation
On Metal, the mat-mat kernels stage f32 activations as f16 for the simdgroup multiply. Models with massive activations (Laguna's attention-sink rows reach |x| ~ 1e6 in the pre-final-norm residual) overflow f16 to inf whenever a ubatch is large enough to take the mat-mat path (e.g. prompt prefill), so the hidden-state features that DFlash extracts contain inf/nan for those rows. One poisoned row NaNs the drafter's KV cache from the first verify round onward, driving draft acceptance to 0% and making DFlash slower than base decoding. Base generation is unaffected because only the last row is sampled. Clamp non-finite feature values (nan -> 0, +-inf -> +-f16 max) before fusing them through the DFlash encoder. On an M4 Max with Laguna-S-2.1 Q4_K_M + laguna-s-2.1-DFlash-BF16 (greedy, --spec-draft-n-max 7): draft acceptance 0% -> 54%, decode 35 -> 49.5 tok/s (~1.4x over base).
joerowell
approved these changes
Jul 21, 2026
Collaborator
|
Thank you for your contribution! Don't know how we missed that. |
|
When I compile for Vulkan I get this error: Compile command: Is there something different I need to do? EDIT TO ADD: This issue isn't happening on the master branch now. So, please disregard. |
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
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.
DFlash speculative decoding on Metal gets 0% draft acceptance (all draft logits NaN), making it slower than base decoding. Sanity checking, the same setup works if the inference is done purely on CPU. The issue seems to be that the Metal mat-mat kernels stage f32 activations as f16, and Laguna's attention-sink activations (~±9e5 in the late residual) overflow to Inf in the hidden states DFlash extracts for the drafter. One poisoned row NaNs the drafter's KV cache from round one. Base generation is unaffected since only the last row gets sampled.
This PR fixes that by clamping non-finite values in the extracted features before the DFlash encoder (NaN→0, ±Inf→±f16 max). On an M4 Max with Laguna-S-2.1 Q4_K_M + the BF16 drafter (greedy, draft-n 7): acceptance 0%→54%, decode 35→49.5 tok/s (~1.4x base).
Repro (Apple Silicon, ≥96 GB):