[inureyes]

Summary

Two defects had one cause: the OpenXLA path pinned its static sequence capacity at 256 and never related it to the request.

• The default CLI invocation always failed. -n defaults to -1, resolves to the checkpoint window minus the prompt, and admission refused it: effective_prompt_len=6 + max_new_tokens=131066 > context_capacity=256.
• Long prompts failed too. At -n 1 any prompt of 255 tokens or more was refused while the same request runs on MLX. Serving had the same problem because a text-only checkpoint derived [256] only.

Changes

• mlxcel-xla: add the ladder [256, 512, 1024, 2048, 4096, 8192, 16384, 32768], resolve_context_capacity, and DEFAULT_XLA_GENERATION_BUDGET = 512. The resolver keeps a pinned MLXCEL_XLA_CONTEXT_CAPACITY first, then picks the smallest rung that admits the request, bounded by the checkpoint window.
• CLI: a defaulted create_session_for_request on the backend seam lets only XlaBackend size the static shape. generate_xla passes prompt plus sizing budget plus worst-case image expansion, prints the chosen capacity, and clamps the admitted budget to the room the shape leaves, mirroring the serve-path clamp from fix: harden the OpenXLA continuous-batching engine after #469 #2848.
• Server: the derived bucket set becomes [256, min(4096, checkpoint window)] plus the image floor when larger, so short requests keep the 256 shape and ordinary prompts are admitted. MLXCEL_XLA_CONTEXT_BUCKETS still overrides.
• Docs: both environment-variable rows updated.

Validation

• cargo test -p mlxcel-xla --lib: 382 passed, 0 failed, 19 ignored.
• cargo test --features xla-backend --bin mlxcel -- commands::generate: 113 passed, 0 failed, 1 ignored.
• cargo check --features xla-iree --all-targets: clean.
• cargo clippy -p mlxcel-xla --all-targets -- -D warnings: clean.
• cargo fmt --all -- --check: clean.
• make verify-env-var-docs: OK, 351 names checked.
• Runtime acceptance cases on gfx1151 are being run after the release build and will be reported in a comment.

Fixes #2855

[inureyes]

Runtime validation on gfx1151 (hip, release build at 1b1ced7)

Checkpoint models/llama-3.2-1b-instruct-4bit, MLXCEL_BACKEND=xla, MLXCEL_XLA_DEVICE=hip, no capacity variable set. Every run under scripts/rocm_gpu_guard.sh --hold.

case	capacity line	result
no -n, 6-token prompt	1024 for 6 prompt + up to 1018	generated 1018 tokens at 11.56 tok/s
-n 64, 6-token prompt	256 for 6 prompt + up to 64	generated 64 tokens
553-token prompt, -n 1	1024 for 553 prompt + up to 1	generated 1 token
553-token prompt, no -n	2048 for 553 prompt + up to 1495	generated 1495 tokens at 11.78 tok/s
-n 2000, 6-token prompt	2048 for 6 prompt + up to 2000	capacity check only, stopped after the line
MLXCEL_XLA_CONTEXT_CAPACITY=300, 553-token prompt	pinned, no sizing line	rejected: effective_prompt_len=553 + max_new_tokens=1 > context_capacity=300

Server on the same checkpoint: the startup log shows capacities=[256, 4096] and OpenXLA capacity buckets compiled capacities=[256, 4096] with image_floor=None. /v1/completions served a 415-prompt-token request (8 completion tokens) and a 6-token request, so a prompt that cannot fit 256 is admitted by the derived 4096 bucket. The per-request bucket choice is not logged, so the short request's 256 routing rests on the existing route_capacity tests rather than on this log.

The release binary was built at 1b1ced7. The follow-up commit f803d6c only adds #[cfg(feature = "xla-backend")] to the sentinel binding, which compiles exactly as before when the feature is on.

Refs #2855

[cla-assistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[inureyes]

recheck