feat: Bucket the OpenXLA context capacity so text and image requests do not share one graph shape

[inureyes]

Problem / Background

The OpenXLA context capacity is a static StableHLO graph shape chosen once when the engine is built. DEFAULT_CONTEXT_CAPACITY = 256 in src/lib/mlxcel-xla/src/context.rs:20, overridable through MLXCEL_XLA_CONTEXT_CAPACITY (src/lib/mlxcel-xla/src/context.rs:23, defaulted at line 52). It is not a per-request parameter and cannot be revisited once the graph is compiled.

The capacity is also the sequence length every decode step attends over, so it is a direct throughput cost, not just a memory ceiling. Measured on the pinned Molmo2 4B checkpoint at /home/inureyes/models/mlx/molmo2-4b, 64-token greedy generations, two runs each, on GB10 (aarch64, sm_121) with CUDA:

capacity	run 1	run 2	relative to 256
256	3.17 tok/s	3.20 tok/s	1.00x
1024	2.16 tok/s	2.17 tok/s	1.47x slower
2048	1.40 tok/s	1.41 tok/s	2.26x slower

Meanwhile a single Molmo2 image expands to between 424 logical tokens (a square image, 1x1 high-resolution tiling) and 1834 tokens (a tall image, 8x1 tiling) on that checkpoint.

The consequence is that no single default is both safe and cheap. A capacity that admits the worst-case image makes every text-only request on a VLM checkpoint more than 2x slower, and a capacity tuned for text cannot serve images at all.

Current Behavior

PR #916 resolves the ambiguity by refusing to start when the default capacity cannot admit one image, and by telling the operator what to set. See ensure_xla_image_context_capacity and xla_image_context_floor in src/multimodal/host_preprocessor.rs: the guard derives the worst-case per-image token count from config alone, compares it against the configured capacity, and returns a startup error naming both the requirement and MLXCEL_XLA_CONTEXT_CAPACITY. An operator-pinned capacity is treated as a decision and passes through untouched.

That makes the tradeoff explicit and moves the failure from per-request (run the whole vision tower, then fail at admission) to startup. It does not remove the tradeoff: the operator still has to pick one number for the whole process.

Proposed Solution

Compile more than one graph shape and route each request to the smallest bucket that fits, so a text-only request pays the small-graph cost while an image request gets the large graph. With the measured numbers above, a two-bucket setup (256 and 2048) would keep text-only requests at 3.2 tok/s on the same VLM checkpoint that today has to run the whole process at 1.4 tok/s to stay image-capable.

The open design questions are genuinely open, and this issue is where they get settled rather than assumed:

• How many buckets, and how are they chosen? Fixed powers of two, a configured list, or derived from the checkpoint (for example the image floor from xla_image_context_floor plus a small text bucket). More buckets means finer routing but more compile time and more resident graphs.
• Compile time and device memory. Each additional shape is another iree-compile invocation at startup and another resident executable plus its workspace on device. On a machine where the model already dominates memory, a second large graph may not fit at all, so there needs to be a policy for when a bucket is dropped.
• Interaction with continuous batching. The XLA path batches rows together. If two rows in one batch want different buckets, either the batch runs at the larger bucket (losing the win for the small row), or rows are partitioned by bucket into separate batches (losing batch occupancy). Which one, and on what criterion, needs to be decided before implementation.
• Migration as a sequence grows. A request admitted into the 256 bucket that generates past 256 tokens has to either stop, or move its KV state into a larger bucket mid-generation. Whether migration is supported at all, and if so whether the KV cache can be copied across graph shapes without a re-prefill, decides how conservative admission has to be.

Scope

In scope: src/lib/mlxcel-xla/src/context.rs (capacity selection becomes a set rather than a scalar), the engine build path that compiles graphs, the XLA serve worker's admission and routing (src/server/batch/xla_worker_admission.rs), and the startup guard in src/multimodal/host_preprocessor.rs which becomes a check that a fitting bucket exists rather than that the single capacity fits.

Out of scope: changing the graph shapes themselves, paged or chunked KV, and any non-XLA backend.

Implementation Notes

• Reuse: xla_image_context_floor in src/multimodal/host_preprocessor.rs already derives the worst-case image expansion from config alone with no weights loaded. Bucket selection for a VLM checkpoint should consume that number rather than recomputing it. See feat: Derive worst-case image context floors for the remaining qualified OpenXLA vision families #1272 for extending the floor to the remaining vision families.
• Constraints: MLXCEL_XLA_CONTEXT_CAPACITY is an existing operator-facing variable and an existing decision signal (the guard treats a pinned value as intentional). Whatever bucketing lands must keep a pinned single value working as a single-bucket configuration, so existing deployments do not change behavior.
• Edge cases: a request whose prompt exceeds the largest bucket must be rejected at admission with a message naming the largest bucket, not truncated; a checkpoint whose image floor exceeds the largest bucket must fail at startup the way the current guard does; a bucket that fails to compile must degrade to the remaining buckets rather than failing the whole engine, if the remaining set still covers the checkpoint.
• Measurement discipline: GB10 single-run decode is bimodal by roughly plus or minus 25 percent even at pinned clocks, so any before/after claim about bucketing needs at least three repeats per configuration, not one.

Acceptance Criteria

• The engine can hold more than one compiled context shape, and a request is routed to the smallest shape whose capacity admits its prompt plus its generation budget.
• On the pinned Molmo2 4B checkpoint with buckets covering both 256 and the image floor, a text-only 64-token greedy generation measures within noise of the 3.17 to 3.20 tok/s single-bucket-256 baseline (three repeats), while an image request in the same process succeeds.
• The four design questions above (bucket selection, compile and memory cost, batching interaction, growth migration) are each answered in the implementation, with the chosen answer and its rejected alternative recorded in the code or the PR body.
• A pinned MLXCEL_XLA_CONTEXT_CAPACITY still produces exactly one bucket and byte-identical behavior to today.
• Routing is integrated into the server admission path, not only the CLI, and a request that exceeds the largest bucket is rejected at admission with the bucket size in the message.
• Unit tests pin bucket selection for representative prompt lengths, and a test covers the reject-above-largest-bucket path.

Verification

eval "$(bash scripts/iree/setup-cuda.sh --env)"
cargo build --release --features cuda,xla-iree

MLX_ENABLE_TF32=0 MLXCEL_BACKEND=xla MLXCEL_XLA_DEVICE=cuda \
  ./target/release/mlxcel generate -m /home/inureyes/models/mlx/molmo2-4b -p "..." -n 64

cargo test --release --features cuda,xla-iree context_capacity -- --test-threads=1

CUDA test runs must use --test-threads=1. A pass is: text-only throughput at the 256 baseline, an image request served in the same process, and the bucket-selection unit tests green.

Technical Considerations

This is the same shape of problem as the already-tracked idea of bucketed B_max for the OpenXLA batching engine, and the batching questions above apply identically to both: if bucketing lands for context capacity, the slot-compaction and batch-partitioning machinery it needs is most of what B_max bucketing needs too. Consider doing them together, or at least building the routing seam once.

Related: PR #916, #871.