# MTP Compact Rollback MTP CR purpouse is to reduce vRAM utilization of MTP so that users can have longer contex while enjoing the benefits of faster token generatioan. When using small models like QWEN 27B or 35B A3B for code generation in constrained vRAM configuration like 16GB GPU this is highly valuable, allowing long ctx sessions for agentic workflow, mitigating the speed penality of dense models. The `--spec-mtp-rs-depth` option lets the user limit how many immediate MTP rollback states the 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 is not constrained to compromise on a lower MTP to preserve a decent ctx lenght, max MTP draft depth like 5 can be used for the same ctx cost. For example, `--spec-mtp-rs-depth 1 --spec-draft-n-max 5` keeps only one immediate rollback snapshot while still allowing MTP to generate five tokens, this works well with code generation that has an hi acceptance rate. With QWEN 27B on a 16GB GPU typical setup: | 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** | ToDo: this table looks surprising good, rechck values with vanilla / patch comparison. The users here gains some 30k (50% increase) contex for his session while the increased MTP depth allows to compensate TG speed for those cases when MTP has to generate further rollbacks. ## 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 currently an opt-in prototype on the isolated `eaman-rs-snapshot-depth` source branch. Production branch `eaman` remains unchanged. ## 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 `--spec-mtp-rs-checkpoint-device` command-line option. On-device storage is hardcoded for this experimental recurrent fallback path. The only new user-facing control is `--spec-mtp-rs-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. ## Usage Use the feature binary and set a compact rollback depth no larger than the MTP draft maximum: ```bash /home/eaman/llama/llama-rs-snapshot-depth/build-rs-gpu/bin/llama-server \ --spec-type draft-mtp,ngram-mod \ --spec-draft-n-max 5 \ --spec-mtp-rs-depth 1 \ --fit-target 20 \ ... ``` The accepted range is: ```text 1 <= --spec-mtp-rs-depth <= --spec-draft-n-max ``` If `--spec-mtp-rs-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_RS_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 - The feature is experimental, opt-in, and isolated from production `eaman`. - 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 feature branch contains the following main commits, rebased onto production `a4e47992`: | Commit | Change | |---|---| | `12938d2e` | Configurable recurrent depth, deep-rejection replay, telemetry, and tests | | `237d104a` | On-device recurrent replay checkpoint proof of concept | | `b4764e61` | Per-device automatic fitter reservation for the lazy checkpoint | | `2eb3b94f` | Fragmented device-save host fallback and additional fitter/parser tests | | `37c29f3c` | Exact aligned checkpoint reservation and catchable allocation failure | | `fff1fc7c` | Startup checkpoint preallocation, sticky host fallback, and optional HIP VEC forcing | CPU/default and combined HIP/Vulkan builds pass at build 1249. Production source repository `/home/eaman/llama/llama.cpp` remains on branch `eaman`; experimental work remains in `/home/eaman/llama/llama-rs-snapshot-depth`. Current `8144f31` artifacts are `latest_rs_cumulative_8144f31.patch` (SHA-256 `339317ad294c33ba6e9c28f2a64cc839c69d899140202586c41de7abd8554209`) and `mtp_rs_only_8144f31.patch` (SHA-256 `902eaa2424444fc3afc414b8ca68f83c0dbf5bc4681fb8f95c06ffd466feea59`). Both cleanly apply to exact base `8144f319`; the cumulative tree exactly matches tip `fff1fc7c`. 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 limitations are in `summary_eaman.md`. Historical `a3b1eff` exports remain pinned test artifacts.