fix(xla): size the static context shape from the request

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=<sku>` 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`. |

src/backend/mod.rs

@@ -155,6 +155,32 @@ pub trait ComputeBackend {
         token_bias: TokenBiasMap,
     ) -> Result<Session>;
 
+    /// 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<usize>,
+    ) -> Result<Session> {
+        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<usize>,
+    ) -> Result<Session> {
+        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]

src/backend/xla.rs

@@ -67,6 +67,28 @@ fn load_unsupported<T>() -> Result<T> {
     )
 }
 
+/// 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<Session> {
+    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,32 +122,42 @@ impl ComputeBackend for XlaBackend {
         _kv_cache_mode: KVCacheMode,
         _token_bias: TokenBiasMap,
     ) -> Result<Session> {
-        // 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(
             model_path,
             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<usize>,
+    ) -> Result<Session> {
+        // 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 {

src/commands/generate.rs

@@ -1588,7 +1588,42 @@ fn decode_xla_cli_images(paths: &[std::path::PathBuf]) -> Result<Vec<image::Dyna
 /// as serving. Gemma3n audio uses the same bounded host front-end and split
 /// IREE prepared-prefill contract as serving; mixed image+audio remains
 /// fail-closed until a joint runtime is qualified.
+/// The budget an OpenXLA session should size its static shape against.
+///
+/// The unlimited `-n` default resolves to the whole checkpoint window, which no
+/// static graph can compile; [`mlxcel_xla::DEFAULT_XLA_GENERATION_BUDGET`]
+/// bounds it. An explicit `-n N` sizes with `N` (issue #2855).
+#[cfg(feature = "xla-backend")]
+fn xla_sizing_budget(max_tokens: usize, budget_was_unlimited: bool) -> 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,18 +1634,57 @@ 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<i32>, GenerationStats)> {
     ensure!(
         !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,14 @@ 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.
+    #[cfg(feature = "xla-backend")]
+    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 +3218,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,
         )?;