use metrale_config::ModelConfig;
use metrale_model_arch::seq_state_reserve::per_sequence_state_bytes;
use crate::cli;
pub(super) fn per_sequence_reserve(args: &cli::ServeArgs, config: &ModelConfig) -> usize {
let spec_on = args.speculative || args.self_speculative || args.dflash;
let per_seq = per_sequence_state_bytes(config, args.max_seq_len, spec_on).unwrap_or_default();
let charge = per_seq.for_batch(args.max_batch_size.ceiling());
if matches!(config.model_type.as_str(), "glm5_next" | "glm5_next_text")
&& metrale_config::glm_dsa_indexer_paged()
{
tracing::info!(
"GLM-5.3 DSA indexer: paged in the V side of each DSA layer's KV pool, {} B per \
token per DSA layer inside the KV price; nothing reserved per sequence \
(METRALE_GLM_DSA_INDEXER_FLAT=1 restores the flat cache)",
4 * config.index_head_dim
);
}
if charge > 0 {
tracing::info!(
"Per-sequence state reserve: {} MB = {} seq x ({} MB target DSA layers + {} MB \
proposer). Owned per sequence, replicated per rank (EP does not shard the \
indexer); previously covered only by cuda_headroom.",
charge / (1024 * 1024),
args.max_batch_size.ceiling().max(1),
per_seq.target_layers / (1024 * 1024),
per_seq.proposer / (1024 * 1024),
);
}
charge
}
#[cfg(test)]
#[path = "per_sequence_state_tests.rs"]
mod tests;