mtp_state_slots

Function mtp_state_slots 

Source
pub fn mtp_state_slots(max_batch_size: usize) -> usize
Expand description

Number of SSM-pool slots the MTP/DFlash VERIFY state pools (per-token intermediates + pre-verify checkpoints) must cover.

Three call sites MUST agree on this number (same contract as the decode ring above):

  • spark-server preflight_reserve — sizes the pre-load GPU reserve;
  • SsmStatePool::new — allocates the intermediate/checkpoint pools;
  • the scheduler’s spec dispatch — gates every speculative step on slot_idx < mtp_state_slots(..) so an uncovered slot can never be verified (uncovered slots plain-decode until retirement-time compaction migrates them under the cap).

WHY a cap exists: the verify pools were sized max_batch_size × K even though spec dispatch is bounded by speculative::mtp_max_seqs() (default 32 — the widest batched-verify chunk, layer::VERIFY_WY_TABLE_SEQS). On the 27B at --max-batch-size 64 with --num-drafts 3 that is 32 dead slots × 5 SSM blobs × 158.9 MB = 25.4 GB of reserve for states no code path can ever touch — the difference between bs=64 refusing at preflight (util 0.70) and booting.

The cap NEVER bites at max_batch_size <= 32: the floor is VERIFY_WY_TABLE_SEQS (32), so bs<=32 sizing and behavior are byte-identical in every env combination (slots are always < bs).

Env contract (read HERE and nowhere else):

  • ATLAS_MTP_POOL_FULL_WIDTH (presence, house convention — =0 is NOT off): restore full-width pools (max_batch_size slots) and make the scheduler guard vacuous. Kill switch for the bs>32 reserve diet.
  • ATLAS_EP_PROTOCOL=v2 implies full width: v2 pins slots in place for the worker mirror (no compaction — see retire_finished_sequences), so a high slot may legitimately speculate forever.
  • ATLAS_MTP_MAX_SEQS participates via crate::speculative::mtp_max_seqs: raising the dispatch cap above 32 widens the pools with it.

★ WHAT THE DIET COSTS, AND THE UTILISATION FLOOR IT SETS (wave 47, dgx3, 27B W4A4). The diet is what makes a single serve able to cover the whole concurrency ladder — speculation is dispatch-capped at 32, so one serve at --max-batch-size 128 --speculative --num-drafts 3 speculates at C<=32 and plain-decodes above it. But the verify pools it keeps are still sized by --num-drafts, and at bs=128 that is not free. Measured preflight reserve, --max-seq-len 4096, blob 151.5 MB:

configbaseverify poolssnapshot/miscreserve
bs=128, spec OFF18.9 GB (128 blobs)5.5 GB24.3 GB
bs=128, spec ON, 3 drafts18.9 GB23.7 GB (32 slots x 5 blobs)8.9 GB51.5 GB

With 39.8 GB already consumed before KV, that reserve REFUSES at --gpu-memory-utilization 0.70 (39.8 + 51.5 = 91.3 GB committed against an 85.2 GB budget) and boots at 0.85 (103.4 GB budget, 13.3 GB left for KV = 217k tokens). The floor for the one-serve ladder is therefore util ~0.82, and it is set HERE, by the verify pools — not by the KV dtype, which moves the answer by well under a GB at these widths. A cheaper diet (row-budget-sized intermediates rather than slot-major) would recover ~9 GB and still not reach 0.70; the reserve, not the speculation regime, is what makes the low-util single config impossible.