Module ssm_reserve

Expand description

SSOT for the SSM/linear-attention GPU reserve terms.

Every term here is computed TWICE — once by spark-server’s preflight_reserve before the weights load, once by the allocating call site (SsmStatePool::new, TransformerModel::new) after — and the two MUST agree or a serve either under-reserves (runtime CUDA alloc failure) or over-reserves (preflight refuses a configuration the runtime could fund). Each function below is the one place that decision is made.

The Phase-C decode-rollback ring DEPTH — its publication cell, the ATLAS_SSM_DECODE_RING / ATLAS_DISABLE_WATCHDOGS contract and the #915 auto-fit — lives in the decode_ring sibling module and is re-exported here, so every existing ssm_reserve::decode_rollback_ring_slots path is unchanged.

Structs§

DecodeRingDecision: Outcome of the ring-depth decision.
MarconiSlotDecision: Outcome of the Marconi snapshot-slot decision.

Enums§

SsmRollbackMode: SSM verify-rollback mode (--ssm-rollback-mode, EXPERIMENTAL scaffold).

Constants§

DECODE_RING_FIT_LADDER: Depths preflight’s #915 auto-fit may choose from, LARGEST FIRST.

Functions§

decode_rollback_ring_slots: Decide the per-sequence decode-rollback ring depth.
decode_rollback_ring_slots_with: Pure core of decode_rollback_ring_slots (env-free, unit-testable).
fit_decode_ring_slots: Largest DECODE_RING_FIT_LADDER depth <= start_slots whose ring term still fits in free_mem alongside everything else the reserve needs (#915). Pure — preflight owns the bytes, this owns the ladder.
marconi_reserve_full: The ATLAS_SSM_MARCONI_FULL kill switch (PRESENCE, house convention).
marconi_snapshot_slots: Number of Marconi SSM-snapshot slots to RESERVE and ALLOCATE.
marconi_snapshot_slots_with: Pure core of marconi_snapshot_slots (env-free, unit-testable).
mtp_pool_full_width: The ATLAS_MTP_POOL_FULL_WIDTH kill switch (PRESENCE, house convention — =0 is NOT off), plus the EP-v2 implication (v2 pins slots in place for the worker mirror, so a high slot may legitimately speculate forever). SSOT for BOTH pool diets it disables: the bs>32 slot-count cap (mtp_state_slots) and the tiered per-slot verify capacity (verify_slot_drafts) — one switch restores the full-width, uniform-K sizing everywhere (pool, preflight, scheduler clamp).
mtp_state_slots: Number of SSM-pool slots the MTP/DFlash VERIFY state pools (per-token intermediates + pre-verify checkpoints) must cover.
mtp_state_slots_with: Pure core of mtp_state_slots (env-free, unit-testable).
parse_decode_ring_slots: SSOT parse for the --ssm-decode-ring-slots value: auto (None — size it from free memory at preflight) or an explicit 0..=8.
prefix_caching_active: Whether the prefix cache this serve will actually install is a REAL cache.
published_decode_ring_slots: The published depth, or None when nothing has been published.
set_decode_ring_slots: Publish the decode-ring depth. Returns the value in force (first write wins, matching gdn_flags::set_from_cli): an explicit --ssm-decode-ring-slots N is published before preflight runs, so the auto-fit’s later write is a no-op and an operator’s explicit depth is never silently shrunk.
set_ssm_rollback_mode: Publish the command line’s mode. Returns the value in force (first write wins, matching gdn_flags::set_from_cli).
ssm_h_prefill_stage_bytes: FP32 h-state PREFILL STAGING bytes (stage 3 of --ssm-h-dtype f16).
ssm_h_stored_bytes: Storage width of one h-state blob in the SSM state pools (stage 3 of --ssm-h-dtype f16): 2 bytes per element under the f16-SIZED pool, the FP32 4 bytes otherwise. SSOT — SsmStatePool::new (allocation strides), preflight_reserve (the pre-load reserve) and every byte-copier that moves h-state between pool regions derive their width from THIS, so sizing and copies cannot disagree.
ssm_pool_reserve_bytes: SSM state-pool reserve bytes for the pre-load preflight — MUST mirror what SsmStatePool::new allocates (modulo the +1 dummy slot per pool, which preflight has never counted; the CUDA headroom term absorbs it):
ssm_replay_ring_bytes: Replay-mode verify-window input ring: k_ceiling - 1 cached rows per covered slot per SSM layer (a partial accept replays at most K-1 tokens — rows 0..K-2; a full accept replays nothing). Reserved by preflight and allocated by SsmStatePool::new through THIS function so the two cannot disagree. Zero when speculation is off or the mode is Snapshot.
ssm_replay_row_bytes: One cached verify-row of GDN inputs for replay, per SSM layer: the deinterleaved qkvz row (qkvz_elems BF16 — what conv1d consumes; Z included, the gated norm needs it) + the gate/beta row (nv * 2 FP32). These are exactly the per-token tensors the WY verify kernels read (ConvGdnArgs::deinterleaved / gates_buf rows), and re-running them through the sequential conv+GDN path from the checkpoint reproduces the snapshot the dropped intermediates used to hold.
ssm_rollback_mode: The mode in force. Snapshot when nothing was published — mirroring the CLI’s explicit default for non-serve contexts (tests, examples), which never carry the flag. Production sizing/pool call sites take the mode as a PARAMETER and read this only at the outermost boundary, so unit tests never depend on the process-global cell.
verify_slot_drafts: Env-reading wrapper of verify_slot_drafts_with: the ladder policy (with its ATLAS_MTP_K_LADDER / ATLAS_NO_MTP_K_LADDER overrides — a disabled ladder returns num_drafts at every width, making the tiers vacuous) plus the mtp_pool_full_width kill switch.
verify_slot_drafts_with: Per-slot verify DRAFT capacity — the tiered half of the verify-pool diet (2026-08-16). Pure core; drafts_at(n) is the ladder policy (speculative::mtp_ladder_drafts).
verify_slot_h_intermediates: Number of per-token H-state intermediates the verify pools allocate for pool slot slot_idx: exactly the slot’s draft capacity (K-1 snapshots for a K-row verify). uniform_verify (DFlash-γ pools, whose verify width does not follow the MTP ladder) sizes every slot at the full num_drafts.
watchdogs_disabled_from_value: ATLAS_DISABLE_WATCHDOGS truthiness — trimmed, case-insensitive, 1 or true only (mirrors spark-server’s parse_disable_watchdogs).