ssm_pool_reserve_bytes

Function ssm_pool_reserve_bytes 

Source
pub fn ssm_pool_reserve_bytes(
    max_batch_size: usize,
    h_blob_bytes: usize,
    conv_blob_bytes: usize,
    spec_on: bool,
    num_drafts: usize,
    mtp_state_slots: usize,
    uniform_verify: bool,
    h_f16_pool: bool,
    rollback: SsmRollbackMode,
) -> usize
Expand description

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

  • base: max_batch_size live per-seq blobs (h_state + conv_state across all SSM layers);
  • spec, per verify slot (mtp_state_slots of them):
    • H intermediates: verify_slot_h_intermediates × h blob (TIERED, and K-1 per K-row verify — index K-1 is never written or read);
    • conv intermediates: num_drafts + 1 × conv blob (uniform AND still K — the fused conv kernels write all K snapshots on-device; see verify_slot_h_intermediates for why conv does not tier);
    • 1 pre-verify checkpoint blob (h + conv).

h_blob_bytes / conv_blob_bytes are the per-seq totals across all SSM layers (num_ssm_layers × ssm_h_state_bytes/ssm_conv_state_bytes), ALWAYS at the FP32 width — h_f16_pool narrows every h term through ssm_h_stored_bytes inside, so preflight and SsmStatePool::new cannot narrow differently. The historical sizing was max_batch × blob × (1 + (num_drafts+1) + 1); today’s uniform mode differs from it by exactly one h blob per slot (the dead K-1 intermediate).