Module aux_state

Module aux_state 

Source
Expand description

2026-09-25: Snapshot and restore of the DSA indexer cache, the part of a GLM-5.3 sequence that a KV-only prefix-cache hit cannot rebuild.

Owner: model-arch (GLM-5.3 DSA). Invariants:

  • restore_blob checks the header, index_head_dim, total size and capacity before any copy, and returns an error for any mismatch.
  • The cursor moves only after every copy is enqueued, so a failed restore leaves len unchanged.
  • A blob is HEADER_BYTES + len * (4 * index_head_dim + 1) bytes: it carries len rows, not capacity.

§Why a blob and not a rewind

rewind_to restores within a live sequence: the rows stay and only the cursor moves. A prefix-cache hit hands the prefix to a sequence whose indexer buffer was never written, and the rows cannot be rebuilt from the MLA latent: k_normed and gate are projections of the hidden state (indexer.wk, index_kpool_compress_gate). So the rows travel with the snapshot.

layers/qsa_snapshot.rs in model-layers does the same for the QSA indexer, and re-pools its block keys on restore. The DSA state stores no pooled keys (pooling runs inside the selector kernel each step), so the blob is the whole reachable cache.

Applying a mismatched blob would make the selector read another sequence’s keys, which is why every mismatch is an error rather than a repair.

§Cost

len * (4 * index_head_dim + 1) bytes per DSA layer, 513 B per token and layer at index_head_dim = 128.