metrale_model_layers/

ssm_reserve.rs

// SPDX-License-Identifier: MIT OR Apache-2.0

//! 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`]).

mod decode_ring;
pub use decode_ring::{
    DECODE_RING_FIT_LADDER, DecodeRingDecision, decode_rollback_ring_slots,
    decode_rollback_ring_slots_with, fit_decode_ring_slots, parse_decode_ring_slots,
    published_decode_ring_slots, set_decode_ring_slots, watchdogs_disabled_from_value,
};

mod copy_tier;
pub use copy_tier::{CopyTier, copy_tier, set_copy_tier, tier_h, tier_rows};
mod pool_plan;
pub use pool_plan::{
    PoolCounts, PoolPlan, PoolShape, PoolState, UnitSource, VerifyCounts, pool_counts,
    pool_counts_tiered, pool_counts_with, recurrent_units, state_dims,
};

mod rollback;
pub use rollback::{
    SsmRollbackMode, set_ssm_rollback_mode, ssm_replay_ring_bytes, ssm_replay_row_bytes,
    ssm_rollback_mode,
};

/// 2026-09-25: Number of SSM-pool slots the verify pools (per-token intermediates and
/// the pre-verify checkpoint) cover.
///
/// Three callers use it: the server's `preflight_reserve` sizes the pre-load
/// reserve with it, `SsmStatePool::new` allocates that many verify slots plus a
/// dummy, and for a model with SSM layers the scheduler takes the plain-decode
/// branch while any active slot is at or above it (`spec_slot_cap`,
/// `lane_decode.rs`). Retirement compacts survivors onto low slots, except
/// under EP protocol v2 (`retire_finished_sequences`).
///
/// It is `max_batch_size` when that is at most 32 ([`MTP_STATE_SLOTS_FLOOR`],
/// the floor) or when [`mtp_pool_full_width`] holds; otherwise
/// `max(mtp_max_seqs(), 32)`, capped at `max_batch_size`. The scheduler never
/// speculates with more than `mtp_max_seqs()` active sequences.
pub fn mtp_state_slots(max_batch_size: usize) -> usize {
    mtp_state_slots_with(
        max_batch_size,
        crate::speculative::mtp_max_seqs(),
        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.
pub fn mtp_pool_full_width() -> bool {
    std::env::var_os("METRALE_MTP_POOL_FULL_WIDTH").is_some()
        || matches!(std::env::var("METRALE_EP_PROTOCOL").as_deref(), Ok("v2"))
}

/// 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).
pub const MTP_STATE_SLOTS_FLOOR: usize = 32;

/// 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.
pub fn mtp_state_slots_with(
    max_batch_size: usize,
    spec_dispatch_cap: usize,
    full_width: bool,
) -> usize {
    if full_width {
        return max_batch_size;
    }
    max_batch_size.min(spec_dispatch_cap.max(MTP_STATE_SLOTS_FLOOR))
}

/// 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.
///
/// Retirement compacts active sequences onto slots `0..n`
/// (`retire_finished_sequences`), so a sequence on slot `slot_idx` usually runs
/// with at least `slot_idx + 1` active, hence the lower bound. When it does not,
/// the scheduler clamps the step's draft count to the smallest capacity among
/// the active slots (`spec_capacity::clamp_drafts_to_slot_capacity`).
///
/// With the default ladder (`4:3,8:3,16:1,32:1`), `--num-drafts 3` and a
/// dispatch cap of 32, slots 0..8 get 3 drafts and slots from 8 on get 1. So
/// `adaptive_rung`'s two-draft lift at widths 9..=16 is clamped back to one
/// draft at every such width: nine or more active sequences hold as many
/// distinct slots, one of them at index 8 or above. [`mtp_pool_full_width`]
/// gives every slot `num_drafts`, which keeps the lift.
pub fn verify_slot_drafts_with(
    slot_idx: usize,
    dispatch_cap: usize,
    num_drafts: usize,
    drafts_at: impl Fn(usize) -> usize,
) -> usize {
    if num_drafts == 0 {
        return 0;
    }
    let hi = dispatch_cap.max(slot_idx + 1);
    ((slot_idx + 1)..=hi)
        .map(&drafts_at)
        .max()
        .unwrap_or(num_drafts)
        .clamp(1, num_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.
pub fn verify_slot_drafts(slot_idx: usize, num_drafts: usize) -> usize {
    if mtp_pool_full_width() {
        return num_drafts;
    }
    verify_slot_drafts_with(
        slot_idx,
        crate::speculative::mtp_max_seqs(),
        num_drafts,
        |n| crate::speculative::mtp_ladder_drafts(n, num_drafts),
    )
}

/// 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`.
///
/// K-1 suffices because the state after the last verified row stays in
/// `h_state`: `commit_accepted_prefix` returns early on a full accept and
/// otherwise reads index `num_accepted - 1 <= K - 2` (`async_chkpt.rs`), and
/// the batched WY path requires only k-1 h intermediates
/// (`trait_decode_batched_conv_gdn_multi.rs`).
///
/// Only the H side is tiered. The conv intermediates stay at `num_drafts + 1`
/// per slot: `gdn_verify_fused_conv_kn_batched` writes all K, and the batched
/// path declines unless every sequence holds k conv intermediates and the
/// per-sequence regions are evenly spaced (`trait_decode_batched_conv_gdn_multi.rs`).
pub fn verify_slot_h_intermediates(
    slot_idx: usize,
    num_drafts: usize,
    uniform_verify: bool,
) -> usize {
    if uniform_verify {
        return num_drafts;
    }
    verify_slot_drafts(slot_idx, num_drafts)
}

/// 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.
///
/// `f16_pool` is `ssm_h_f16_pool_enabled()` (`--ssm-h-dtype f16-pool`) at the
/// production call sites; it is a parameter so pool geometry is testable
/// without the process-global flag.
///
/// Panics when `h_f32_bytes` is not a multiple of 4.
pub fn ssm_h_stored_bytes(h_f32_bytes: usize, f16_pool: bool) -> usize {
    assert!(
        h_f32_bytes.is_multiple_of(4),
        "h-state blobs are FP32-element sized"
    );
    if f16_pool {
        h_f32_bytes / 2
    } else {
        h_f32_bytes
    }
}

/// 2026-09-25: Bytes of the FP32 h-state prefill staging arena: `slots *
/// h_layer_f32_bytes` when `f16_pool`, else 0.
///
/// On an f16-sized pool the layer widens a slot's h-state into its FP32
/// staging blob before the GDN prefill kernels run and narrows it back after
/// (`ssm_h_fp16::prefill_h_begin` / `prefill_h_end`). There is one blob per
/// slot, shared by all layers, so `h_layer_f32_bytes` is one layer's FP32 h
/// blob, not the across-layers total the other reserve terms use.
///
/// `SsmStatePool::new` passes its slot count including the dummy slot
/// (`max_slots + 1`); the server's preflight passes `max_batch_size`.
pub fn ssm_h_prefill_stage_bytes(slots: usize, h_layer_f32_bytes: usize, f16_pool: bool) -> usize {
    if f16_pool {
        slots * h_layer_f32_bytes
    } else {
        0
    }
}

/// 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.
pub struct MarconiSlotDecision {
    pub slots: usize,
    pub skip_reason: Option<&'static str>,
}

/// 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.
///
/// The server's `preflight_reserve` and `TransformerModel::new` (`impl_a1.rs`)
/// both call this. A slot is restored from a prefix-cache match
/// (`PrefixMatch::ssm_snapshot`); with prefix caching inactive the server
/// installs `NoPrefixCaching` (`build_prefix_cache`), whose lookup returns an
/// empty match.
///
/// `METRALE_SSM_MARCONI_FULL`, present with any value (`0` included), keeps
/// `requested` regardless.
pub fn marconi_snapshot_slots(
    requested: usize,
    prefix_caching_active: bool,
) -> MarconiSlotDecision {
    marconi_snapshot_slots_with(requested, prefix_caching_active, marconi_reserve_full())
}

/// 2026-09-25: Whether `METRALE_SSM_MARCONI_FULL` is present, with any value.
pub fn marconi_reserve_full() -> bool {
    std::env::var_os("METRALE_SSM_MARCONI_FULL").is_some()
}

/// 2026-09-25: Env-free core of [`marconi_snapshot_slots`].
pub fn marconi_snapshot_slots_with(
    requested: usize,
    prefix_caching_active: bool,
    full_reserve: bool,
) -> MarconiSlotDecision {
    if requested == 0 || prefix_caching_active || full_reserve {
        return MarconiSlotDecision {
            slots: requested,
            skip_reason: None,
        };
    }
    MarconiSlotDecision {
        slots: 0,
        skip_reason: Some("prefix caching inactive — Marconi snapshot slots are unreachable"),
    }
}

/// 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`).
pub fn prefix_caching_active(enable_flag: bool, kv_only_prefix_cache_is_safe: bool) -> bool {
    enable_flag && kv_only_prefix_cache_is_safe
}

#[cfg(test)]
#[path = "ssm_reserve_tests.rs"]
mod mtp_state_slot_tests;