Module ssm_reserve

Expand description

2026-09-25: GPU reserve terms for SSM (linear-attention) state: verify-pool slot counts and capacities, pool bytes, Marconi snapshot slots, and the decode-rollback ring depth (decode_ring).

Owner: model-layers (SSM reserve). Invariants: none beyond the types.

The server’s preflight_reserve, which runs before the weights load, and the allocating code (SsmStatePool::new, TransformerModel::new) call the same functions here. If the two computed different sizes, a serve would fail a CUDA allocation after load, or refuse a configuration it could run. 2026-09-30: the pool itself is one plan both read (PoolPlan).

Structs§

CopyTier: 2026-10-02: Copy-tier geometry.
DecodeRingDecision: 2026-09-25: Outcome of the ring-depth decision. skip_reason is Some only for the implicit skips (speculative decode, watchdogs disabled), so the allocating call site can log the saving.
MarconiSlotDecision: 2026-09-25: Outcome of the Marconi snapshot-slot decision. skip_reason is Some only when the slots are dropped because prefix caching is inactive, so the allocating call site can log the saving.
PoolCounts: 2026-09-30: The pool’s unit counts.
PoolPlan: 2026-09-30: The SSM pool of one model, sized.
PoolShape: 2026-09-30: What decides the pool’s unit counts.
VerifyCounts: 2026-09-30: The verify pools’ unit counts.

Enums§

PoolState: 2026-09-30: A recurrent state of the pool.
SsmRollbackMode: 2026-09-25: SSM verify-rollback mode (--ssm-rollback-mode, experimental).
UnitSource: 2026-09-30: Where a pool’s per-unit sizes come from.

Constants§

DECODE_RING_FIT_LADDER: 2026-09-25: Depths the preflight fit may choose, largest first. The first equals metrale_kernels::DECODE_ROLLBACK_RING_SLOTS (8); each next depth halves the ring’s share of the reserve.
MTP_STATE_SLOTS_FLOOR: 2026-09-29: The fewest slots the verify pools cover (when max_batch_size allows): the verify table width before 2026-09-29, kept apart from it so a wider verify table does not grow every serve’s pools (137.6 MB per slot on Qwen3.6-35B-A3B).

Functions§

copy_tier: 2026-10-02: The published tier; None without prompt lookup.
decode_rollback_ring_slots: 2026-09-25: The per-sequence ring depth, from the published depth and the environment (decode_rollback_ring_slots_with). Preflight passes args.speculative || args.dflash as use_speculative; the allocation must receive the same flag, or the reserve and the ring differ.
decode_rollback_ring_slots_with: 2026-09-25: Env-free core of decode_rollback_ring_slots. Precedence, highest first:
fit_decode_ring_slots: 2026-09-25: Largest DECODE_RING_FIT_LADDER depth, at most start_slots, with reserve_without_ring + depth * bytes_per_ring_slot <= free_mem (saturating arithmetic). bytes_per_ring_slot is one depth unit of the ring. Returns 0 also when no depth fits; the caller tells the two apart by comparing the total against its limit (preflight/decode_ring.rs).
marconi_reserve_full: 2026-09-25: Whether METRALE_SSM_MARCONI_FULL is present, with any value.
marconi_snapshot_slots: 2026-09-25: Number of Marconi SSM-snapshot slots to reserve and allocate: requested (--ssm-cache-slots, default 16), or 0 when prefix caching is inactive.
marconi_snapshot_slots_with: 2026-09-25: Env-free core of marconi_snapshot_slots.
mtp_pool_full_width: 2026-09-25: True when METRALE_MTP_POOL_FULL_WIDTH is present (any value, 0 included) or METRALE_EP_PROTOCOL is v2. Then mtp_state_slots is max_batch_size and verify_slot_drafts is num_drafts for every slot.
mtp_state_slots: 2026-09-25: Number of SSM-pool slots the verify pools (per-token intermediates and the pre-verify checkpoint) cover.
mtp_state_slots_with: 2026-09-25: Env-free core of mtp_state_slots. The MTP_STATE_SLOTS_FLOOR (32) covers every slot when max_batch_size is at most 32, also when METRALE_NO_MTP_K_LADDER lowers the dispatch cap to 4.
parse_decode_ring_slots: 2026-09-25: Parses --ssm-decode-ring-slots: auto gives None (fit at preflight), a number up to DECODE_ROLLBACK_RING_SLOTS gives Some, and anything else an error. validate_serve_args and publish_kernel_flags both parse through this.
pool_counts: 2026-09-30: pool_counts_with with the verify slots mtp_state_slots and the per-slot h intermediates verify_slot_h_intermediates (the MTP ladder).
pool_counts_tiered: 2026-10-02: pool_counts with the copy tier passed in, so tests need not publish it. The tier widens num_intermediates (every slot’s conv intermediates and the dummy’s h) and raises the copy slots’ h; a uniform-h pool takes the widened depth everywhere.
pool_counts_with: 2026-09-30: Pure: the counts for mtp_slots verify slots, slot s tiered to slot_h(s) h intermediates (capped at K - 1). The verify dummy, at index mtp_slots, always gets K - 1: batched-verify padding rows may write any live token index. Replay keeps no intermediates, only the checkpoints.
prefix_caching_active: 2026-09-25: Whether the prefix cache the server will install is a real cache: the flag and ModelConfig::kv_only_prefix_cache_is_safe, the predicate build_prefix_cache uses. Preflight calls this before the cache exists; TransformerModel::new asks the constructed cache (PrefixCache::is_active).
published_decode_ring_slots: 2026-09-25: The published depth, or None. Reads the cell without setting it.
recurrent_units: 2026-09-30: One recurrent layer’s h and conv declarations, and where they came from.
set_copy_tier: 2026-10-02: Publish the copy tier, before the preflight reserve.
set_decode_ring_slots: 2026-09-25: Publishes the ring depth and returns the depth in force; the first write wins. Preflight skips its fit when a depth is already published (published_decode_ring_slots), so an explicit depth is not shrunk.
set_ssm_rollback_mode: 2026-09-25: Publish the command line’s mode. Returns the mode in force (first write wins).
ssm_h_prefill_stage_bytes: 2026-09-25: Bytes of the FP32 h-state prefill staging arena: slots * h_layer_f32_bytes when f16_pool, else 0.
ssm_h_stored_bytes: 2026-09-25: Bytes of one stored h-state blob: half of h_f32_bytes when f16_pool (FP16 elements), else h_f32_bytes. SsmStatePool::new, the preflight reserve (PoolPlan) and the SSM layer’s slot stride (ssm_h_fp16.rs) all take the stored width from here.
ssm_replay_ring_bytes: 2026-09-25: Replay-mode verify-window input ring: mtp_state_slots * (k_ceiling - 1) * num_ssm_layers * row_bytes, since a partial accept replays at most K-1 tokens. Preflight reserves and SsmStatePool::new allocates through this function, with different arguments: preflight passes num_drafts + 1 and mtp_state_slots, the pool its num_intermediates and mtp_slots + 1 (the dummy slot).
ssm_replay_row_bytes: 2026-09-25: Bytes of one cached verify row of GDN inputs, per SSM layer: the deinterleaved qkvz row (qkvz_elems BF16, the rows of ConvGdnArgs::deinterleaved) plus the gate/beta row (nv * 2 FP32, the rows of ConvGdnArgs::gates_buf).
ssm_rollback_mode: 2026-09-25: The mode in force. Snapshot when nothing was published (tests, examples); reading it that way also fixes it for the process.
state_dims: 2026-09-30: The dims the circuits’ recurrent state shapes read, from this process’s config (after tensor-parallel sharding, topology.rs).
tier_h: 2026-10-02: Slot s’s h intermediates under tier: base raised to the copy depth on a copy slot.
tier_rows: 2026-10-02: Rows of the widest verify the pools must hold: the MTP k (drafts + 1), widened to a copy’s.
verify_slot_drafts: 2026-09-25: verify_slot_drafts_with with the dispatch cap mtp_max_seqs() and the ladder mtp_ladder_drafts. Every slot gets num_drafts when mtp_pool_full_width holds or the ladder is disabled.
verify_slot_drafts_with: 2026-09-25: Verify draft capacity of pool slot slot_idx: the largest count drafts_at(n) gives for a width n in slot_idx + 1..=max(dispatch_cap, slot_idx + 1), clamped to 1..=num_drafts; 0 when num_drafts is 0.
verify_slot_h_intermediates: 2026-09-25: H-state intermediates the verify pools hold for pool slot slot_idx: the slot’s draft capacity, which is K-1 for a K-row verify. With uniform_verify (set by the server’s preflight for DFlash) every slot gets num_drafts.
watchdogs_disabled_from_value: 2026-09-25: METRALE_DISABLE_WATCHDOGS truthiness: 1 or true, trimmed and case-insensitive, as the server’s parse_disable_watchdogs reads it.