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,
) -> usizeExpand 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_sizelive per-seq blobs (h_state + conv_state across all SSM layers); - spec, per verify slot (
mtp_state_slotsof 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; seeverify_slot_h_intermediatesfor why conv does not tier); - 1 pre-verify checkpoint blob (h + conv).
- H intermediates:
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).