From bd2d8e10290e9aafc518bac8b8ea53905a0c0dc5 Mon Sep 17 00:00:00 2001 From: replikeit <56370676+replikeit@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:40:30 +0300 Subject: [PATCH] speculative : fix MTP warmup conditioning row 0 on a future hidden state (#2222) common_speculative_on_target_batch stored this batch's last hidden into target_hidden_by_seq before reading the map back for the shifted warmup conditioning, so row 0 was conditioned on this batch's last hidden (a future state) instead of the previous call's, and the position-0 zeros fallback was unreachable. Snapshot the previous value before the store; other readers are unaffected. Warmup-only; affects draft acceptance, not correctness. --- common/speculative.cpp | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index d08fab8e..53c75e64 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2872,6 +2872,15 @@ int32_t common_speculative_on_target_batch( return -1; } + // snapshot the previous call's last hidden before the store below overwrites it (row-0 conditioning needs it) + std::vector prev_call_last_hidden; + { + const auto prev_it = mtp_state->target_hidden_by_seq.find(seq_id); + if (prev_it != mtp_state->target_hidden_by_seq.end()) { + prev_call_last_hidden = prev_it->second; + } + } + const float * last_hidden = hidden_rows_storage.data() + (size_t) (batch.n_tokens - 1) * features.width; mtp_store_target_hidden(*mtp_state, seq_id, last_hidden, features.width); @@ -2899,9 +2908,8 @@ int32_t common_speculative_on_target_batch( const bool uses_shifted_hidden_rows = mtp_model_uses_recurrent_conditioning(*mtp_state); std::vector previous_hidden_storage; if (uses_shifted_hidden_rows) { - const auto hidden_it = mtp_state->target_hidden_by_seq.find(seq_id); - if (hidden_it != mtp_state->target_hidden_by_seq.end() && (int32_t) hidden_it->second.size() == features.width) { - previous_hidden_storage = hidden_it->second; + if ((int32_t) prev_call_last_hidden.size() == features.width) { + previous_hidden_storage = std::move(prev_call_last_hidden); } else { previous_hidden_storage.assign(features.width, 0.0f); }