From cf7fe540ad615130971022a975ef6f6efe2750cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?apolin=C3=A1rio?= Date: Tue, 21 Jul 2026 23:40:23 +0200 Subject: [PATCH] speculative : sanitize non-finite DFlash target features 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). --- common/speculative.cpp | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index 365372005cdc..e6ec8f6ac62f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1080,6 +1080,29 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } } + // sanitize non-finite feature values before fusing. on Metal, the + // mat-mat kernels stage f32 activations as f16 for the simdgroup + // multiply; Laguna's massive-activation rows (attention-sink tokens, + // |x| ~ 1e6 in the pre-final-norm residual) overflow f16 -> inf/nan. + // one poisoned row would otherwise NaN the whole drafter KV cache. + { + size_t n_bad = 0; + for (auto & v : features_buf) { + if (!std::isfinite(v)) { + v = v != v ? 0.0f : (v > 0.0f ? 65504.0f : -65504.0f); + n_bad++; + } + } + if (n_bad > 0) { + static bool warned = false; + if (!warned) { + LOG_WRN("%s: sanitized %zu non-finite target feature values (f16 overflow on massive activations); " + "draft quality may degrade slightly on affected rows\n", __func__, n_bad); + warned = true; + } + } + } + // fuse extracted features through DFlash encoder llama_batch enc_batch = { /*.n_tokens =*/ n_chunk,