Skip to content

feat(xla): parameterize static context capacity#880

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-857-context-capacity
Jul 22, 2026
Merged

feat(xla): parameterize static context capacity#880
inureyes merged 1 commit into
mainfrom
feature/issue-857-context-capacity

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Parameterize the OpenXLA static context capacity across graph emission, compiled artifacts, IREE runtime buffers, single-session generation, continuous batching, and server startup.

Changes

  • Add MLXCEL_XLA_CONTEXT_CAPACITY with a compatibility default of 256 and validate explicit capacities against the signed IREE ABI.
  • Derive prefill tensors, RoPE tables, masks, single/ragged KV caches, slot buffers, and decode bounds from one canonical capacity.
  • Include capacity in compiled MLIR/VMFB artifact names and hashes, and validate returned native KV shapes before accepting them.
  • Centralize text and future post-expansion multimodal admission around effective_prompt_len + max_new_tokens <= context_capacity with typed errors before scheduler mutation.
  • Add 256/1024 emitter schema coverage, admission boundary and overflow coverage, a no-slot-mutation regression, operator documentation, and native Rust/C bounds checks.

Validation

  • cargo fmt --all -- --check
  • cargo clippy -p mlxcel-xla --lib --tests -- -D warnings
  • cargo check -p mlxcel-xla --lib
  • cargo check -p mlxcel --features xla-backend --lib
  • IREE_DIST=/home/inureyes/Development/mlxcel/spike/iree-ffi/iree-dist cargo check -p mlxcel --features xla-iree --lib
  • cargo test -p mlxcel-xla --lib context (9 passed, including 256/1024 graph schemas and admission boundaries)
  • cargo test -p mlxcel-xla --lib admission
  • cargo test -p mlxcel-xla --lib rejected_request_does_not_mutate_queue_or_live_slot
  • cargo test -p mlxcel-xla --lib from_json_reproduces_bundled_assets_byte_for_byte
  • MLXCEL_XLA_CONTEXT_CAPACITY=256 IREE_DIST=/home/inureyes/Development/mlxcel/spike/iree-ffi/iree-dist cargo run --features xla-iree --example xla_oracle_check -- --model /home/inureyes/models/smollm-135m-4bit --oracle /tmp/issue857-smollm-oracle.json --device local-task (Linux aarch64 CPU local-task, 2-token HF fp32 oracle, TOKEN-EXACT PASS)
  • MLXCEL_XLA_CONTEXT_CAPACITY=1024 IREE_DIST=/home/inureyes/Development/mlxcel/spike/iree-ffi/iree-dist cargo run --features xla-iree --example xla_oracle_check -- --model /home/inureyes/models/smollm-135m-4bit --oracle /tmp/issue857-smollm-oracle.json --device local-task (TOKEN-EXACT PASS; distinct -c1024- artifacts from the -c256- pair)

The full crate-wide library suite and validate_arch.sh module-wide structural gate are deferred to the epic integration wave; the issue branch uses the focused tests above to avoid duplicating broad wave validation.

Closes #857

Thread one validated capacity through StableHLO emission, compiled-artifact identity, IREE buffers, single-session generation, and continuous batching so modules and runtime state cannot silently disagree.

Reject prompt-plus-generation overflow with a typed error before scheduler or slot mutation, guard decode positions at the Rust and C boundaries, and document the static-shape memory and compile-time tradeoffs.

Refs #857
@inureyes inureyes added status:review Under review type:enhancement New features, capabilities, or significant additions priority:high High priority area:architecture Architecture and code structure changes area:inference Generation, sampling, decoding (incl. speculative, DRY) status:done Completed and removed status:review Under review labels Jul 22, 2026
@inureyes
inureyes merged commit 25136ca into main Jul 22, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:architecture Architecture and code structure changes area:inference Generation, sampling, decoding (incl. speculative, DRY) priority:high High priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(xla): parameterize context and multimodal prefill capacity

1 participant