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,