From 1b1ced730d43593bde927ae0f24609c0834e7c3d Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 5 Nov 2026 12:24:58 +0900 Subject: [PATCH 1/2] fix(xla): size the static context shape from the request The OpenXLA static capacity defaulted to 256 tokens, so any prompt of 255 tokens or more was refused at admission, and the default `-n` (unlimited) resolved to the whole checkpoint window and was refused even for short prompts. Size the CLI shape from the request through a capacity ladder bounded by the checkpoint window, bound the unlimited budget to a 512-token sizing default and clamp the admitted budget to the room the shape leaves, and derive a 4096 serving bucket next to the 256 floor while keeping the short-request shape. A pinned `MLXCEL_XLA_CONTEXT_CAPACITY` and a pinned `MLXCEL_XLA_CONTEXT_BUCKETS` keep their existing precedence and single-shape behavior. Refs #2855 --- docs/environment-variables.md | 4 +- src/backend/mod.rs | 66 +++++++++ src/backend/xla.rs | 66 ++++++--- src/commands/generate.rs | 86 ++++++++++- src/commands/generate_tests.rs | 24 +++ src/lib/mlxcel-xla/src/context.rs | 236 ++++++++++++++++++++++++++++-- src/lib/mlxcel-xla/src/lib.rs | 5 +- src/server/model_worker.rs | 22 +-- 8 files changed, 461 insertions(+), 48 deletions(-) diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 575986ff7..97fe2b209 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -459,8 +459,8 @@ for how to build the backend. | `MLXCEL_XLA_LOCAL_TASK_WORKERS` | integer, 1 to 256 | unset (one worker per physical core of the current NUMA node, at most 64) | Worker count of the `local-task` (CPU) device. The shim builds that device itself with 8 MiB worker stacks instead of reading IREE's global task flags, whose 128 KiB default stack cannot hold the static TLS of a process that has loaded the CUDA or ROCm runtime libraries (`thread creation failed with 22`). Read only when the device is `local-task`; the valid range is 1 to 256 (IREE's `IREE_TASK_TOPOLOGY_MAX_GROUP_COUNT`); `0`, a negative, a non-integer or a value above 256 fails the load with an error naming the variable instead of being clamped silently. | | `MLXCEL_XLA_ROCM_TARGET` | single `gfx*` chip SKU (e.g. `gfx1151`, `gfx1100`, `gfx942`) | unset | Required when `MLXCEL_XLA_DEVICE=hip`. Passed as `--iree-rocm-target=` to `iree-compile`. Unset, empty, comma-separated, or non-`gfx*` values are rejected before spawning `iree-compile` with a `[compile]` diagnostic. | | `MLXCEL_XLA_VISION_BACKEND` | `auto`, `iree`, `host` | `auto` | LLaVA vision-tower/projector selector for the shared CLI/server prepared-prefill path. `auto` tries the resident IREE module and emits the exact startup failure before using the host implementation; `iree` fails startup instead of falling back; `host` explicitly keeps vision on MLX. The selected backend and per-request IREE transfer/timing metrics are emitted through tracing. | -| `MLXCEL_XLA_CONTEXT_CAPACITY` | integer tokens in `1..=2147483647` | `256` | Static sequence capacity compiled into the OpenXLA prefill/decode pair and every KV buffer. Requests are rejected before native execution when `effective_prompt_len + max_new_tokens` exceeds this value. Increasing it grows KV memory linearly and the prefill attention mask quadratically, and produces a distinct compiled artifact that may take longer to compile. StableHLO dynamic sequence shapes are not used. | -| `MLXCEL_XLA_CONTEXT_BUCKETS` | comma-separated integers, each in `1..=2147483647` | derived (`256`, plus the worst-case image floor when larger) | Server OpenXLA continuous-batching worker only (`xla-iree`): the static context shapes it compiles, sorted and deduplicated, for an operator who knows the traffic better than the derivation. Ignored when `MLXCEL_XLA_CONTEXT_CAPACITY` is set, which yields a single bucket. An empty list, a `0`, or a non-integer aborts worker startup. The CLI path does not read it. | +| `MLXCEL_XLA_CONTEXT_CAPACITY` | integer tokens in `1..=2147483647` | unset (request-sized on the CLI, derived buckets on the server) | When set, pins exactly one static sequence capacity compiled into the OpenXLA prefill/decode pair and every KV buffer, and requests are rejected before native execution when `effective_prompt_len + max_new_tokens` exceeds it. When unset, the CLI sizes the shape from the request: the smallest rung of `[256, 512, 1024, 2048, 4096, 8192, 16384, 32768]` that admits the prompt plus the generation budget, capped by the checkpoint's `max_position_embeddings`; an unlimited `-n -1` sizes with a 512-token budget and the admitted budget is clamped to the room the shape leaves. Increasing the capacity grows KV memory linearly and the prefill attention mask quadratically, and produces a distinct compiled artifact that may take longer to compile. StableHLO dynamic sequence shapes are not used. | +| `MLXCEL_XLA_CONTEXT_BUCKETS` | comma-separated integers, each in `1..=2147483647` | derived (`256`, `min(4096, checkpoint window)`, plus the worst-case image floor when larger) | Server OpenXLA continuous-batching worker only (`xla-iree`): the static context shapes it compiles, sorted and deduplicated, for an operator who knows the traffic better than the derivation. Requests route to the smallest bucket that admits them, so short requests keep the 256 shape. Ignored when `MLXCEL_XLA_CONTEXT_CAPACITY` is set, which yields a single bucket. An empty list, a `0`, or a non-integer aborts worker startup. The CLI path does not read it. | | `MLXCEL_XLA_PRECISION` | `f32`, `f16`, `bf16` | `f16` on `cuda`/`hip`, `f32` elsewhere | Contraction precision the StableHLO emitter uses for matmuls (norms and softmax stay in f32). Read at graph-emit time. An explicit value forces that precision even on a GPU device whose default would differ; an unset value falls back to the per-device default (`f16` on `cuda` and `hip`, `f32` on `metal` and `local-task`), and an unrecognized non-empty value is rejected at runtime with a `[compile]` diagnostic. The committed byte-exact goldens are the `f32` graphs, so a byte-exactness check rejects a non-default precision. | | `MLXCEL_XLA_QUANT` | `packed` | unset | `packed` keeps quantized weights packed inside the graph (device-side dequant) instead of dequantizing at load. Read both at emit time and by the loader so uploaded buffers match the emitted args. A no-op on unquantized checkpoints. Supported on `cuda`, `hip`, and `local-task` (not supported on `metal`, where packed int8 dequant prefill faults on the Metal HAL driver); note that on `hip` (`gfx1151`), `iree-compile` lowers `dequant -> matmul` as two separate dispatches so host-dequantized `f16` achieves higher decode throughput while `packed` reduces host upload bytes by ~3.5x. | | `MLXCEL_XLA_IREE_COMPILE` | path to `iree-compile` | baked at build time | Runtime override for the `iree-compile` binary used to lower the bundled graphs to vmfbs. The CUDA/ROCm source-runtime builds rely on a matching `iree-compile` (set via `IREE_CUDA_COMPILE` / `IREE_ROCM_COMPILE` at build time or `MLXCEL_XLA_IREE_COMPILE` at runtime); the CPU/Vulkan dist build falls back to the dist's own `bin/iree-compile`. | diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 54bd63307..6c6d48d3d 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -155,6 +155,32 @@ pub trait ComputeBackend { token_bias: TokenBiasMap, ) -> Result; + /// Produce a single-sequence [`Session`] sized for a known request. + /// + /// `required_positions` is the prompt length plus the generation budget, + /// including any multimodal expansion the caller can predict. Backends with + /// dynamic shapes ignore it and fall through to [`Self::create_session`]; + /// the OpenXLA backend compiles a static graph, so it uses the hint to pick + /// the smallest shape that admits the request instead of a fixed default + /// (issue #2855). The default body keeps every existing implementation and + /// caller unchanged. + /// + /// # Errors + /// + /// Returns an error if the backend cannot construct a session, including + /// when the request cannot fit any shape the backend may compile. + fn create_session_for_request( + &self, + model_path: &Path, + num_layers: usize, + kv_cache_mode: KVCacheMode, + token_bias: TokenBiasMap, + required_positions: Option, + ) -> Result { + let _ = required_positions; + self.create_session(model_path, num_layers, kv_cache_mode, token_bias) + } + /// Whether this backend supports cross-sequence batched serving (the server /// `BatchScheduler` path). The single-sequence [`Session`] never batches; /// this advertises the backend-level capability the extended load-boundary @@ -298,6 +324,46 @@ impl Backend { } } + /// Produce a single-sequence [`Session`] sized for a known request. + /// + /// See [`ComputeBackend::create_session_for_request`]; the dispatch folds + /// to the MLX backend under default features, where the hint is ignored. + #[inline] + pub fn create_session_for_request( + &self, + model_path: &Path, + num_layers: usize, + kv_cache_mode: KVCacheMode, + token_bias: TokenBiasMap, + required_positions: Option, + ) -> Result { + match self { + Backend::Mlx(b) => b.create_session_for_request( + model_path, + num_layers, + kv_cache_mode, + token_bias, + required_positions, + ), + #[cfg(feature = "experimental-backend")] + Backend::Experimental(b) => b.create_session_for_request( + model_path, + num_layers, + kv_cache_mode, + token_bias, + required_positions, + ), + #[cfg(feature = "xla-backend")] + Backend::Xla(b) => b.create_session_for_request( + model_path, + num_layers, + kv_cache_mode, + token_bias, + required_positions, + ), + } + } + /// Whether the active backend supports cross-sequence batched serving. #[inline] #[must_use] diff --git a/src/backend/xla.rs b/src/backend/xla.rs index 572ff5ba7..1f4257708 100644 --- a/src/backend/xla.rs +++ b/src/backend/xla.rs @@ -67,6 +67,28 @@ fn load_unsupported() -> Result { ) } +/// Build the OpenXLA session for a resolved static capacity. +/// +/// The engine drives generation from the session, so it loads its own weights +/// and config from the model directory here (rather than the MLX `load_model` +/// boundary). Under the `iree` feature this compiles the bundled prefill / +/// decode_step graphs and uploads the weights resident; without it, the session +/// loads but `prefill` / `decode_step` report that the `iree` feature is off. +/// KV mode and token bias do not apply: the session owns its own KV and samples +/// greedily on-device. +fn build_xla_session( + model_path: &Path, + num_layers: usize, + context_capacity: usize, +) -> Result { + let image_preprocessor = crate::load_xla_image_preprocessor(model_path) + .map_err(|error| anyhow::anyhow!("OpenXLA image preprocessor load failed: {error}"))?; + let session = + XlaInferenceSession::load_with_context_capacity(model_path, num_layers, context_capacity) + .map_err(|e| anyhow::anyhow!(e))?; + Ok(Session::xla(session, image_preprocessor)) +} + impl ComputeBackend for XlaBackend { fn name(&self) -> &'static str { "xla" @@ -100,18 +122,11 @@ impl ComputeBackend for XlaBackend { _kv_cache_mode: KVCacheMode, _token_bias: TokenBiasMap, ) -> Result { - // The OpenXLA engine drives generation from the session, so it loads its - // own weights and config from the model directory here (rather than the - // MLX `load_model` boundary). Under the `iree` feature this compiles the - // bundled prefill / decode_step graphs and uploads the weights resident; - // without it, the session loads but `prefill` / `decode_step` report that - // the `iree` feature is off. KV mode and token bias do not apply: the - // session owns its own KV and samples greedily on-device. - let image_preprocessor = crate::load_xla_image_preprocessor(model_path) - .map_err(|error| anyhow::anyhow!("OpenXLA image preprocessor load failed: {error}"))?; // Match the server path: reject a graph that could never admit one // image before generation starts, instead of after the vision tower has - // already run on the first image request. + // already run on the first image request. The request-sized path below + // includes the image expansion in its hint, so the guard is only needed + // here, where the request is unknown. let context_capacity = mlxcel_xla::context_capacity_from_env().map_err(|e| anyhow::anyhow!(e))?; crate::ensure_xla_image_context_capacity( @@ -119,13 +134,30 @@ impl ComputeBackend for XlaBackend { context_capacity, std::env::var_os(mlxcel_xla::CONTEXT_CAPACITY_ENV).is_some(), )?; - let session = XlaInferenceSession::load_with_context_capacity( - model_path, - num_layers, - context_capacity, - ) - .map_err(|e| anyhow::anyhow!(e))?; - Ok(Session::xla(session, image_preprocessor)) + build_xla_session(model_path, num_layers, context_capacity) + } + + fn create_session_for_request( + &self, + model_path: &Path, + num_layers: usize, + kv_cache_mode: KVCacheMode, + token_bias: TokenBiasMap, + required_positions: Option, + ) -> Result { + // The static shape has to be chosen before the graph compiles, so this + // is the only place that can size it for the request instead of the + // 256-token compatibility default (issue #2855). A pinned + // `MLXCEL_XLA_CONTEXT_CAPACITY` still wins inside the resolver and + // ignores the hint. The image guard from `create_session` is omitted on + // purpose: the caller's hint already includes the worst-case image + // expansion, so a too-small shape can no longer be compiled here. + let model_window = crate::read_model_context_window(model_path); + let context_capacity = + mlxcel_xla::resolve_context_capacity(required_positions, model_window) + .map_err(|error| anyhow::anyhow!(error))?; + let _ = (kv_cache_mode, token_bias); + build_xla_session(model_path, num_layers, context_capacity) } fn supports_batched_serving(&self) -> bool { diff --git a/src/commands/generate.rs b/src/commands/generate.rs index a391cd196..e0ccaa720 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1588,7 +1588,42 @@ fn decode_xla_cli_images(paths: &[std::path::PathBuf]) -> Result usize { + if budget_was_unlimited && max_tokens > mlxcel_xla::DEFAULT_XLA_GENERATION_BUDGET { + mlxcel_xla::DEFAULT_XLA_GENERATION_BUDGET + } else { + max_tokens + } +} + +/// The positions a request needs from an OpenXLA static shape: prompt plus +/// generation budget plus any worst-case image expansion the checkpoint +/// declares. +#[cfg(feature = "xla-backend")] +fn xla_required_positions( + prompt_len: usize, + sizing_budget: usize, + image_expansion: usize, +) -> usize { + prompt_len + .saturating_add(sizing_budget) + .saturating_add(image_expansion) +} + +/// The worst-case image expansion to include in an OpenXLA sizing hint. +#[cfg(feature = "xla-backend")] +fn xla_image_expansion(model_path: &Path, image_count: usize) -> usize { + image_count.saturating_mul(mlxcel::xla_image_context_floor(model_path).unwrap_or(0)) +} + #[cfg(feature = "xla-backend")] +#[allow(clippy::too_many_arguments)] fn generate_xla( model_path: &Path, num_layers: usize, @@ -1599,6 +1634,7 @@ fn generate_xla( audio_path: Option<&Path>, has_video: bool, max_tokens: usize, + sizing_budget: usize, kv_cache_mode: KVCacheMode, token_bias: TokenBiasMap, ) -> Result<(Vec, GenerationStats)> { @@ -1606,11 +1642,49 @@ fn generate_xla( !has_video, "the OpenXLA backend does not support video input yet" ); - let mut session = - select_backend().create_session(model_path, num_layers, kv_cache_mode, token_bias)?; + let required_positions = xla_required_positions( + prompt_tokens.len(), + sizing_budget, + xla_image_expansion(model_path, image_paths.len()), + ); + let mut session = select_backend().create_session_for_request( + model_path, + num_layers, + kv_cache_mode, + token_bias, + Some(required_positions), + )?; let start_time = Instant::now(); let tokens = match &mut session { mlxcel::Session::Xla(s) => { + let capacity = s.context_capacity(); + let requested_budget = max_tokens; + // The prompt the graph actually admits: image placeholders expand + // inside the session, so their worst case is part of the room the + // generation budget may use. + let effective_prompt_len = required_positions.saturating_sub(sizing_budget); + // The shape is static, so an unlimited (or over-wide) budget has to + // be clamped to the room it leaves after the prompt. The work the + // request asks for is unchanged: generation still stops at EOS. + let admitted = requested_budget.min(capacity.saturating_sub(effective_prompt_len)); + if std::env::var_os(mlxcel_xla::CONTEXT_CAPACITY_ENV).is_none() { + println!( + "OpenXLA context capacity: {capacity} for {effective_prompt_len} prompt + up to \ + {admitted} generated tokens (request-sized; MLXCEL_XLA_CONTEXT_CAPACITY pins \ + one shape)" + ); + } + let max_tokens = admitted; + if max_tokens == 0 { + anyhow::bail!( + "{}", + mlxcel_xla::ContextCapacityError { + effective_prompt_len, + max_new_tokens: requested_budget, + context_capacity: capacity, + } + ); + } let eos = s.eos_token_ids().to_vec(); let phi4mm_media = mlxcel::models::get_model_type(model_path)? == mlxcel::models::ModelType::Phi4MMVLM @@ -2980,6 +3054,13 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { )?; let mut prompt_tokens = tokenize_prompt(&tokenizer, &prompt)?; + // The OpenXLA static shape has to be sized before the graph compiles, and + // the unlimited sentinel would otherwise size it against the whole window + // (issue #2855). Capture the sentinel before `resolve_cli_max_tokens` + // rewrites it; the XLA branch uses it to pick a bounded sizing budget. + let budget_was_unlimited = + args.generation.max_tokens == mlxcel::cli::max_tokens::UNLIMITED_MAX_TOKENS; + // llama.cpp parity (issue #476): resolve an unlimited `-n -1` into a // concrete budget (model context window minus the rendered prompt) now, so // every downstream consumer (the memory preflight below, plus the standard / @@ -3136,6 +3217,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { args.generation.audio.as_deref(), !args.generation.video.is_empty(), args.generation.max_tokens, + xla_sizing_budget(args.generation.max_tokens, budget_was_unlimited), kv_cache_mode, token_bias, )?; diff --git a/src/commands/generate_tests.rs b/src/commands/generate_tests.rs index a52382bf5..9762511f2 100644 --- a/src/commands/generate_tests.rs +++ b/src/commands/generate_tests.rs @@ -24,6 +24,8 @@ use super::{ validate_pipeline_parallel_args, validate_prompt_lookup_args, validate_tensor_parallel_args, validate_xla_cli_image_cardinality, validate_xla_output_audio, }; +#[cfg(feature = "xla-backend")] +use super::{xla_required_positions, xla_sizing_budget}; use mlxcel::server::chat_template::{ChatMessage, ChatTemplateProcessor, flatten_template_text}; use mlxcel_core::cache::KVCacheMode; use mlxcel_core::drafter::{DrafterKind, resolve_drafter_kind}; @@ -68,6 +70,28 @@ fn xla_cli_accepts_empty_image_list_for_audio_only_requests() { assert!(decode_xla_cli_images(&[]).unwrap().is_empty()); } +#[cfg(feature = "xla-backend")] +#[test] +fn xla_sizing_budget_bounds_only_the_unlimited_default() { + // The default -n resolves to the whole window, so it sizes on the bounded + // default instead; an explicit budget is kept verbatim. + assert_eq!( + xla_sizing_budget(131_066, true), + mlxcel_xla::DEFAULT_XLA_GENERATION_BUDGET + ); + assert_eq!(xla_sizing_budget(2_000, false), 2_000); + assert_eq!(xla_sizing_budget(100, true), 100); +} + +#[cfg(feature = "xla-backend")] +#[test] +fn xla_required_positions_adds_prompt_budget_and_image_expansion() { + assert_eq!(xla_required_positions(6, 512, 0), 518); + assert_eq!(xla_required_positions(6, 512, 576), 1_094); + assert_eq!(xla_required_positions(0, 0, 0), 0); + assert_eq!(xla_required_positions(usize::MAX, 1, 1), usize::MAX); +} + // issue #166 / #1165: the offline MTP loop is constructed whenever the // RESOLVED kind is MTP — an explicit `--draft-kind mtp`, or an auto-detected // MTP `model_type` (`gemma4_assistant`, `gemma4_unified_assistant`, diff --git a/src/lib/mlxcel-xla/src/context.rs b/src/lib/mlxcel-xla/src/context.rs index 7b8f2d5a3..4631ffa66 100644 --- a/src/lib/mlxcel-xla/src/context.rs +++ b/src/lib/mlxcel-xla/src/context.rs @@ -16,9 +16,41 @@ use std::fmt; -/// Compatibility capacity used when the operator does not select one. +/// Smallest static shape the request-sizing ladder will compile. +/// +/// Used when the operator pins nothing and the caller has no request to size +/// against (the legacy [`context_capacity_from_env`] path), and kept as the +/// floor of [`resolve_context_capacity`] so a short request still takes the +/// cheapest shape. pub const DEFAULT_CONTEXT_CAPACITY: usize = 256; +/// The static shapes a request-sized session may compile, ascending. +/// +/// The ladder bounds the padding a request pays for: the resolver picks the +/// smallest rung at or above the required positions. The top rung is far below +/// the IREE C ABI's signed 32-bit position limit, and a capacity is also the +/// sequence length every decode step attends over, so the ladder deliberately +/// stops well short of the largest checkpoint windows. +pub const CONTEXT_CAPACITY_LADDER: [usize; 8] = [256, 512, 1024, 2048, 4096, 8192, 16384, 32768]; + +/// The larger text shape the serving bucket set derives by default. +/// +/// A text-only checkpoint used to derive `[256]`, so a prompt that fit any +/// other backend was refused at admission. Serving compiles this second shape +/// next to [`DEFAULT_CONTEXT_CAPACITY`] and routes by request size, so short +/// requests keep the cheap shape and ordinary prompts still run. +pub const SERVER_TEXT_CONTEXT_CAPACITY: usize = 4096; + +/// Generation budget the CLI sizes an OpenXLA static shape against when `-n` is +/// the unlimited default (issue #2855). +/// +/// The unlimited sentinel resolves to the whole checkpoint window, which no +/// static graph can compile. Sizing with a bounded budget keeps the default +/// invocation cheap while still leaving room for a normal completion; an +/// explicit `-n N` sizes with `N` instead. Generation still stops at EOS, and +/// the admitted budget is clamped to the compiled shape's room. +pub const DEFAULT_XLA_GENERATION_BUDGET: usize = 512; + /// Environment variable selecting the static StableHLO context shape. pub const CONTEXT_CAPACITY_ENV: &str = "MLXCEL_XLA_CONTEXT_CAPACITY"; @@ -54,6 +86,71 @@ pub fn context_capacity_from_env() -> Result { } } +/// Resolve the static graph capacity for a single request. +/// +/// The precedence is: +/// +/// 1. [`CONTEXT_CAPACITY_ENV`], which still pins exactly one shape and ignores +/// the request, as it always has. +/// 2. No request hint: [`DEFAULT_CONTEXT_CAPACITY`], the historical shape. +/// 3. Otherwise the smallest [`CONTEXT_CAPACITY_LADDER`] rung that admits +/// `required_positions` (prompt plus generation budget), bounded by the +/// checkpoint window and by the top rung. The window check runs first, so a +/// returned capacity always admits the request, and a window below the +/// ladder floor yields the window itself. +/// +/// # Errors +/// +/// Returns an error for a malformed [`CONTEXT_CAPACITY_ENV`] value, and for a +/// request that needs more positions than the checkpoint window or the ladder +/// can provide. Nothing is silently truncated or clamped. +pub fn resolve_context_capacity( + required_positions: Option, + model_window: Option, +) -> Result { + let pinned = read_env_capacity(CONTEXT_CAPACITY_ENV)?; + resolve_context_capacity_with(required_positions, model_window, pinned) +} + +/// The pure form of [`resolve_context_capacity`], split out so the rule is +/// testable without touching the process environment. +pub(crate) fn resolve_context_capacity_with( + required_positions: Option, + model_window: Option, + pinned: Option, +) -> Result { + if let Some(pinned) = pinned { + return Ok(pinned); + } + let Some(required) = required_positions else { + return Ok(DEFAULT_CONTEXT_CAPACITY); + }; + let ladder_top = *CONTEXT_CAPACITY_LADDER + .last() + .unwrap_or(&DEFAULT_CONTEXT_CAPACITY); + let ceiling = model_window.unwrap_or(ladder_top).min(ladder_top); + if required > ceiling { + return Err(match model_window { + Some(window) => format!( + "this request needs {required} positions but the checkpoint context window is \ + {window}; shorten the prompt or lower -n" + ), + None => format!( + "this request needs {required} positions but the OpenXLA context capacity ladder \ + tops out at {ladder_top}; shorten the prompt or lower -n" + ), + }); + } + // Never compile below the floor, but never above the window either. + let target = required.max(DEFAULT_CONTEXT_CAPACITY).min(ceiling); + let rung = CONTEXT_CAPACITY_LADDER + .iter() + .copied() + .find(|&capacity| capacity >= target) + .unwrap_or(target); + Ok(rung.min(ceiling)) +} + /// Validate a capacity supplied through an API rather than the environment. pub(crate) fn validate_context_capacity_value(context_capacity: usize) -> Result { if context_capacity == 0 { @@ -92,19 +189,24 @@ pub const CONTEXT_BUCKETS_ENV: &str = "MLXCEL_XLA_CONTEXT_BUCKETS"; /// wins over everything else rather than becoming a floor or a hint. /// 2. [`CONTEXT_BUCKETS_ENV`], a comma-separated list, for an operator who /// knows their traffic better than any derivation can. -/// 3. Derived from the checkpoint: the text default, plus the worst-case image -/// expansion when the caller knows one. This is the case that makes an -/// image-capable checkpoint usable without the operator having to discover -/// the number themselves. +/// 3. Derived from the defaults: the cheap text floor, a larger serving text +/// shape bounded by the checkpoint window, plus the worst-case image +/// expansion when the caller knows one. This is the case that makes a +/// text-only checkpoint serve ordinary prompts and an image-capable +/// checkpoint usable without the operator having to discover the numbers +/// themselves. /// /// `image_floor` is `None` for a text-only checkpoint or a family whose worst -/// case is not derived, and a floor at or below the text default adds no -/// bucket, since the default already admits it. +/// case is not derived, and a floor at or below the text floor adds no bucket, +/// since the floor already admits it. /// /// # Errors /// /// Returns an error for a malformed or out-of-range value in either variable. -pub fn context_capacity_buckets_from_env(image_floor: Option) -> Result, String> { +pub fn context_capacity_buckets_from_env( + image_floor: Option, + model_window: Option, +) -> Result, String> { if let Some(pinned) = read_env_capacity(CONTEXT_CAPACITY_ENV)? { return Ok(vec![pinned]); } @@ -113,19 +215,38 @@ pub fn context_capacity_buckets_from_env(image_floor: Option) -> Result) -> Vec { - let mut buckets = vec![text_default]; +pub fn derive_context_buckets( + text_default: usize, + model_window: Option, + image_floor: Option, +) -> Vec { + // Never compile a shape wider than the checkpoint can use; when the window + // is below the floor, the window is the only shape left. + let ceiling = match model_window { + Some(window) if window < text_default => window, + Some(window) => SERVER_TEXT_CONTEXT_CAPACITY.min(window).max(text_default), + None => SERVER_TEXT_CONTEXT_CAPACITY.max(text_default), + }; + let mut buckets = vec![text_default.min(ceiling), ceiling]; if let Some(floor) = image_floor && floor > text_default { buckets.push(floor); } + buckets.sort_unstable(); + buckets.dedup(); buckets } @@ -227,12 +348,95 @@ mod tests { } #[test] - fn a_derived_set_adds_an_image_bucket_only_when_it_is_larger() { - assert_eq!(derive_context_buckets(256, None), vec![256]); - assert_eq!(derive_context_buckets(256, Some(1834)), vec![256, 1834]); - // A floor the text default already admits adds nothing to compile. - assert_eq!(derive_context_buckets(256, Some(200)), vec![256]); - assert_eq!(derive_context_buckets(256, Some(256)), vec![256]); + fn a_derived_set_has_a_text_ceiling_and_an_image_bucket_when_larger() { + let wide = Some(131_072); + assert_eq!(derive_context_buckets(256, wide, None), vec![256, 4096]); + assert_eq!( + derive_context_buckets(256, wide, Some(1834)), + vec![256, 1834, 4096] + ); + // A floor the text floor already admits adds nothing to compile. + assert_eq!( + derive_context_buckets(256, wide, Some(200)), + vec![256, 4096] + ); + assert_eq!( + derive_context_buckets(256, wide, Some(256)), + vec![256, 4096] + ); + // A floor above the text ceiling is still served with its own shape. + assert_eq!( + derive_context_buckets(256, wide, Some(8192)), + vec![256, 4096, 8192] + ); + } + + #[test] + fn a_derived_set_never_exceeds_the_checkpoint_window() { + assert_eq!( + derive_context_buckets(256, Some(2048), None), + vec![256, 2048] + ); + // A window below the floor leaves only the window. + assert_eq!(derive_context_buckets(256, Some(100), None), vec![100]); + } + + #[test] + fn the_resolver_takes_the_next_rung_and_keeps_the_floor() { + let window = Some(131_072); + assert_eq!( + resolve_context_capacity_with(Some(1), window, None), + Ok(256) + ); + assert_eq!( + resolve_context_capacity_with(Some(256), window, None), + Ok(256) + ); + assert_eq!( + resolve_context_capacity_with(Some(257), window, None), + Ok(512) + ); + assert_eq!( + resolve_context_capacity_with(Some(554), window, None), + Ok(1024) + ); + assert_eq!( + resolve_context_capacity_with(Some(32_768), window, None), + Ok(32_768) + ); + } + + #[test] + fn the_resolver_never_exceeds_the_model_window() { + // A window below the ladder floor yields the window itself, but only + // when the request fits it. + assert_eq!( + resolve_context_capacity_with(Some(50), Some(100), None), + Ok(100) + ); + let err = resolve_context_capacity_with(Some(400), Some(300), None).unwrap_err(); + assert!(err.contains("needs 400 positions"), "{err}"); + assert!(err.contains("window is 300"), "{err}"); + } + + #[test] + fn the_resolver_rejects_requests_past_the_ladder_top() { + let err = resolve_context_capacity_with(Some(40_000), None, None).unwrap_err(); + assert!(err.contains("needs 40000 positions"), "{err}"); + assert!(err.contains("tops out at 32768"), "{err}"); + } + + #[test] + fn a_pinned_capacity_wins_over_the_request() { + assert_eq!( + resolve_context_capacity_with(Some(600), Some(131_072), Some(2_048)), + Ok(2_048) + ); + assert_eq!( + resolve_context_capacity_with(None, Some(131_072), Some(2_048)), + Ok(2_048) + ); + assert_eq!(resolve_context_capacity_with(None, None, None), Ok(256)); } #[test] diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 787bada3a..64fc68d27 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -181,9 +181,10 @@ pub use batch::{ run_gemma3n_prefix_decode_diagnostic, }; pub use context::{ - CONTEXT_BUCKETS_ENV, CONTEXT_CAPACITY_ENV, ContextCapacityError, DEFAULT_CONTEXT_CAPACITY, + CONTEXT_BUCKETS_ENV, CONTEXT_CAPACITY_ENV, CONTEXT_CAPACITY_LADDER, ContextCapacityError, + DEFAULT_CONTEXT_CAPACITY, DEFAULT_XLA_GENERATION_BUDGET, SERVER_TEXT_CONTEXT_CAPACITY, context_capacity_buckets_from_env, context_capacity_from_env, derive_context_buckets, - validate_request_capacity, + resolve_context_capacity, validate_request_capacity, }; #[cfg(feature = "diagnostics")] pub use diagnostic_flags::{ diff --git a/src/server/model_worker.rs b/src/server/model_worker.rs index 3971e5ef1..7ca04734c 100644 --- a/src/server/model_worker.rs +++ b/src/server/model_worker.rs @@ -1499,16 +1499,20 @@ pub(crate) fn spawn_xla_model_worker( .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); // Capacities are static graph shapes, so the set has to be decided // before any engine is built and cannot be revisited per request. - // The worst-case image expansion is what makes a second, larger - // shape worth compiling, and it is derived from config alone. + // The worst-case image expansion is what makes an extra, larger + // shape worth compiling, and it is derived from config alone; the + // text shapes are bounded by the checkpoint window so no shape is + // compiled that the model itself could never fill (#2855). let image_floor = crate::xla_image_context_floor(&model_path); - let capacities = match mlxcel_xla::context_capacity_buckets_from_env(image_floor) { - Ok(value) => value, - Err(err) => crate::worker_failfast::exit_on_worker_startup_failure( - "model-worker-xla", - &format!("failed to configure the OpenXLA engine: {err}"), - ), - }; + let model_window = crate::read_model_context_window(&model_path); + let capacities = + match mlxcel_xla::context_capacity_buckets_from_env(image_floor, model_window) { + Ok(value) => value, + Err(err) => crate::worker_failfast::exit_on_worker_startup_failure( + "model-worker-xla", + &format!("failed to configure the OpenXLA engine: {err}"), + ), + }; let operator_pinned = std::env::var_os(mlxcel_xla::CONTEXT_CAPACITY_ENV).is_some(); // The guard asks whether ANY shape can admit an image, so it takes // the largest. With a derived set that is the image floor itself, so From f803d6c6432c083732727407e3b98213a70063fa Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 5 Nov 2026 12:45:36 +0900 Subject: [PATCH 2/2] fix(xla): gate the unlimited-budget probe behind the XLA feature The sentinel capture is only read by the xla-backend branch, so default-feature builds saw an unused binding and CI clippy failed under -D warnings. Refs #2855 --- src/commands/generate.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/commands/generate.rs b/src/commands/generate.rs index e0ccaa720..5f1433f72 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -3058,6 +3058,7 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { // the unlimited sentinel would otherwise size it against the whole window // (issue #2855). Capture the sentinel before `resolve_cli_max_tokens` // rewrites it; the XLA branch uses it to pick a bounded sizing budget. + #[cfg(feature = "xla-backend")] let budget_was_unlimited = args.generation.max_tokens == mlxcel::cli::max_tokens::UNLIMITED_MAX_TOKENS;