From 34e545f5040bef90e5c46dc86b79d29301630cb1 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Fri, 2 Oct 2026 22:56:30 -0300 Subject: [PATCH] spec: MTP bridges the previous h-row into a fresh sequence (engine #290) common_speculative_impl_draft_mtp keeps pending_h, the target hidden row that precedes the first token of the next batch, and writes it into embl row i_batch_beg[seq_id] of every draft batch it processes. Row 0 of that batch has no other writer: process() fills rows 1..n_tokens-1 from the target's shifted nextn embeddings, so the bridge is what produces row 0. The bridge is unconditional. i_batch_beg[seq_id] is the first token of the sequence in *this* batch, so a fresh request at pos 0 also takes the bridge, and pending_h at that moment still holds the previous request's last hidden row - it is only initialised to zeros in the constructor. The MTP head is therefore seeded at pos 0 with another request's state, which is exactly what engine #290 reports: "draft candidate 0 at pos 0 has p = 0.862 vs 0.847", before any draft has been accepted. eagle3 has the same deferred-boundary concept and guards it: it stores pending_pos_last with the row and only uses it when "pending_pos_last + 1 == pos[beg]" (common/speculative.cpp). draft-mtp carries no position, so it cannot tell a continuation from a fresh sequence. Carry the position with the row: * pending_pos[seq_id] tracks the position pending_h belongs to (-1 = none). Set wherever pending_h is written: the end of process() and accept(). * The bridge is used only when pending_pos[seq_id] + 1 == pos of this batch's first row for that sequence; otherwise row 0 is zeroed, which is the state a brand new sequence starts from. * verify_pos_first[seq_id] records the position of verify_h[seq_id][0] so accept() can derive the position of the row it selects. This keeps every legitimate bridge: within a prefill split into ubatches (chunk 2 at pos 5 follows pending_pos 4) and from prefill into the first verify batch, while a new request at pos 0 no longer inherits the previous one. Also implement get_state/set_state for the pair, mirroring eagle3. Checkpoints already stash the speculative state (tools/server/server-context.cpp, "stash the draft's speculative state with the checkpoint", data_spec) and the field is even documented as "e.g. eagle3's deferred-boundary g_embd row"; mtp returned the base-class false, so a restored checkpoint left pending_h/pending_pos stale for the restored position. Verified: builds clean (Release, cpu-only llama-common, which compiles common/speculative.cpp and links). Not verified on hardware - the #290 repro needs Qwen3.8-27B with the MTP head on Strix Halo. --- common/speculative.cpp | 56 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 5653a90b889c..b060b250cca1 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1271,6 +1271,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // call to pair with, so it's stashed here until that next call fires. std::vector> pending_h; // [n_seq][n_embd] + // Position that pending_h belongs to, or -1 when there is no carryover. + // The bridge is only valid for a batch that continues at pending_pos + 1; + // a fresh sequence starts at pos 0 while pending_h still holds the previous + // request's last row (engine #290). + std::vector pending_pos; + std::vector i_batch_beg; std::vector i_batch_end; @@ -1278,6 +1284,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { // Row 0 corresponds to the sampled token, row N to the Nth accepted draft token. std::vector> verify_h; std::vector verify_h_rows; + std::vector verify_pos_first; // [n_seq] — pos of verify_h[seq][0] std::vector i_last; std::vector> chain_h; @@ -1352,6 +1359,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } pending_h.assign(n_seq, std::vector(n_embd, 0.0f)); + pending_pos.assign(n_seq, -1); i_last.assign(n_seq, -1); i_batch_beg.assign(n_seq, -1); @@ -1359,6 +1367,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { verify_h.assign(n_seq, {}); verify_h_rows.assign(n_seq, 0); + verify_pos_first.assign(n_seq, -1); } ~common_speculative_impl_draft_mtp() override { @@ -1461,7 +1470,19 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { continue; } - set_h(i_batch_beg[seq_id], pending_h[seq_id].data()); + // pending_h is the h-row that precedes the first token of this + // batch, so it only pairs with that token when the batch really + // continues from it. A fresh sequence starts at pos 0 while + // pending_h still holds the previous request's last row; seeding + // the MTP head with it there made the draft's first prediction + // depend on the previous request (engine #290). This is the same + // continuity guard eagle3 applies with pending_pos_last. + const llama_pos pos_beg = batch_in.pos[i_batch_beg[seq_id]]; + if (pending_pos[seq_id] >= 0 && pending_pos[seq_id] + 1 == pos_beg) { + set_h(i_batch_beg[seq_id], pending_h[seq_id].data()); + } else { + std::memset(batch.embd + (size_t) i_batch_beg[seq_id] * n_embd, 0, row_bytes); + } } auto * mem_dft = llama_get_memory(ctx_dft); @@ -1504,6 +1525,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1; verify_h_rows[seq_id] = n_rows; verify_h[seq_id].resize((size_t) n_rows * n_embd); + verify_pos_first[seq_id] = batch_in.pos[i_batch_beg[seq_id]]; for (int32_t i = 0; i < n_rows; ++i) { const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, i_batch_beg[seq_id] + i); @@ -1512,6 +1534,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes); + pending_pos[seq_id] = batch_in.pos[i_batch_end[seq_id]]; } return true; @@ -1681,6 +1704,37 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const int32_t i_h = std::min(n_accepted, n_rows - 1); const size_t row_bytes = (size_t) n_embd * sizeof(float); std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); + pending_pos[seq_id] = verify_pos_first[seq_id] + i_h; + } + + bool get_state(llama_seq_id seq_id, std::vector & data) const override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || pending_pos[seq_id] < 0) { + return false; + } + + const llama_pos pos = pending_pos[seq_id]; + const std::vector & h = pending_h[seq_id]; + + data.resize(sizeof(llama_pos) + h.size() * sizeof(float)); + std::memcpy(data.data(), &pos, sizeof(llama_pos)); + std::memcpy(data.data() + sizeof(llama_pos), h.data(), h.size() * sizeof(float)); + return true; + } + + void set_state(llama_seq_id seq_id, const std::vector & data) override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { + return; + } + if (data.size() != sizeof(llama_pos) + (size_t) n_embd * sizeof(float)) { + return; + } + + llama_pos pos = -1; + std::memcpy(&pos, data.data(), sizeof(llama_pos)); + + pending_pos[seq_id] = pos; + pending_h[seq_id].resize(n_embd); + std::memcpy(pending_h[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd * sizeof(float)); } bool need_embd() const override {