Module ssm_reserve

Module ssm_reserve 

Source
Expand description

SSOT for the Phase-C decode-rollback ring depth.

Two call sites MUST agree on this number or a serve either under-reserves (runtime CUDA alloc failure after weights load) or over-reserves (preflight refuses batch sizes the runtime could fund):

  • spark-server preflight_reserve — sizes the SSM-snapshot GPU reservation before weights load;
  • TransformerModel::new (impl_a1.rs) — allocates the actual ring.

The ring’s ONLY writer (scheduler snapshot_boundary_if_ssm) and reader (content-loop rollback_to_boundary) live on the PLAIN decode path — the speculative path does its rejection rollback through the verify snapshot, never this ring. Under --speculative the ring is unreachable, and it is NOT cheap: 8 slots × max_batch × the full SSM blob (27B: 158.9 MB) is ~19 GB at batch 16 and ~38 GB at batch 32. Reserving it unconditionally while the runtime skipped it capped the native batch at ~20 on GB10 (SSM reserve 75.2 GB vs an 85.2 GB budget at util 0.70).

Env contract (read HERE and nowhere else):

  • ATLAS_SSM_DECODE_RING=1 force-allocates the ring even under spec (mixed workloads whose grammar-bound sequences fall to plain decode and should keep loop re-steer); =0 force-disables it even without spec.
  • ATLAS_DISABLE_WATCHDOGS=1|true (trimmed, case-insensitive — mirrors spark-server’s parse_disable_watchdogs): the ring’s only reader can never fire, so the ring is skipped.

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).

Functions§

decode_rollback_ring_slots
Decide the per-sequence decode-rollback ring depth.
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).
prefix_caching_active
Whether the prefix cache this serve will actually install is a REAL cache.
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.