use std::sync::atomic::{AtomicU64, Ordering};
use ferrox_models::{Ceiling, ContextFit, KvBudget, KvElem, KvShape};
use crate::generate::DecodeError;
#[derive(Debug)]
pub struct ContextCeiling {
limit: Option<usize>,
shape: KvShape,
refused: AtomicU64,
}
impl ContextCeiling {
pub fn new(limit: Option<usize>, shape: KvShape) -> Self {
ContextCeiling {
limit,
shape,
refused: AtomicU64::new(0),
}
}
pub fn bytes_for(&self, positions: usize) -> u64 {
self.shape.kv_bytes_for_tokens(positions)
}
pub fn refused(&self) -> u64 {
self.refused.load(Ordering::Relaxed)
}
pub fn limit(&self) -> Option<usize> {
self.limit
}
pub fn prompt_refusal(&self, prompt_tokens: usize) -> Option<DecodeError> {
let limit = self.limit?;
if prompt_tokens < limit {
return None;
}
self.refused.fetch_add(1, Ordering::Relaxed);
Some(DecodeError::KvBudgetExceeded {
binding: Ceiling::ContextLength.code(),
estimated_bytes: self.bytes_for(prompt_tokens),
limit_bytes: self.bytes_for(limit),
positions: prompt_tokens,
positions_limit: limit,
detail: format!("prompt is too long: {prompt_tokens} tokens > {limit} maximum"),
})
}
pub fn overflow_refusal(&self, prompt_tokens: usize, max_tokens: usize) -> Option<DecodeError> {
match prompt_tokens.checked_add(max_tokens) {
Some(_) => None,
None => {
self.refused.fetch_add(1, Ordering::Relaxed);
Some(DecodeError::KvBudgetExceeded {
binding: Ceiling::ContextLength.code(),
estimated_bytes: 0,
limit_bytes: 0,
positions: usize::MAX,
positions_limit: self.limit.unwrap_or(usize::MAX),
detail: format!(
"prompt of {prompt_tokens} tokens plus max_tokens of {max_tokens} \
overflows the position counter, so this request cannot be served by \
any deployment. Send a max_tokens that fits the model's context"
),
})
}
}
}
pub fn refusal(&self, positions: usize) -> Option<DecodeError> {
let limit = self.limit?;
if positions <= limit {
return None;
}
self.refused.fetch_add(1, Ordering::Relaxed);
Some(DecodeError::KvBudgetExceeded {
binding: Ceiling::ContextLength.code(),
estimated_bytes: self.bytes_for(positions),
limit_bytes: self.bytes_for(limit),
positions,
positions_limit: limit,
detail: format!(
"request asks for {positions} token positions (prompt + max_tokens) but this \
deployment admits {limit} per request; shorten the prompt or lower max_tokens"
),
})
}
}
#[derive(Debug, Clone, Copy)]
pub struct DerivedLimits {
pub max_context: usize,
pub kv_blocks: usize,
pub fit: ContextFit,
}
pub fn derive_limits(
budget: &KvBudget,
gguf_ctx: usize,
block_size: usize,
) -> Option<DerivedLimits> {
assert!(block_size > 0, "kv block size must be positive");
let cap = gguf_ctx.max(1);
let fit = budget.max_context(cap, ferrox_models::CTX_AUTO_GRANULARITY);
if fit.tokens == 0 {
return None;
}
Some(DerivedLimits {
max_context: fit.tokens,
kv_blocks: fit.tokens / block_size,
fit,
})
}
pub fn apply_derived(
config: &mut crate::serving::batch::BatcherConfig,
derived: &DerivedLimits,
) -> Adopted {
let mut adopted = Adopted::default();
if config.max_context.is_none() {
config.max_context = Some(derived.max_context);
adopted.max_context = true;
}
if config.kv_blocks.is_none() && derived.kv_blocks > 0 {
config.kv_blocks = Some(derived.kv_blocks);
adopted.kv_blocks = true;
}
adopted
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct Adopted {
pub max_context: bool,
pub kv_blocks: bool,
}
pub fn price_gguf(
path: &str,
kv_elem: KvElem,
concurrent_requests: usize,
) -> Option<(KvBudget, usize, String)> {
use ferrox_models::residency_report::{ResidencyAssumptions, ResidencyReport};
use ferrox_models::{BudgetBackend, DeviceBudget};
let backend = if cfg!(feature = "metal") {
BudgetBackend::Metal
} else if cfg!(feature = "cuda") {
BudgetBackend::Cuda
} else {
BudgetBackend::Cpu
};
let device = DeviceBudget::detect(backend);
if device.is_unknown() {
tracing::info!("{device}; serving with no derived context ceiling");
return None;
}
let gguf_ctx = gguf_context_length(path)?;
let assumptions = ResidencyAssumptions {
context_tokens: gguf_ctx,
concurrent_requests: concurrent_requests.max(1),
kv_elem,
..ResidencyAssumptions::default()
};
match ResidencyReport::from_gguf(path, assumptions, device.usable_bytes) {
Ok(report) => Some((report.kv_budget(), gguf_ctx, device.to_string())),
Err(e) => {
tracing::info!(
"KV budget not computed for this checkpoint ({e}); serving with no derived \
context ceiling"
);
None
}
}
}
fn gguf_context_length(path: &str) -> Option<usize> {
let file = ferrox_gguf::ShardedGguf::open(path).ok()?;
let arch = file
.metadata_str("general.architecture")
.unwrap_or("unknown")
.to_string();
Some(
file.metadata_u64(&format!("{arch}.context_length"))
.map(|v| v as usize)
.unwrap_or(4096),
)
}
#[cfg(test)]
mod tests {
use super::*;
use ferrox_models::KvLayout;
fn shape() -> KvShape {
KvShape {
n_layers: 2,
layout: KvLayout::Gqa {
n_kv_heads: 1,
head_dim: 4,
},
elem: KvElem::F32,
}
}
fn budget(device_bytes: u64, weights_bytes: u64) -> KvBudget {
KvBudget {
weights_bytes,
activation_headroom_bytes: 0,
device_budget_bytes: device_bytes,
shape: shape(),
concurrent_requests: 1,
}
}
#[test]
fn per_token_bytes_are_what_the_hand_computation_says() {
assert_eq!(shape().per_token_kv_bytes(), 64);
}
#[test]
fn a_request_past_the_ceiling_is_refused_in_bytes_and_is_not_retryable() {
let ceiling = ContextCeiling::new(Some(100), shape());
assert!(ceiling.refusal(100).is_none(), "exactly at the limit fits");
let err = ceiling.refusal(101).expect("101 > 100 must be refused");
match &err {
DecodeError::KvBudgetExceeded {
binding,
estimated_bytes,
limit_bytes,
positions,
positions_limit,
..
} => {
assert_eq!(*binding, Ceiling::ContextLength.code());
assert_eq!(*estimated_bytes, 101 * 64);
assert_eq!(*limit_bytes, 100 * 64);
assert_eq!(*positions, 101);
assert_eq!(*positions_limit, 100);
}
other => panic!("expected KvBudgetExceeded, got {other:?}"),
}
assert_eq!(
err.retry_after_secs(),
None,
"an idle server refuses this identically, so 'retry shortly' would be a lie"
);
assert_eq!(ceiling.refused(), 1, "the refusal must be counted");
}
#[test]
fn no_ceiling_refuses_nothing() {
let ceiling = ContextCeiling::new(None, shape());
assert!(ceiling.refusal(usize::MAX / 2).is_none());
assert_eq!(ceiling.refused(), 0);
}
#[test]
fn the_derived_context_is_the_room_left_after_weights_divided_by_the_per_token_cost() {
let derived = derive_limits(&budget(8192, 4096), 100_000, 16)
.expect("4096 bytes of KV room fits some context");
assert_eq!(derived.max_context, 64);
assert_eq!(derived.kv_blocks, 4, "64 positions / 16 per block");
}
#[test]
fn a_roomy_machine_is_still_capped_at_the_models_trained_context() {
let derived = derive_limits(&budget(1 << 40, 0), 4096, 256)
.expect("a terabyte of room fits the model's whole context");
assert_eq!(derived.max_context, 4096);
assert_eq!(derived.kv_blocks, 16);
}
#[test]
fn a_partial_block_is_floored_away_rather_than_promised() {
let derived = derive_limits(&budget(6400, 0), 100_000, 64).expect("100 tokens fit");
assert_eq!(derived.max_context, 100);
assert_eq!(derived.kv_blocks, 1);
}
#[test]
fn a_model_that_leaves_no_room_derives_no_ceiling_at_all() {
assert!(derive_limits(&budget(4096, 4096), 100_000, 16).is_none());
assert!(derive_limits(&budget(4096, 8192), 100_000, 16).is_none());
}
fn derived(max_context: usize, kv_blocks: usize) -> DerivedLimits {
DerivedLimits {
max_context,
kv_blocks,
fit: budget(1 << 30, 0).max_context(max_context.max(1), 1),
}
}
#[test]
fn a_configured_ceiling_is_never_overridden_by_a_derived_one() {
let mut config = crate::serving::batch::BatcherConfig {
max_context: Some(999),
kv_blocks: Some(7),
..Default::default()
};
let adopted = apply_derived(&mut config, &derived(4096, 16));
assert_eq!(config.max_context, Some(999));
assert_eq!(config.kv_blocks, Some(7));
assert_eq!(adopted, Adopted::default(), "nothing was adopted");
}
#[test]
fn an_absent_ceiling_is_filled_and_the_two_slots_are_independent() {
let mut both = crate::serving::batch::BatcherConfig {
max_context: None,
kv_blocks: None,
..Default::default()
};
let adopted = apply_derived(&mut both, &derived(4096, 16));
assert_eq!(both.max_context, Some(4096));
assert_eq!(both.kv_blocks, Some(16));
assert_eq!(
adopted,
Adopted {
max_context: true,
kv_blocks: true
}
);
let mut half = crate::serving::batch::BatcherConfig {
max_context: Some(512),
kv_blocks: None,
..Default::default()
};
let adopted = apply_derived(&mut half, &derived(4096, 16));
assert_eq!(half.max_context, Some(512), "the set one survives");
assert_eq!(half.kv_blocks, Some(16), "the unset one is still derived");
assert_eq!(
adopted,
Adopted {
max_context: false,
kv_blocks: true
}
);
}
#[test]
fn a_fit_smaller_than_one_block_leaves_the_ledger_absent() {
let mut config = crate::serving::batch::BatcherConfig {
max_context: None,
kv_blocks: None,
..Default::default()
};
apply_derived(&mut config, &derived(100, 0));
assert_eq!(config.max_context, Some(100));
assert_eq!(config.kv_blocks, None);
}
#[test]
fn a_sliding_model_derives_its_ceiling_from_memory_like_any_other() {
let mut cfg = ferrox_models::config::test_dense_fixture();
cfg.n_layers = 2;
cfg.n_kv_heads = 1;
cfg.head_dim = 4;
cfg.sliding_window = Some(8);
cfg.swa_pattern = None; let windowed = KvShape::from_config(&cfg, KvElem::F32);
assert_eq!(windowed, shape(), "64 bytes/token, window or no window");
let b = KvBudget {
shape: windowed,
..budget(64 * 1024, 0)
};
let derived = derive_limits(&b, 8192, 256).expect("a fit of 1024 tokens is a real fit");
assert_eq!(derived.max_context, 1024);
}
#[test]
fn a_max_tokens_that_wraps_the_position_sum_is_refused() {
let ceiling = ContextCeiling::new(Some(100), shape());
let err = ceiling
.overflow_refusal(10, usize::MAX)
.expect("usize::MAX must be refused");
assert!(
format!("{err}").contains("max_tokens"),
"the refusal must name the field the caller sent"
);
assert!(
10usize.wrapping_add(usize::MAX) < 100,
"the wrap this refusal exists to catch"
);
}
#[test]
fn a_wrapping_sum_is_refused_even_with_no_ceiling_configured() {
let ceiling = ContextCeiling::new(None, shape());
assert!(
ceiling.overflow_refusal(10, usize::MAX).is_some(),
"no ceiling is not a licence to accept a request no machine could serve"
);
assert!(ceiling.overflow_refusal(10, 100).is_none());
}
#[test]
fn an_ordinary_request_past_the_ceiling_is_left_for_the_clamp() {
let ceiling = ContextCeiling::new(Some(100), shape());
assert!(ceiling.overflow_refusal(10, 50).is_none(), "10 + 50 fits");
assert!(
ceiling.overflow_refusal(10, 200).is_none(),
"over the ceiling but representable: the clamp handles it"
);
assert_eq!(ceiling.refused(), 0, "a clamp is not a refusal");
}
#[test]
fn a_huge_but_representable_max_tokens_reports_bytes_instead_of_panicking() {
let ceiling = ContextCeiling::new(Some(100), shape());
let positions = 10 + (u64::MAX / 64) as usize;
assert!(ceiling.bytes_for(positions) > u64::MAX / 2);
let err = ceiling
.refusal(positions)
.expect("far past a 100-position ceiling");
assert!(format!("{err}").contains("100"));
}
}