# MTP Compact Rollback The `--spec-mtp-cr-depth` option lets the user limit how many immediate MTP rollback states the target model keeps in VRAM. This reduces MTP's persistent memory use and leaves more room for the KV cache, allowing a larger context. Because rollback-state depth is no longer tied to `--spec-draft-n-max`, the user can also keep an aggressive MTP draft maximum instead of reducing it to recover context. For example, `--spec-mtp-cr-depth 1 --spec-draft-n-max 5` keeps one immediate rollback snapshot while still allowing MTP to propose as many as five tokens. The measured single-16-GiB-GPU Vulkan example was: | MTP mode | Maximum draft (`n`) | Available context | TG | |---|---:|---:|---:| | Standard Rollback | 3 | 56,000 | 44.34 t/s | | **MTP Compact Rollback** | **5** | **85,760** | **46.76 t/s** | This recorded Qwen 27B launcher result is a practical, directional example, not a final controlled A/B benchmark. It shows the intended trade: more context and a larger MTP draft without a throughput sacrifice. ## Overview MTP (multi-token prediction) speculative decoding proposes several likely next tokens and lets the target model verify them together. If several are accepted, generation advances by several tokens for one target verification pass. The maximum proposal length is controlled by `--spec-draft-n-max`; a larger value can improve token-generation speed when the model accepts the later positions often enough. Code is frequently predictable enough for MTP to work well. Across the recorded Qwen 27B and Qwen 35B-A3B code-generation benchmarks, aggregate draft acceptance ranged from about 80% to 98%, with many runs in the mid-80s to mid-90s. Acceptance depends on the prompt, model, sampler, and MTP settings, and later draft positions are normally accepted less often than the first. Aggregate acceptance and per-position acceptance should therefore be reported separately. On models with recurrent or hybrid recurrent layers, each accepted token also changes internal model state. Standard rollback retains one state snapshot for every possible MTP draft position, so raising the draft maximum also raises persistent VRAM use: ```text 1 live recurrent state + N rollback snapshots ``` In the tested dense 27B Q6 configuration, each additional state group consumed 149.625 MiB across the target devices. This amount is model-dependent: the tested 9B model used 50.25 MiB per group. It is not determined by MTP KV-cache quantization; KV precision changes a separate memory cost that also affects the available context. MTP Compact Rollback separates the maximum draft length from the number of immediate rollback snapshots. With compact depth `D` below draft maximum `N`, the target retains only `D` normal snapshots plus one reusable checkpoint. At depth 1, the rollback-state allocation stays fixed when the draft maximum is raised from 3 to 4 or 5. Other draft-related memory and computation can still vary, but the extra persistent rollback snapshots no longer consume context. ## What happens after a rejection MTP Compact Rollback does not discard the accepted result or generate the whole draft again. The tokens already chosen by verification and sampling are kept. Only the target model state needed to continue from those tokens may need to be reconstructed. There are two paths: ### Direct rollback If the rejected suffix fits within the retained compact depth, the target uses the existing direct rollback path. There is no checkpoint restore or replay. ### Checkpoint restore and replay Before verifying a draft that can exceed the retained depth, the server saves a `PARTIAL_ONLY` checkpoint of the target's recurrent state. The checkpoint stays on the target devices instead of being copied through host memory. If the target must roll back farther than the retained snapshots allow, the server: 1. restores the target state from before draft verification; 2. keeps the previously selected output token, accepted draft prefix, and replacement token; 3. evaluates those known tokens again through the target to reconstruct the correct recurrent state; and 4. continues generation from that reconstructed state. The draft model does not propose the tokens again, and sampling is not repeated. The extra target replay occurs only on deep rejections. Fully accepted drafts and shallow rollbacks remain on the existing faster paths. The implementation is integrated into both current Eaman patches. The standalone `eaman-mtp-compact-rollback` branch contains Compact Rollback plus opt-in adaptive MTP sizing; cumulative `eaman` adds the ROCm fitting, pipeline, HIP VEC, and MoE work. ## Why the checkpoint is on-device The first depth-limited prototype used a host-backed checkpoint. It saved GPU memory but copied a large recurrent state through host memory during speculative rounds, causing a substantial throughput loss. The current prototype uses: ```text LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE ``` The checkpoint buffers are allocated and pinned on their target devices during server startup, before evaluation begins. Each speculative round reuses those buffers, retaining most of the performance recovered over host checkpoints. If startup allocation fails, or a later save unexpectedly falls back, host storage remains selected for the rest of that server run. There is no separate checkpoint-device command-line option. On-device storage is selected internally, with sticky host fallback if device checkpoint storage cannot be used. The Compact Rollback control is `--spec-mtp-cr-depth`. ## Memory accounting With recurrent depth `D` below draft maximum `N`, the current on-device design uses: ```text 1 live state + D rollback snapshots + 1 replay checkpoint ``` At depth 1, this is three recurrent-state groups regardless of whether the MTP draft maximum is 3, 4, or 5. Compared with the original `1 + N` allocation, the net saving is `N - D - 1` state groups. For the tested dense Q6 model, one group is 149.625 MiB, split across both GPUs: | MTP draft maximum | Standard rollback states | Compact depth 1 + checkpoint | VRAM saved | |---:|---:|---:|---:| | 3 | 598.500 MiB | 448.875 MiB | 149.625 MiB | | 4 | 748.125 MiB | 448.875 MiB | 299.250 MiB | | 5 | 897.750 MiB | 448.875 MiB | 448.875 MiB | This is why larger MTP draft maxima no longer reduce fitted context through additional target recurrent snapshots. ## Automatic fitter support The device checkpoint is a separate backend allocation, so it is not present in the normal context memory breakdown. Without explicit accounting and early allocation, auto-fit could fill the GPUs and fail later when speculative generation first requested the checkpoint. Feature commit `be535675` introduced the first reservation. When fitting is enabled and the configured MTP recurrent depth is below the draft maximum, startup measures the target context twice: ```text n_rs_seq = D n_rs_seq = D + 1 ``` That per-device context-memory difference is useful as a recurrent-plane diagnostic, but commit `37c29f3c` showed that it is not an exact representation of the separately allocated checkpoint. Backend alignment and actual tensor placement can produce a different device split even when the total is similar. The recurrent memory implementation now builds the real one-row checkpoint copy layout and calls `ggml_backend_alloc_ctx_tensors_from_buft_size()` for each buffer type. It multiplies each result by the configured sequence count because every slot can own an independent checkpoint. The server reserves those exact sizes in the model's device order and logs the old plane delta only for comparison. Measurement, device-order, or mapping failures stop model loading instead of continuing with an unsafe fit. A zero or decreasing measured reservation is also treated as an error. For the F16 target-KV MoE profile that previously failed after prefill, the old proxy and corrected reservation were: ```text exact checkpoint recurrent-plane delta ROCm0 37.69 MiB 43.97 MiB ROCm1 25.12 MiB 18.84 MiB total 62.81 MiB 62.81 MiB ``` Commit `fff1fc7c` now allocates this exact layout for every configured slot before evaluation. In the final production-equivalent `--fit-target 70` run, the server retained 151,552 context, pinned 37.69/25.12 MiB, and reached the listening state. The earlier 140,288 result used `--fit-target 200`; it was not a target-70 before/after comparison. The remaining long-prompt confirmation is being combined with performance testing. The tested dense configuration reported: ```text ROCm0 reservation: 93.52 MiB ROCm1 reservation: 56.11 MiB Total reservation: 149.62 MiB Fitted context: 149504 tokens ``` The original validation allocated 93.516 MiB and 56.109 MiB on the first request. The current implementation performs the equivalent allocation at startup so runtime cannot consume that fitted headroom first. ## Measured performance On the dense Q6 model, the initial host-checkpoint depth-1 run achieved 23.96 t/s. Moving the checkpoint on-device improved the fixed-context result to 26.50 t/s. A later 14,046-token auto-fit run reached 27.09 t/s with 93.030% aggregate draft acceptance and 348 replayed tokens in 101 events. Preliminary on-device tests kept the same 149,504-token context for draft maxima 3, 4, and 5: | Draft maximum | TG | Mean accepted length | Replay events / tokens | |---:|---:|---:|---:| | 3 | 27.09 t/s | 3.63 | 101 / 348 | | 4 | 26.11 t/s | 4.05 | 159 / 528 | | 5 | 27.48 t/s | 4.97 | 191 / 772 | These workloads were not fully matched, so the table is directional rather than a final n=3/4/5 benchmark. It does show that the recurrent-memory cost remains fixed while draft length and replay work change. The n=5 run is currently the strongest preliminary throughput result, with about 60% acceptance at the fifth MTP position. ## Adaptive MTP sizing The standalone and cumulative patches also provide: ```text --spec-draft-adaptive ``` Adaptive sizing is opt-in and defaults to off. Each sequence starts at the configured `--spec-draft-n-max`. Two consecutive ceiling-reaching drafts that accept fewer than three tokens reduce the ceiling by one; a fully accepted draft raises it by one. Drafts stopped early by `--spec-draft-p-min`, a slot cap, or decode failure do not cause backoff. MTP CR accounting reports the original logical acceptance to the controller exactly once, while physical replay accounting still receives the reconstructed token count. The validated general settings are: ```text --spec-mtp-cr-depth 1 --spec-draft-n-max 7 --spec-draft-p-min 0.75 --spec-draft-adaptive ``` Do not force `--spec-draft-n-min 3`: matched testing showed that discarding one- and two-token drafts caused the earlier slowdown. Leaving `n-min` at its normal zero value allows those short drafts; the adaptive ceiling itself never falls below one. The broader build-1268 runtime matrix showed a workload-dependent result. Adaptive sizing cost about 2% on long code, but creative prose improved from 14.38 to 14.92 t/s at 68,864 context and from 10.49 to 14.20 t/s at 144,384 context. It is therefore included in both patches but remains default-off for workload-specific use. ## Usage Use the feature binary and set a compact rollback depth no larger than the MTP draft maximum: ```bash /home/eaman/llama/llama-mtp-cr-adaptive/build-adaptive-gpu/bin/llama-server \ --spec-type draft-mtp,ngram-mod \ --spec-draft-n-max 5 \ --spec-mtp-cr-depth 1 \ --fit-target 20 \ ... ``` The accepted range is: ```text 1 <= --spec-mtp-cr-depth <= --spec-draft-n-max ``` If `--spec-mtp-cr-depth` is omitted, the original full-depth behavior remains the default. The option does not reduce recurrent depth for EAGLE3, DFlash, or DSpark when those methods are combined with MTP. The environment-variable equivalent is: ```bash LLAMA_ARG_SPEC_MTP_CR_DEPTH=1 ``` Useful runtime telemetry includes: ```text draft acceptance acc per pos MTP replays = events / tokens ``` For comparisons, record fitted context, TG, per-position acceptance, mean accepted length, replay events/tokens, and the per-device checkpoint reservation. Restart the server between configurations so every run performs a fresh fit against comparable free VRAM. ## Hardening validation Feature commit `732dd501` makes the on-device checkpoint path safe when a recurrent checkpoint cannot be represented by one device range. A failed on-device save is now catchable and the server retries that checkpoint with host storage. The checkpoint records where it was actually stored, so restore uses the matching mode. A forced fragmented whole-cache fixture validates the fallback; normal positive per-slot recurrent allocation remains contiguous. Parser and fitter-input tests now cover default margins, fit disabled, split-mode none, asymmetric per-device targets, reduced and full recurrent depth, combined EAGLE3 use, and dynamic two-device argument order. The deterministic Vulkan 9B rollback matrix still passes every accepted-prefix position from 0 through 4. Its only later failure is the known, separate dirty full-checkpoint case. A real ROCm `-np 2` run on the Qwopus3.5 9B model exercised two concurrent slots. Each slot created an independent device checkpoint and lazily allocated 50.250 MiB on ROCm0. Both completed 192 generated tokens without OOM or host fallback: | Slot | Generation | Draft accepted | Deep replay | |---:|---:|---:|---:| | 0 | 35.99 t/s | 118/222 | 27 events / 90 tokens | | 1 | 28.73 t/s | 91/283 | 44 events / 123 tokens | The raw log contains the runtime section from the original server plus an incidental second-launch port-bind failure. It does not contain the original startup fitter lines, so the expected 100.50 MiB two-slot reservation is an inference from two measured 50.250 MiB allocations, not a directly captured reservation line. ## Current limitations and follow-up - Compact Rollback and adaptive sizing are integrated into the standalone and cumulative patches; both controls remain opt-in/default-preserving. - The on-device checkpoint behavior is currently hardcoded rather than exposed as a general state-storage policy. - Replay adds work, so the best compact depth and draft maximum depend on rejection distribution and workload. - The current evidence covers Vulkan/HIP correctness tests, a real two-slot 9B ROCm run, a 9B ROCm comparison, and dense dual-GPU runs. Repeated matched dense tests, MoE tests, and longer stability runs remain valuable. - Startup performs additional fitter measurement passes to calculate the checkpoint reservation. An audit found that safely reducing them needs a load-scoped cache for baseline and minimum-context scratch measurements; final fitting must remain fresh. This efficiency change is deferred and is not required for correctness. - Fitter tests validate reservation inputs and device mapping, including asymmetric placement and split-mode none. A future captured `-np 2` startup should also verify the reservation and preallocation lines directly. - A separate legacy test involving a full checkpoint restored over a dirty hybrid context still has a mismatch. The server path described here uses `PARTIAL_ONLY` checkpoints and is not that full-state restore path. - Long speculative output is not guaranteed to be byte-identical across all configurations because the existing MTP path and target-only generation can diverge over long runs. Correct rollback tests and stability measurements are more meaningful than long-output byte identity alone. ## Implementation status The current exact base is llama.cpp `3737e413`. | Branch / commit | Contents | |---|---| | `eaman-mtp-compact-rollback` / `6cb89357` | Standalone MTP Compact Rollback plus opt-in adaptive sizing | | `eaman` / `527be7e6` | Complete cumulative Eaman patch including the standalone feature | The matching standalone combined HIP/Vulkan binary reports build 1268 and is currently located under the historically named validation worktree: ```text /home/eaman/llama/llama-mtp-cr-adaptive/build-adaptive-gpu/bin/llama-server ``` The cumulative HIP/Vulkan server linked at build 1282 and its focused parser test passed. Both exported patches apply cleanly to `3737e413` and reproduce the tracked tree of their source branch exactly: ```text mtp_compact_rollback_3737e41.patch SHA-256: 759ad384f22859f084aa4a1c11252e63a423568f7e31cb9f623d4b473bd526ba latest_rocm_improvement_3737e41.patch SHA-256: 266efe7114dc74f0b947e98d5a355c2827a353317e9ea2cb6f6b746a4d0f4425 ``` Matched dense depth 3/1 fitting reached 136,704/149,504 context. Matched MoE depth 4/1 reached 233,216/262,144. Direct `-np 2` startup reservation, concurrent two-slot replay, and a 1,024-token depth-1 run passed. Detailed measurements and the current adaptive matrix are in `summary_eaman.md`.