nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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.
This commit is contained in:
replikeit 2026-08-01 16:40:30 +03:00 committed by GitHub
parent f2bde5749b
commit bd2d8e1029
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
common/speculative.cpp
View File

@ -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<float> 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<float> 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);
}