Skip to content

speculative : fix DFlash 0% draft acceptance on Metal (sanitize non-finite target features) - #2

Merged
joerowell merged 1 commit into
poolsideai:lagunafrom
apolinario:fix-dflash-metal-nonfinite-features
Jul 21, 2026
Merged

speculative : fix DFlash 0% draft acceptance on Metal (sanitize non-finite target features)#2
joerowell merged 1 commit into
poolsideai:lagunafrom
apolinario:fix-dflash-metal-nonfinite-features

Conversation

@apolinario

Copy link
Copy Markdown

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):

./build/bin/llama-cli -m laguna-s-2.1-Q4_K_M.gguf -md laguna-s-2.1-DFlash-BF16.gguf \
  --spec-type draft-dflash --spec-draft-n-max 7 -c 8192 -st -fa on --temp 0 --top-k 1 \
  -p "def fibonacci(n):" -n 96 --no-warmup -v 2>&1 | grep "draft acceptance"
# draft acceptance = 0.00000 (0 accepted / 637 generated)
# with -ngl 0 -ngld 0 (CPU): acceptance ~0.29
# with -ub 8 (prefill via mat-vec kernels, f32 reads): acceptance recovers

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

Copy link
Copy Markdown
Collaborator

Thank you for your contribution! Don't know how we missed that.

@joerowell
joerowell merged commit 04b2b72 into poolsideai:laguna Jul 21, 2026
1 check passed
@usofrob

usofrob commented Jul 22, 2026

Copy link
Copy Markdown

When I compile for Vulkan I get this error:
common/speculative.cpp:1091:35: error: ‘isfinite’ is not a member of ‘std’
1091 | if (!std::isfinite(v)) {

Compile command:
cmake -B build -DGGML_VULKAN=ON
cmake --build build --config Release -- -j 24

Is there something different I need to do?
It worked before I took in this change.
I'm compiling on Ubuntu 24.04.

EDIT TO ADD: This issue isn't happening on the master branch now. So, please disregard.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants