speculative : save and restore MTP hidden state

[renewboy]

Overview

MTP does not implement get_state/set_state, so restoring target and draft KV checkpoints leaves the later batch's hidden state in the speculative driver.

Save and restore pending_h and the verification hidden rows through the existing hooks. Validate the model width, row count and payload size before changing state, and clear transient batch data after a successful restore.

Validation

• Added test-speculative-mtp-state.
• Without the patch at 08246a2: fails with MTP hidden state was not saved.
• With the patch: passes on Qwen3.5-2B (CPU and Metal) and Gemma4 E2B with its assistant model (Metal, shared KV).
• Checks draft token equality after restore, verification row restoration, malformed input rejection and isolation between sequences.
• Full CI has not been run.

Requirements

• I have read and agree with the contributing guidelines.
• AI usage disclosure: YES. Codex assisted with the upstream implementation and regression tests under human direction.

[pwilkin]

/bot review

[ggml-gh-bot]

❌ Code review failed.

Error: 429 "Too Many Requests"

[renewboy]

/bot review

@pwilkin Thanks! The bot review hit HTTP 429. Could you please retry /bot review and rerun the failed jobs?

[pwilkin]

/bot review

[ggml-gh-bot]

Automated code review

Full review below. Note: the previous bot review on this PR errored out with HTTP 429 and produced no findings, so this is a full pass over the diff rather than a delta review.

Scope: common/speculative.cpp (MTP speculative driver state hooks) + new test tests/test-speculative-mtp-state.cpp + CMake registration. Checklists run: scope gate, security, general. No public API, ggml, or server changes.

Scope gate: clean. This is a targeted bug fix (MTP had no get_state/set_state, so server checkpoint restore left stale hidden state), fits the existing hook design used by the EAGLE3 impl, and is single-purpose. AI usage is disclosed in the PR description.

Security / correctness of the new state hooks (common/speculative.cpp:1829-1887): the validation is done in the right order, which is the part most such patches get wrong:

• n_rows < 0 is checked before the (uint32_t) n_rows cast, so no sign-extension bypass.
• n_embd <= 0 is rejected before payload_bytes % row_bytes, so no division by zero.
• The exact payload check (payload_bytes / row_bytes != n_rows + 1) makes the allocation strictly proportional to the input size, so a crafted blob cannot force an oversized verify_h resize.
• State is left untouched when validation fails (early return before any mutation), which the test verifies via roundtrip.
• get_state/set_state guard the memcpy of verify_h behind n_rows > 0, avoiding the memcpy(dst, nullptr, 0) UB on empty vectors.
• The magic value prevents a blob saved by a different impl (e.g. EAGLE3's, which has no magic) from being misparsed here.

No blocking issues found. Findings below.

Will slow the review

(point 1) tests/test-speculative-mtp-state.cpp is registered with llama_build only (tests/CMakeLists.txt:153), so it compiles in CI but never runs there. It also needs a real MTP-capable model via -m, and the generate-models fixture dummies have no nextn layers, so they cannot drive it. As-is the regression coverage is manual-only. Either state in the PR how it is meant to be run (the exact command plus a model pointer), or see whether a minimal nextn-capable dummy can be generated so the test can be wired into ctest. Otherwise a future regression in this code will not be caught automatically.

(point 2) The magic 0x4d545031 is duplicated in get_state (common/speculative.cpp:1832) and set_state (common/speculative.cpp:1858). Hoist it to a single named constant (e.g. static constexpr uint32_t MTP_STATE_MAGIC) so the format tag cannot drift between writer and reader.

(point 3) The test's save/pollute/restore/re-draft sequencing is hard to follow: expected is captured before accept() dirties pending_h, restore_kv() is needed before each get_draft() comparison because draft() writes into the draft KV, and the final double set_state exists so accept(0) reproduces accepted. None of this is commented. Add a few one-line comments per phase explaining the invariant being checked, otherwise reviewers cannot tell a deliberate step from an accidental one.

(point 4) Memory cost per checkpoint: get_state stashes the full verify_h (up to n_batch * n_embd floats for a prefill-round checkpoint) into every server checkpoint entry, per sequence, times n_ctx_checkpoints. For decode rounds it is small (1 + n_accepted rows), but prefill checkpoints on wide models can add several MiB per stash. Worth confirming this is acceptable, or considering capping what is saved (e.g. only the rows that accept() can still select).

Nits

(point 5) set_state resets i_batch_beg/i_batch_end/i_last, clears chain_h and batch (common/speculative.cpp:1879-1884). As far as I can tell every consumer re-initializes these before use (process() refills i_batch_*, draft() clears batch and re-assigns chain_h per drafting seq, i_last is set before first use), so the resets look redundant. Also batch and chain_h are impl-wide while set_state is per-seq, so a per-seq restore wipes scratch for the other sequences. Either trim these lines or add one short comment stating a concrete case where stale values would be read.

(point 6) set_state silently drops malformed data. A rejected restore means the server continues with the current (wrong) pending_h, which shows up only as degraded draft quality. A single SPC_WRN on the rejection path (magic/width/size mismatch) would make this diagnosable.

(point 7) The bound (uint32_t) n_rows > llama_n_batch(params.ctx_tgt) is fine as a sanity check, but confirm n_batch is a true upper bound for verify_h_rows: process() sets rows from the per-seq token count of one batch passed to common_speculative_process, which llama_process caps at n_batch. If a caller ever passes a larger batch (e.g. a pre-decode batch assembled above n_batch), a legitimately saved state would be rejected on restore. The exact payload-size check alone already constrains memory, so the bound could also be dropped if it is not load-bearing.

Style: clean. ASCII only, comments concise, naming and formatting match the surrounding file; the test follows the existing common_params_parse / common_init pattern used by other tests.

This review was generated automatically by pi coding agent using zai-org/GLM-5.3. It may contain mistakes. Maintainers make the final call.

[ggerganov]

Could you rebase on latest master? Do you see improved performance after the change and if yes, in which use case?

@ruixiang63 I wonder, was there any reason to not implement get/set state for MTP? I see we've added it for eagle3 and not sure why we skipped MTP then.

[renewboy]

Could you rebase on latest master? Do you see improved performance after the change and if yes, in which use case?

@ruixiang63 I wonder, was there any reason to not implement get/set state for MTP? I see we've added it for eagle3 and not sure why we skipped MTP then.

@ggerganov Rebased onto master at 62a6f74.
I ran a small llama-server A/B using HumanEval, GSM8K, and MT-Bench prompts, with MTP enabled in both builds. I did not observe a consistent performance improvement.
The main benefit is checkpoint-restore correctness when revisiting or editing a cached prefix.

[ruixiang63]

I don’t think there was an intentional reason to skip MTP. My read is that EAGLE3’s deferred-boundary state made the need for get_state/set_state obvious, while MTP is covered by restoring the target/draft KV caches without needing deferred-boundary and state update.

The motiveation of adding state functions for eagle3 is here: #24593. @ggerganov

@renewboy Could you share some data showing the impact of this change? For example, speed-bench results comparing acceptance rate and speedup before and after the patch, ideally on a hybrid target model with a multi-turn workload that actually triggers checkpoint restores.

[ruixiang63]

I don't think we need a test for this. The test model doesn't trigger a hybrid model checkpoint restore, so the state restore path is never exercised.

[renewboy]

@ruixiang63 I ran the requested comparison on a hybrid target with actual server checkpoint restores.

Local environment

• Hardware: Apple M4 Pro, 48 GiB, Metal; macOS 15.6.1.
• Model: Qwen3.5-9B Q4_K_M with MTP.

Settings

• Upstream speed-bench request and metric code, using four selected two-turn SPEED-Bench samples: two coding and two math.
• Two repetitions per configuration.
• Greedy decoding, seed 42, output limit 256.
• Draft max/min: 3/0; draft probability threshold: 0.
• One slot, context 8192, batch/ubatch 256, eight threads.
• Four checkpoints, checkpoint minimum step 64.

Benchmark results

These results cover second-turn requests only: four requests per workload/configuration, all confirmed to restore a checkpoint. “Before” and “after” both enable MTP.

Workload	Acceptance before / after	Decode tok/s before / after	MTP/AR speedup before / after
Coding	56.06% / 56.06%	46.07 / 45.54	1.24× / 1.23×
Math	85.09% / 85.09%	60.38 / 58.77	1.60× / 1.56×

Patched/unpatched MTP inputs, histories, acceptance counts and response text matched. AR (Autoregressive) ratios are workload references, since one AR first-turn response differed.

No performance improvement was measured from the patch. The small timing differences do not establish a precise performance cost.

Restore-state and draft-logit diagnostic

I separately used the same instrumented binary in two modes, differing only in whether pending_h was restored. Both modes restored target/draft KV and verification state. This diagnostic was excluded from throughput measurements.

The first replayed token’s hidden input was compared against the checkpoint’s saved boundary vector:

Checkpoint position	Max absolute hidden-state error: restore omitted / enabled	Changed subsequent draft-logit vectors
34	11.075701 / 0	8/8
35	33.730278 / 0	8/8
59	19.777798 / 0	8/8
76	22.411842 / 0	8/8

All 32/32 raw draft-logit vector checksums differed, while 0/32 top-1 tokens changed. For example, one recorded top-1 candidate probability changed from 0.865745842 with restoration omitted to 0.868228912 with restoration enabled. Final target token IDs remained identical. Repeating the diagnostic reproduced the same boundary values and logit checksums.

MTP catch-up consumes pending_h for the first replayed token before updating it from the new target batch. Restoring it removes that stale boundary input.

Conclusion

The evidence supports a restore-state consistency issue affecting draft logits. It does not demonstrate a target-output correctness bug, acceptance improvement or server speedup.

The original motivation was checkpoint resumption in a direct llama-common integration. I also agree that the standalone test does not cover the server’s hybrid checkpoint-restore path.