Skip to content

Opt-in codec for recurrent-state context checkpoints (2x less host RAM, off by default) #27211

Description

@hojin12312

Problem

For recurrent and hybrid models the server keeps up to --ctx-checkpoints (default 32) context checkpoints per slot, and each one serializes the full recurrent state as fp32 into host RAM (server_prompt::checkpoints). On Qwen3.8-27B (arch qwen35, 65 blocks, SSM state size 128 across 16 groups, inner size 6144) one checkpoint's recurrent plane is 149.6 MiB, so a single slot at the default limit holds about 4.7 GiB, and that scales linearly with slots.

I could not find any upstream attempt to compress this. The checkpoint-related work I found is all functional — SWA generalization (#16382), speculative checkpointing (#19493, #22227), slot persistence (#20819), restore fixes (#22384) — and 4b836fd points the other way, warning about nondeterminism when a quantized V cache meets speculative checkpoints. Before opening a PR I'd like to know whether the storage side is something you'd entertain at all.

What I've been running

A codec at the state serialization boundary, off by default:

--rs-checkpoint-codec N    # 0 = fp32 (default), 1 = RHT-INT16, 2 = RHT-INT8

Each row is rotated by a random Hadamard transform (sign vector derived from a fixed seed, so it is reproducible and carries no side data), then stored as int16 or int8 with a per-row scale. Only the write path is affected. The read path detects a magic header, so checkpoints written with and without the flag load interchangeably. Non-finite values or degenerate scales fall back to fp32 for that block, and a metadata mismatch is fail-closed. Restored state runs in fp32; nothing in the live compute path changes.

It hooks the recurrent memory's state io, so it applies to recurrent and hybrid memory generally, though I have only tested the two models below.

Setup for everything that follows: Qwen3.8-27B (qwen35, 65 blocks, one MTP layer) on 4 GPUs and Qwen3.6-35B-A3B (qwen35moe) on 2, both Q8_0 GGUF with -sm tensor and MTP speculation on, server at ece963f. The weight quant is incidental to these numbers, since the recurrent state is fp32 either way. It has been running on all three instances in production here.

Measurements

Roundtrip, forced continuation of 48 tokens:

model codec checkpoint mean rel. L2 worst token flips
27B fp32 149.63 MiB 0/48
27B RHT-INT16 74.82 MiB (2.00x) 0.0049% 0.0067% 0/48
27B RHT-INT8 37.41 MiB (4.00x) 1.27% 1.74% 0/48
35B RHT-INT16 31.41 MiB (2.00x) 0.0037% 0.0050% 0/48
35B RHT-INT8 15.70 MiB (4.00x) 0.96% 1.28% 1/48

Server-level A/B on two identical instances, greedy sampling, every continuation going through a real restore:

27B RHT-INT16 35B RHT-INT16 35B RHT-INT8
token flips vs fp32 0/360 (0.0%) 65/349 (18.6%) 234/338 (69.2%)
divergence/rewind test 5/5 identical 5/5 correct, 4/5 differ textually

fp32 against fp32 on the same pair is 0/72, so the baseline really is deterministic. Across 42 restores there were no decode errors. Restore costs about 0.35 s for a 120–160 MiB checkpoint, and model load time is unchanged.

The caveat

Weight and KV quantization error is per-token and self-limiting. Checkpoint error is not: the restored state is the starting point for everything decoded after it, so the error feeds the entire continuation. That is what the two models above show — same codec, same L2, 0% versus 18.6% token divergence. Sensitivity is model-dependent and cannot be read off the L2 figure, which is why I want this opt-in and off by default, with the caveat stated in the flag help, in the same spirit as the 4b836fd warning.

INT8 fails my own gate on the 35B (69%) and I would keep it non-default, but I'd rather include the measurement than hide it.

On "why not just f16": at the same 2 bytes per element, int16 with a per-row scale after the rotation spends its whole mantissa on the row's actual dynamic range, while f16 gives up bits to an exponent the data doesn't need. If a plain f16 path would be the more welcome version, that is a much smaller change and I'd rather learn it now than after a PR. For what it's worth, the same codec ships as the default state codec in oMLX (an unrelated MLX-based project I also work on) at 1.93x and 0.0025% mean relative L2, which lines up with the numbers here.

Questions

  1. Is a lossy, default-off storage codec for checkpoints something you would consider?
  2. Scope. There is a second piece I have running. The draft context's KV and speculative state go into the same checkpoint at whatever --spec-draft-type-k/-v is set to, F16 by default, and unlike the recurrent plane that part grows with position: 84.7 MiB of a 234 MiB 27B checkpoint at 21.5k tokens, about 4.1 KiB per token, so it dominates at long context. Storing it as q8_0 while the live draft cache stays f16, plus skipping the per-cycle draft snapshot, takes the 35B checkpoint down a further 48.9%. One PR or two?
  3. Naming. I use --rs-checkpoint-codec locally, but --checkpoint-codec seems closer to the existing flags. Preference?

The branch is 18 files modified for about 220 insertions, plus a new codec source pair and two tests (a unit test and a roundtrip harness). I'll rebase onto master and re-run the gates before sending anything.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions