Struct StateBudget 

pub struct StateBudget {
    pub kda_recurrent: usize,
    pub kda_conv: usize,
    pub dsa_kv_per_token: usize,
    pub dsa_indexer_per_token: usize,
    pub mhc_highway_per_token: usize,
    pub moe_routing_per_token: usize,
}

Per-sequence state contract (Slice 11 gate 3).

Derived from the checkpoint config, not from a formula chosen to look tidy. Every field is bytes for ONE sequence at EP=1; per_rank halves only what EP actually shards.

Fields

kda_recurrent: usize

KDA recurrent state — [heads, head_dim, head_dim] fp32 mandatory (HF casts to float32 and vLLM hardcodes it), 34 layers. FIXED: does not grow with sequence length.

kda_conv: usize

KDA causal-conv window, bf16, conv_dim x (kernel - 1 + num_spec). FIXED.

dsa_kv_per_token: usize

DSA MLA KV per TOKEN across the 11 text layers — kv_lora_rank bf16, NoPE so there is no rope section. GROWS with sequence length.

dsa_indexer_per_token: usize

Indexer key + gate state per TOKEN across the 11 text layers. GROWS. 🪤 REPLICATED, not sharded: the indexer is DsaShard::Replicated (dsa/tp.rs), so Self::per_rank must not divide it.

mhc_highway_per_token: usize

mHC highway per TOKEN — hc_mult x hidden fp32 in Atlas (bf16 in HF; see the OPEN highway-dtype item). Activation-lifetime, not persistent across steps.

moe_routing_per_token: usize

MoE routing scratch per token: logits + top-k ids + weights.

Implementations

impl StateBudget

pub fn fixed(&self) -> usize

Fixed (sequence-length-independent) bytes per sequence.

pub fn per_token(&self) -> usize

Bytes that grow with every token of context.

pub fn for_sequence(&self, tokens: usize) -> usize

Total persistent state for a sequence of tokens.

pub fn per_rank(&self, ep: usize) -> StateBudget

EP shards the KDA head dimension and the KV heads; the DSA INDEXER cache, the mHC highway and the routing scratch are replicated. Only what EP actually shards is divided.

Trait Implementations

impl Clone for StateBudget

fn clone(&self) -> StateBudget

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for StateBudget

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Copy for StateBudget