diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 4a62fb8b..e9b6ab7d 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -95,6 +95,7 @@ for how to build the backend. |----------|--------|---------|-------| | `MLXCEL_BACKEND` | `mlx`, `xla` | `mlx` | Runtime compute-backend selector read by `select_backend()`. `xla` routes `load_model` / `forward` through the OpenXLA/IREE engine; it takes effect only when the binary was built with `xla-backend` (else it is ignored and MLX runs). Any other value selects MLX. | | `MLXCEL_XLA_DEVICE` | `metal`, `cuda`, `local-task` | `metal` on macOS, `local-task` elsewhere | IREE HAL device for the XLA engine. CUDA is never auto-selected (it needs a CUDA-enabled runtime build), so set `cuda` explicitly on a GB10-class host. `local-task` is the CPU device. | +| `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_PRECISION` | `f32`, `f16`, `bf16` | `f32` | 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 or unrecognized value falls back to the per-device default. 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. Not supported on the Metal target (packed int8 dequant prefill faults on the Metal HAL driver); use the CUDA or `local-task` target, or leave it unset to dequant at load. | | `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 source-runtime build ships no compiler, so this must point at a cuda-capable `iree-compile` matching the runtime version; the CPU/Vulkan dist build falls back to the dist's own `bin/iree-compile`. | diff --git a/src/lib/mlxcel-xla/README.md b/src/lib/mlxcel-xla/README.md index 8909a006..b6142050 100644 --- a/src/lib/mlxcel-xla/README.md +++ b/src/lib/mlxcel-xla/README.md @@ -156,7 +156,7 @@ correctness-first lever; it does not close the gap to MLX (see the perf note abo - The bundled graphs are authored for **Llama-3.2-1B-Instruct** specifically; `load` verifies `config.json` matches and errors otherwise. -- Prompt length is capped at the prefill bucket (`MAX_SEQ = 256` tokens). +- The context capacity is a static graph shape selected with `MLXCEL_XLA_CONTEXT_CAPACITY` (compatibility default `256`). - Greedy sampling only; text-only (no VLM / draft). - `XlaInferenceSession` is single-sequence; `XlaBatchEngine` (below) adds multi-sequence throughput. @@ -175,6 +175,14 @@ correctness-first lever; it does not close the gap to MLX (see the perf note abo ## Batched continuous batching (Stage 2b) +### Static context capacity + +Set `MLXCEL_XLA_CONTEXT_CAPACITY` before loading a model to select the sequence dimension shared by the prefill graph, decode graph, RoPE tables, masks, and single/ragged KV caches. For example, `MLXCEL_XLA_CONTEXT_CAPACITY=1024` provides enough static space for a 729-token image expansion plus text and generation headroom. The selected value is part of the compiled-artifact identity; modules and runtime buffers with a different capacity are rejected instead of being paired. + +Admission uses `effective_prompt_len + max_new_tokens <= context_capacity`. Text requests use their token count. A multimodal preprocessor must supply the length after placeholder expansion to the same validation helper before native execution. The error reports the effective prompt length, requested generation budget, and configured capacity, and batch rejection occurs before a request id or slot is consumed. + +Larger capacities trade flexibility for resources: KV memory grows linearly with capacity, while the static prefill attention mask and score tensors grow quadratically and compilation can take longer. This backend emits a separate static StableHLO graph for each capacity; it does not claim or rely on dynamic StableHLO sequence shapes. + `XlaBatchEngine` runs many sequences at once: `B_max` slots share one rank-5 KV cache and serve a request stream, so the device stays full. Requests of different lengths join and leave the batch at different times; a freed slot is recycled by diff --git a/src/lib/mlxcel-xla/csrc/xla_iree.c b/src/lib/mlxcel-xla/csrc/xla_iree.c index dfe106d6..4ce8dba6 100644 --- a/src/lib/mlxcel-xla/csrc/xla_iree.c +++ b/src/lib/mlxcel-xla/csrc/xla_iree.c @@ -91,6 +91,7 @@ typedef struct xla_ctx { iree_runtime_session_t* session; // holds both @prefill and @decode_step iree_hal_allocator_t* allocator; // device allocator (shared by both calls) int32_t n_weights; + int32_t context_capacity; iree_hal_buffer_view_t** weights; // resident weights, uploaded once iree_hal_buffer_view_t* kcache; // threaded KV (set by prefill, advanced by decode) iree_hal_buffer_view_t* vcache; @@ -186,12 +187,49 @@ static iree_status_t xla_alloc_bv(xla_ctx* c, iree_host_size_t rank, iree_make_const_byte_span(data, nbytes), out); } +// Validate the static KV contract shared by the paired prefill/decode modules. +// `context_axis` is 1 for rank-4 [L,S,nkv,d] and 2 for rank-5 [B,L,S,nkv,d]. +static int xla_validate_kv_pair(xla_ctx* c, iree_hal_buffer_view_t* kc, + iree_hal_buffer_view_t* vc, + iree_host_size_t rank, + iree_host_size_t context_axis) { + if (!kc || !vc || iree_hal_buffer_view_shape_rank(kc) != rank || + iree_hal_buffer_view_shape_rank(vc) != rank) { + fprintf(stderr, "xla_iree: expected matching rank-%zu K/V buffers\n", + (size_t)rank); + return 1; + } + for (iree_host_size_t i = 0; i < rank; ++i) { + iree_hal_dim_t kd = iree_hal_buffer_view_shape_dim(kc, i); + iree_hal_dim_t vd = iree_hal_buffer_view_shape_dim(vc, i); + if (kd != vd) { + fprintf(stderr, "xla_iree: K/V shape mismatch at axis %zu (%lld vs %lld)\n", + (size_t)i, (long long)kd, (long long)vd); + return 1; + } + } + iree_hal_dim_t actual = + iree_hal_buffer_view_shape_dim(kc, context_axis); + if (actual != (iree_hal_dim_t)c->context_capacity) { + fprintf(stderr, + "xla_iree: module KV context dimension %lld disagrees with configured context_capacity=%d\n", + (long long)actual, c->context_capacity); + return 1; + } + return 0; +} + static iree_status_t xla_llama_create_impl( xla_ctx* c, const char* device_uri, const char* prefill_vmfb, const char* decode_vmfb, int32_t n_weights, const void* const* weight_data, const int32_t* weight_dtypes, const int32_t* weight_ranks, - const int64_t* weight_dims) { + const int64_t* weight_dims, int32_t context_capacity) { + if (context_capacity <= 0) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "context_capacity must be positive"); + } c->n_weights = n_weights; + c->context_capacity = context_capacity; #ifdef XLA_GATE_CUDA // Register the CUDA driver so use_all_available_drivers exposes it for a @@ -279,12 +317,14 @@ xla_ctx* xla_llama_create(const char* device_uri, const char* prefill_vmfb, const void* const* weight_data, const int32_t* weight_dtypes, const int32_t* weight_ranks, - const int64_t* weight_dims) { + const int64_t* weight_dims, + int32_t context_capacity) { xla_ctx* c = (xla_ctx*)calloc(1, sizeof(xla_ctx)); if (!c) return NULL; iree_status_t s = xla_llama_create_impl(c, device_uri, prefill_vmfb, decode_vmfb, n_weights, - weight_data, weight_dtypes, weight_ranks, weight_dims); + weight_data, weight_dtypes, weight_ranks, weight_dims, + context_capacity); if (!iree_status_is_ok(s)) { char buf[512]; iree_host_size_t got = 0; @@ -303,6 +343,12 @@ xla_ctx* xla_llama_create(const char* device_uri, const char* prefill_vmfb, int xla_llama_prefill(xla_ctx* c, const int32_t* tokens, int32_t lp, const int32_t* positions, int32_t real_len, int32_t* out_token) { + if (lp != c->context_capacity || real_len < 0 || real_len > lp) { + fprintf(stderr, + "xla_llama_prefill: lp=%d real_len=%d context_capacity=%d\n", lp, + real_len, c->context_capacity); + return 1; + } iree_runtime_call_t call; XLA_CHECK(iree_runtime_call_initialize_by_name( c->session, iree_make_cstring_view("prefill.main"), &call)); @@ -333,6 +379,16 @@ int xla_llama_prefill(xla_ctx* c, const int32_t* tokens, int32_t lp, XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &token_out)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &kc)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &vc)); + if (xla_validate_kv_pair(c, kc, vc, 4, 1) != 0) { + iree_hal_buffer_view_release(token_out); + iree_hal_buffer_view_release(kc); + iree_hal_buffer_view_release(vc); + iree_hal_buffer_view_release(tok_bv); + iree_hal_buffer_view_release(pos_bv); + iree_hal_buffer_view_release(len_bv); + iree_runtime_call_deinitialize(&call); + return 1; + } if (c->kcache) iree_hal_buffer_view_release(c->kcache); if (c->vcache) iree_hal_buffer_view_release(c->vcache); c->kcache = kc; @@ -352,6 +408,13 @@ int xla_llama_prefill(xla_ctx* c, const int32_t* tokens, int32_t lp, // the resident KV cache in place. int xla_llama_decode(xla_ctx* c, int32_t token, int32_t pos, int32_t cache_len, int32_t* out_token) { + if (!c->kcache || !c->vcache || pos != cache_len || pos < 0 || + pos >= c->context_capacity) { + fprintf(stderr, + "xla_llama_decode: pos=%d cache_len=%d context_capacity=%d or KV is uninitialized\n", + pos, cache_len, c->context_capacity); + return 1; + } iree_runtime_call_t call; XLA_CHECK(iree_runtime_call_initialize_by_name( c->session, iree_make_cstring_view("decode_step.main"), &call)); @@ -382,6 +445,16 @@ int xla_llama_decode(xla_ctx* c, int32_t token, int32_t pos, int32_t cache_len, XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &token_out)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &kc)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &vc)); + if (xla_validate_kv_pair(c, kc, vc, 4, 1) != 0) { + iree_hal_buffer_view_release(token_out); + iree_hal_buffer_view_release(kc); + iree_hal_buffer_view_release(vc); + iree_hal_buffer_view_release(tok_bv); + iree_hal_buffer_view_release(pos_bv); + iree_hal_buffer_view_release(len_bv); + iree_runtime_call_deinitialize(&call); + return 1; + } iree_hal_buffer_view_t* old_k = c->kcache; iree_hal_buffer_view_t* old_v = c->vcache; c->kcache = kc; @@ -438,6 +511,10 @@ static iree_status_t xla_ensure_batch_kv(xla_ctx* c) { // per-slot dims (re-captured by the next prefill_slot, which lazily reallocates // the rank-5 KV once the dims are known). Call once before seeding slots. int xla_llama_ragged_reset(xla_ctx* c, int32_t bsz) { + if (bsz <= 0) { + fprintf(stderr, "xla_llama_ragged_reset: bsz must be positive\n"); + return 1; + } c->rg_bsz = bsz; c->rg_per = 0; if (c->kcache_b) { @@ -466,6 +543,12 @@ int xla_llama_prefill_slot_logits(xla_ctx* c, int32_t slot, const int32_t* token slot, c->rg_bsz); return 1; } + if (lp != c->context_capacity || real_len <= 0 || real_len > lp) { + fprintf(stderr, + "xla_llama_prefill_slot_logits: lp=%d real_len=%d context_capacity=%d\n", + lp, real_len, c->context_capacity); + return 1; + } iree_runtime_call_t call; XLA_CHECK(iree_runtime_call_initialize_by_name( c->session, iree_make_cstring_view("prefill.main"), &call)); @@ -496,6 +579,16 @@ int xla_llama_prefill_slot_logits(xla_ctx* c, int32_t slot, const int32_t* token XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &logits)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &kc)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &vc)); + if (xla_validate_kv_pair(c, kc, vc, 4, 1) != 0) { + iree_hal_buffer_view_release(logits); + iree_hal_buffer_view_release(kc); + iree_hal_buffer_view_release(vc); + iree_hal_buffer_view_release(tok_bv); + iree_hal_buffer_view_release(pos_bv); + iree_hal_buffer_view_release(len_bv); + iree_runtime_call_deinitialize(&call); + return 1; + } if (c->kcache) iree_hal_buffer_view_release(c->kcache); if (c->vcache) iree_hal_buffer_view_release(c->vcache); c->kcache = kc; @@ -514,19 +607,25 @@ int xla_llama_prefill_slot_logits(xla_ctx* c, int32_t slot, const int32_t* token return code ? code : 1; } - if (iree_hal_buffer_view_shape_rank(c->kcache) != 4) { - fprintf(stderr, "xla_llama_prefill_slot_logits: expected rank-4 single-seq KV\n"); - return 1; - } iree_host_size_t per = iree_hal_buffer_view_element_count(c->kcache); if (c->rg_per == 0) { c->rg_per = per; for (int32_t i = 0; i < 4; i++) { c->rg_dims[i] = iree_hal_buffer_view_shape_dim(c->kcache, i); } - } else if (per != c->rg_per) { - fprintf(stderr, "xla_llama_prefill_slot_logits: KV element count changed\n"); - return 1; + } else { + if (per != c->rg_per) { + fprintf(stderr, "xla_llama_prefill_slot_logits: KV element count changed\n"); + return 1; + } + for (int32_t i = 0; i < 4; i++) { + if (c->rg_dims[i] != iree_hal_buffer_view_shape_dim(c->kcache, i)) { + fprintf(stderr, + "xla_llama_prefill_slot_logits: KV shape changed at axis %d\n", + i); + return 1; + } + } } XLA_CHECK(xla_ensure_batch_kv(c)); iree_device_size_t off = (iree_device_size_t)slot * per * sizeof(float); @@ -553,6 +652,21 @@ int xla_llama_prefill_slot_logits(xla_ctx* c, int32_t slot, const int32_t* token int xla_llama_decode_ragged_logits(xla_ctx* c, int32_t bsz, const int32_t* tokens, const int32_t* pos, const int32_t* cache_len, int32_t vocab, float* out_logits) { + if (bsz != c->rg_bsz || !c->kcache_b || !c->vcache_b) { + fprintf(stderr, + "xla_llama_decode_ragged_logits: bsz=%d configured=%d or KV is uninitialized\n", + bsz, c->rg_bsz); + return 1; + } + for (int32_t i = 0; i < bsz; ++i) { + if (pos[i] != cache_len[i] || pos[i] < 0 || + pos[i] >= c->context_capacity) { + fprintf(stderr, + "xla_llama_decode_ragged_logits: row=%d pos=%d cache_len=%d context_capacity=%d\n", + i, pos[i], cache_len[i], c->context_capacity); + return 1; + } + } iree_runtime_call_t call; XLA_CHECK(iree_runtime_call_initialize_by_name( c->session, iree_make_cstring_view("decode_step.main"), &call)); @@ -585,6 +699,18 @@ int xla_llama_decode_ragged_logits(xla_ctx* c, int32_t bsz, const int32_t* token XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &logits_out)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &kc)); XLA_CHECK(iree_runtime_call_outputs_pop_front_buffer_view(&call, &vc)); + if (xla_validate_kv_pair(c, kc, vc, 5, 2) != 0 || + iree_hal_buffer_view_shape_dim(kc, 0) != (iree_hal_dim_t)c->rg_bsz) { + fprintf(stderr, "xla_llama_decode_ragged_logits: output batch shape mismatch\n"); + iree_hal_buffer_view_release(logits_out); + iree_hal_buffer_view_release(kc); + iree_hal_buffer_view_release(vc); + iree_hal_buffer_view_release(tok_bv); + iree_hal_buffer_view_release(pos_bv); + iree_hal_buffer_view_release(len_bv); + iree_runtime_call_deinitialize(&call); + return 1; + } iree_hal_buffer_view_t* old_k = c->kcache_b; iree_hal_buffer_view_t* old_v = c->vcache_b; c->kcache_b = kc; diff --git a/src/lib/mlxcel-xla/src/batch.rs b/src/lib/mlxcel-xla/src/batch.rs index e9b01a53..33a4d23c 100644 --- a/src/lib/mlxcel-xla/src/batch.rs +++ b/src/lib/mlxcel-xla/src/batch.rs @@ -36,15 +36,43 @@ //! tests cannot link the IREE runtime; see the `iree.rs` test note). use std::collections::VecDeque; +use std::fmt; #[cfg(feature = "iree")] use std::path::Path; #[cfg(feature = "iree")] -use crate::iree::{IreeLlama, IreeRaggedLlama, PREFILL_LP}; +use crate::iree::{IreeLlama, IreeRaggedLlama}; use crate::sampler::SampleParams; #[cfg(feature = "iree")] use crate::sampler::sample; +use crate::{ContextCapacityError, validate_request_capacity}; + +/// Typed validation failure returned before a request enters the scheduler. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum XlaAdmissionError { + EmptyPrompt, + ZeroMaxNewTokens, + ContextCapacity(ContextCapacityError), +} + +impl fmt::Display for XlaAdmissionError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::EmptyPrompt => f.write_str("XLA batched submit requires a non-empty prompt"), + Self::ZeroMaxNewTokens => f.write_str("max_new_tokens must be >= 1"), + Self::ContextCapacity(err) => err.fmt(f), + } + } +} + +impl std::error::Error for XlaAdmissionError {} + +impl From for XlaAdmissionError { + fn from(value: ContextCapacityError) -> Self { + Self::ContextCapacity(value) + } +} /// Why a request stopped generating. Cancellation is silent (the caller that /// called [`XlaBatchEngine::cancel`] already knows), so it is not a finish reason. @@ -202,6 +230,25 @@ impl Scheduler { } } +/// Validate and queue atomically: every failure returns before scheduler state is +/// changed, which keeps both free and active slots untouched. +fn queue_request( + sched: &mut Scheduler, + prompt: &[i32], + max_new_tokens: usize, + params: SampleParams, + context_capacity: usize, +) -> Result { + if prompt.is_empty() { + return Err(XlaAdmissionError::EmptyPrompt); + } + if max_new_tokens == 0 { + return Err(XlaAdmissionError::ZeroMaxNewTokens); + } + validate_request_capacity(prompt.len(), max_new_tokens, context_capacity)?; + Ok(sched.submit(prompt.to_vec(), max_new_tokens, params)) +} + /// The continuous-batching engine: `B_max` slots over one ragged decode graph, /// fed by a FIFO queue. See the module docs. #[cfg(feature = "iree")] @@ -222,7 +269,19 @@ impl XlaBatchEngine { /// Propagates load/compile failures, or an unsupported `b_max` (must be one /// of the bundled ragged graphs). pub fn load(model_path: &Path, b_max: usize, device: &str) -> Result { - let engine = IreeRaggedLlama::load(model_path, device, b_max)?; + let context_capacity = crate::context_capacity_from_env()?; + Self::load_with_context_capacity(model_path, b_max, device, context_capacity) + } + + /// Load with an explicitly selected static graph and KV-cache capacity. + pub fn load_with_context_capacity( + model_path: &Path, + b_max: usize, + device: &str, + context_capacity: usize, + ) -> Result { + let context_capacity = crate::context::validate_context_capacity_value(context_capacity)?; + let engine = IreeRaggedLlama::load(model_path, device, b_max, context_capacity)?; let eos = crate::read_eos(model_path); Ok(Self { engine, @@ -236,6 +295,12 @@ impl XlaBatchEngine { self.engine.b_max() } + /// Static sequence capacity compiled into the graph and per-slot KV cache. + #[must_use] + pub fn context_capacity(&self) -> usize { + self.engine.context_capacity() + } + /// The model's EOS token ids (from `generation_config.json`). #[must_use] pub fn eos_token_ids(&self) -> &[i32] { @@ -274,20 +339,15 @@ impl XlaBatchEngine { prompt: &[i32], max_new_tokens: usize, params: SampleParams, - ) -> Result { - if prompt.is_empty() { - return Err("XLA batched submit requires a non-empty prompt".to_string()); - } - if prompt.len() > PREFILL_LP { - return Err(format!( - "prompt of {} exceeds the {PREFILL_LP}-token prefill bucket", - prompt.len() - )); - } - if max_new_tokens == 0 { - return Err("max_new_tokens must be >= 1".to_string()); - } - Ok(self.sched.submit(prompt.to_vec(), max_new_tokens, params)) + ) -> Result { + let context_capacity = self.context_capacity(); + queue_request( + &mut self.sched, + prompt, + max_new_tokens, + params, + context_capacity, + ) } /// Cancel a request by id (frees its slot or drops it from the queue). @@ -424,8 +484,18 @@ impl XlaReferenceEngine { /// /// Propagates the underlying load/compile failures. pub fn load(model_path: &Path, device: &str) -> Result { + let context_capacity = crate::context_capacity_from_env()?; + Self::load_with_context_capacity(model_path, device, context_capacity) + } + + /// Load a reference engine at the same explicit capacity as a batch engine. + pub fn load_with_context_capacity( + model_path: &Path, + device: &str, + context_capacity: usize, + ) -> Result { Ok(Self { - inner: IreeLlama::load(model_path, device)?, + inner: IreeLlama::load(model_path, device, context_capacity)?, }) } @@ -443,6 +513,14 @@ impl XlaReferenceEngine { max_new_tokens: usize, eos: &[i32], ) -> Result, String> { + if prompt.is_empty() { + return Err("XLA reference generation requires a non-empty prompt".to_string()); + } + validate_request_capacity(prompt.len(), max_new_tokens, self.inner.context_capacity()) + .map_err(|err| err.to_string())?; + if max_new_tokens == 0 { + return Ok(Vec::new()); + } let first = self.inner.prefill_first(prompt)?; let mut out = vec![first]; let mut cache_len = prompt.len() as i32; @@ -491,6 +569,50 @@ mod tests { assert!(!s.any_active()); } + #[test] + fn request_admission_checks_prompt_plus_generation_budget() { + let mut s = Scheduler::new(2, EOS.to_vec()); + let id = queue_request(&mut s, &[1; 768], 256, g(), 1024).expect("exact fit"); + assert_eq!(id, 0); + + let err = queue_request(&mut s, &[1; 769], 256, g(), 1024).unwrap_err(); + assert_eq!( + err, + XlaAdmissionError::ContextCapacity(ContextCapacityError { + effective_prompt_len: 769, + max_new_tokens: 256, + context_capacity: 1024, + }) + ); + } + + #[test] + fn rejected_request_does_not_mutate_queue_or_live_slot() { + let mut s = Scheduler::new(2, EOS.to_vec()); + s.slots[0] = Some(Slot { + req_id: 41, + cur: 7, + cache_len: 300, + produced: 4, + cap: 100, + params: g(), + rng: 9, + history: vec![1, 2, 3], + }); + let before_next_id = s.next_id; + let before_queue_len = s.queue.len(); + + let err = queue_request(&mut s, &[5; 900], 125, g(), 1024).unwrap_err(); + assert!(matches!(err, XlaAdmissionError::ContextCapacity(_))); + assert_eq!(s.next_id, before_next_id, "no request id was consumed"); + assert_eq!(s.queue.len(), before_queue_len, "nothing was queued"); + let live = s.slots[0].as_ref().expect("existing slot remains active"); + assert_eq!(live.req_id, 41); + assert_eq!(live.cur, 7); + assert_eq!(live.cache_len, 300); + assert!(s.slots[1].is_none(), "a free slot remains free"); + } + #[test] fn pop_next_pending_skips_cancelled() { let mut s = Scheduler::new(2, EOS.to_vec()); diff --git a/src/lib/mlxcel-xla/src/context.rs b/src/lib/mlxcel-xla/src/context.rs new file mode 100644 index 00000000..7d412761 --- /dev/null +++ b/src/lib/mlxcel-xla/src/context.rs @@ -0,0 +1,142 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Static context-capacity configuration and request admission for OpenXLA. + +use std::fmt; + +/// Compatibility capacity used when the operator does not select one. +pub const DEFAULT_CONTEXT_CAPACITY: usize = 256; + +/// Environment variable selecting the static StableHLO context shape. +pub const CONTEXT_CAPACITY_ENV: &str = "MLXCEL_XLA_CONTEXT_CAPACITY"; + +/// A request whose effective prompt plus generation budget cannot fit the graph. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ContextCapacityError { + pub effective_prompt_len: usize, + pub max_new_tokens: usize, + pub context_capacity: usize, +} + +impl fmt::Display for ContextCapacityError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "request exceeds the OpenXLA context capacity: effective_prompt_len={} + max_new_tokens={} > context_capacity={}", + self.effective_prompt_len, self.max_new_tokens, self.context_capacity + ) + } +} + +impl std::error::Error for ContextCapacityError {} + +/// Resolve the static graph capacity selected by the operator. +/// +/// The value must fit the IREE C ABI's signed 32-bit position/count arguments. +/// An unset variable keeps the historical 256-token graph shape. +pub fn context_capacity_from_env() -> Result { + match std::env::var(CONTEXT_CAPACITY_ENV) { + Ok(raw) => parse_context_capacity(&raw), + Err(std::env::VarError::NotPresent) => Ok(DEFAULT_CONTEXT_CAPACITY), + Err(err) => Err(format!("read {CONTEXT_CAPACITY_ENV}: {err}")), + } +} + +/// 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 { + return Err("OpenXLA context capacity must be at least 1 token".to_string()); + } + if context_capacity > i32::MAX as usize { + return Err(format!( + "OpenXLA context capacity {context_capacity} exceeds the IREE ABI maximum {}", + i32::MAX + )); + } + Ok(context_capacity) +} + +fn parse_context_capacity(raw: &str) -> Result { + let value = raw.parse::().map_err(|_| { + format!( + "{CONTEXT_CAPACITY_ENV} must be an integer in 1..={}, got {raw:?}", + i32::MAX + ) + })?; + validate_context_capacity_value(value) +} + +/// Enforce the common text/multimodal admission invariant. +/// +/// `effective_prompt_len` is the token count after any placeholder expansion. +/// The checked addition also rejects an overflowing generation budget. +pub fn validate_request_capacity( + effective_prompt_len: usize, + max_new_tokens: usize, + context_capacity: usize, +) -> Result<(), ContextCapacityError> { + if context_capacity > 0 + && effective_prompt_len + .checked_add(max_new_tokens) + .is_some_and(|needed| needed <= context_capacity) + { + Ok(()) + } else { + Err(ContextCapacityError { + effective_prompt_len, + max_new_tokens, + context_capacity, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exact_fit_is_admitted() { + assert_eq!(validate_request_capacity(768, 256, 1024), Ok(())); + } + + #[test] + fn one_token_overflow_reports_all_values() { + let err = validate_request_capacity(769, 256, 1024).unwrap_err(); + assert_eq!(err.effective_prompt_len, 769); + assert_eq!(err.max_new_tokens, 256); + assert_eq!(err.context_capacity, 1024); + assert_eq!( + err.to_string(), + "request exceeds the OpenXLA context capacity: effective_prompt_len=769 + max_new_tokens=256 > context_capacity=1024" + ); + } + + #[test] + fn expanded_multimodal_length_uses_the_same_invariant() { + assert!(validate_request_capacity(729 + 32, 128, 1024).is_ok()); + assert!(validate_request_capacity(729 + 200, 128, 1024).is_err()); + } + + #[test] + fn overflowing_generation_budget_is_rejected() { + let err = validate_request_capacity(1, usize::MAX, usize::MAX).unwrap_err(); + assert_eq!(err.max_new_tokens, usize::MAX); + } + + #[test] + fn zero_capacity_is_never_admitted() { + assert!(validate_request_capacity(0, 0, 0).is_err()); + } +} diff --git a/src/lib/mlxcel-xla/src/emitter/config.rs b/src/lib/mlxcel-xla/src/emitter/config.rs index b737a2b7..b2273239 100644 --- a/src/lib/mlxcel-xla/src/emitter/config.rs +++ b/src/lib/mlxcel-xla/src/emitter/config.rs @@ -160,6 +160,9 @@ pub struct MoeConfig { #[derive(Clone, Debug)] pub struct Config { + /// Static sequence dimension shared by prefill, decode, and every KV cache. + /// StableHLO shapes remain static; changing this value emits a distinct graph. + pub context_capacity: usize, pub hidden: usize, pub inter: usize, pub n_layers: usize, @@ -307,6 +310,7 @@ impl Config { /// Hard-coded Llama-3.2-1B-Instruct values (config.json of the spike model). pub fn llama_3_2_1b() -> Self { Config { + context_capacity: crate::DEFAULT_CONTEXT_CAPACITY, hidden: 2048, inter: 8192, n_layers: 16, @@ -1016,6 +1020,7 @@ impl Config { let tie_word_embeddings = ob("tie_word_embeddings").unwrap_or(tie_default); Ok(Config { + context_capacity: crate::DEFAULT_CONTEXT_CAPACITY, hidden, inter: u("intermediate_size")?, n_layers, @@ -1069,6 +1074,12 @@ impl Config { Self::from_json_str(&s).map_err(|e| format!("{}: {e}", p.display())) } + /// Select the static sequence shape used by all emitted graphs and caches. + pub fn with_context_capacity(mut self, context_capacity: usize) -> Result { + self.context_capacity = crate::context::validate_context_capacity_value(context_capacity)?; + Ok(self) + } + pub fn group(&self) -> usize { self.n_q / self.n_kv } diff --git a/src/lib/mlxcel-xla/src/emitter/context_tests.rs b/src/lib/mlxcel-xla/src/emitter/context_tests.rs new file mode 100644 index 00000000..7206f9e1 --- /dev/null +++ b/src/lib/mlxcel-xla/src/emitter/context_tests.rs @@ -0,0 +1,85 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{Config, emit_decode, emit_decode_ragged, emit_prefill}; + +fn tiny_config(context_capacity: usize) -> Config { + Config { + context_capacity, + hidden: 8, + inter: 16, + n_layers: 2, + n_q: 2, + n_kv: 1, + head_dim: 4, + vocab: 16, + ..Config::llama_3_2_1b() + } +} + +fn assert_capacity_shapes(context_capacity: usize) { + let cfg = tiny_config(context_capacity); + let prefill = emit_prefill(&cfg, false); + let decode = emit_decode(&cfg, false); + let ragged = emit_decode_ragged(&cfg, 4, false); + let single_cache = format!("tensor<2x{context_capacity}x1x4xf32>"); + let batch_cache = format!("tensor<4x2x{context_capacity}x1x4xf32>"); + let prompt = format!("tensor<{context_capacity}xi32>"); + let causal_mask = format!("tensor<{context_capacity}x{context_capacity}xf32>"); + let ragged_mask = format!("tensor<4x{context_capacity}xf32>"); + let rope = format!("tensor<{context_capacity}x4xf32>"); + + assert!( + prefill.matches(&prompt).count() >= 2, + "tokens and positions use {prompt}" + ); + assert!( + prefill.matches(&single_cache).count() >= 3, + "prefill returns matching K/V" + ); + assert!( + prefill.contains(&causal_mask), + "prefill mask uses {causal_mask}" + ); + assert!( + prefill.contains(&rope), + "prefill position table uses {rope}" + ); + + assert!( + decode.matches(&single_cache).count() >= 5, + "decode input/output K/V agree" + ); + assert!(decode.contains(&rope), "decode position table uses {rope}"); + + assert!( + ragged.matches(&batch_cache).count() >= 5, + "ragged input/output K/V agree" + ); + assert!( + ragged.contains(&ragged_mask), + "ragged mask uses {ragged_mask}" + ); + assert!(ragged.contains(&rope), "ragged position table uses {rope}"); +} + +#[test] +fn graph_schema_is_consistent_at_compatibility_capacity() { + assert_capacity_shapes(256); +} + +#[test] +fn graph_schema_is_consistent_at_multimodal_capacity() { + assert_capacity_shapes(1024); +} diff --git a/src/lib/mlxcel-xla/src/emitter/mod.rs b/src/lib/mlxcel-xla/src/emitter/mod.rs index e3e7848b..3283913c 100644 --- a/src/lib/mlxcel-xla/src/emitter/mod.rs +++ b/src/lib/mlxcel-xla/src/emitter/mod.rs @@ -51,6 +51,9 @@ mod model; mod moe; mod rope; +#[cfg(test)] +mod context_tests; + pub(crate) use config::{Config, NormStyle, QkNorm, WeightScheme}; // MoE FFN config types (issue #500), read by the weight loader (`iree.rs`) and the // validation harness. `SharedExpertConfig` is only named in some build cfgs. @@ -1235,7 +1238,7 @@ mod tests { /// Gemma3 pairs the Gemma2 four-norm layer and a per-head `(1+w)` q/k norm with a /// DUAL RoPE base: the sliding layers rotate on a local-base table distinct from - /// the global one, so the graph carries two extra `[MAX_SEQ, head_dim]` constant + /// the global one, so the graph carries two extra `[context_capacity, head_dim]` constant /// tables (cos_local, sin_local) versus a single-RoPE twin. #[test] fn gemma3_dual_rope_and_qk_norm_reach_the_shared_core() { diff --git a/src/lib/mlxcel-xla/src/emitter/model.rs b/src/lib/mlxcel-xla/src/emitter/model.rs index 7631f4a7..07fd9bd1 100644 --- a/src/lib/mlxcel-xla/src/emitter/model.rs +++ b/src/lib/mlxcel-xla/src/emitter/model.rs @@ -24,15 +24,6 @@ use super::config::{Config, NormStyle}; use super::moe::{self, MoeLayerW, MoeSharedW}; use super::rope; -const MAX_SEQ: usize = 256; - -/// Bucketed padded prompt length for prefill. Set to MAX_SEQ so the one bucket -/// covers any prompt the cache holds (e.g. the ~103-token Llama-3.2 chat-template -/// prompt), not just the 46-token spike prompt. KV for real positions is -/// identical to a smaller bucket (padding positions are causally masked), so the -/// generated tokens are unchanged. -const PREFILL_LP: usize = MAX_SEQ; - /// Per-layer weight handles (JAX alphabetical order: down, gate, in_ln, /// post_ln, up, wk, wo, wq, wv). `bk`/`bq`/`bv` are the q/k/v projection biases, /// present only for an architecture with `qkv_bias` (Qwen2); `None` for Llama, @@ -473,13 +464,13 @@ fn build_arg_schema(b: &mut Builder, c: &Config) -> (Vec, Args) { let kcache = take_arg( &mut decls, &mut idx, - Ty::f32(vec![c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]), + Ty::f32(vec![c.n_layers, c.context_capacity, c.n_kv, c.head_dim]), "kcache".into(), ); let vcache = take_arg( &mut decls, &mut idx, - Ty::f32(vec![c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]), + Ty::f32(vec![c.n_layers, c.context_capacity, c.n_kv, c.head_dim]), "vcache".into(), ); @@ -533,20 +524,21 @@ pub(crate) struct Consts { } fn emit_consts(b: &mut Builder, c: &Config) -> Consts { - let (cos, sin) = rope::rope_tables(c, MAX_SEQ); + let seq = c.context_capacity; + let (cos, sin) = rope::rope_tables(c, seq); // The rotary width (`head_dim`, or the smaller `rotary_dim` for a partial-RoPE // arch like StableLM). The Llama family has `rot == head_dim`, so its tables are // byte-identical. let rot = c.rotary_width(); - let cos_table = b.const_tensor_f32(&cos, vec![MAX_SEQ, rot]); - let sin_table = b.const_tensor_f32(&sin, vec![MAX_SEQ, rot]); + let cos_table = b.const_tensor_f32(&cos, vec![seq, rot]); + let sin_table = b.const_tensor_f32(&sin, vec![seq, rot]); // Gemma3 / OLMo3 local RoPE tables (distinct base for the sliding layers). let (cos_local, sin_local) = match c.rope_local_base { Some(base) => { - let (cl, sl) = rope::rope_tables_local(c, MAX_SEQ, base); + let (cl, sl) = rope::rope_tables_local(c, seq, base); ( - Some(b.const_tensor_f32(&cl, vec![MAX_SEQ, rot])), - Some(b.const_tensor_f32(&sl, vec![MAX_SEQ, rot])), + Some(b.const_tensor_f32(&cl, vec![seq, rot])), + Some(b.const_tensor_f32(&sl, vec![seq, rot])), ) } None => (None, None), @@ -1282,6 +1274,7 @@ impl AttnLayout { ) -> (Val, Val) { let d = c.head_dim; let nkv = c.n_kv; + let seq = c.context_capacity; match self { AttnLayout::Single { cache_len, .. } => { let k_upd = b.reshape(kk, vec![1, 1, nkv, d]); @@ -1296,10 +1289,10 @@ impl AttnLayout { &v_upd, &[&k.layer_idx[li], cache_len, &k.c0, &k.c0], ); - let kl = b.slice(&*kcache, &[(li, li + 1), (0, MAX_SEQ), (0, nkv), (0, d)]); - let kl = b.reshape(&kl, vec![MAX_SEQ, nkv, d]); - let vl = b.slice(&*vcache, &[(li, li + 1), (0, MAX_SEQ), (0, nkv), (0, d)]); - let vl = b.reshape(&vl, vec![MAX_SEQ, nkv, d]); + let kl = b.slice(&*kcache, &[(li, li + 1), (0, seq), (0, nkv), (0, d)]); + let kl = b.reshape(&kl, vec![seq, nkv, d]); + let vl = b.slice(&*vcache, &[(li, li + 1), (0, seq), (0, nkv), (0, d)]); + let vl = b.reshape(&vl, vec![seq, nkv, d]); (kl, vl) } AttnLayout::Ragged { @@ -1329,14 +1322,14 @@ impl AttnLayout { } let kl = b.slice( &*kcache, - &[(0, *bsz), (li, li + 1), (0, MAX_SEQ), (0, nkv), (0, d)], + &[(0, *bsz), (li, li + 1), (0, seq), (0, nkv), (0, d)], ); - let kl = b.reshape(&kl, vec![*bsz, MAX_SEQ, nkv, d]); + let kl = b.reshape(&kl, vec![*bsz, seq, nkv, d]); let vl = b.slice( &*vcache, - &[(0, *bsz), (li, li + 1), (0, MAX_SEQ), (0, nkv), (0, d)], + &[(0, *bsz), (li, li + 1), (0, seq), (0, nkv), (0, d)], ); - let vl = b.reshape(&vl, vec![*bsz, MAX_SEQ, nkv, d]); + let vl = b.reshape(&vl, vec![*bsz, seq, nkv, d]); (kl, vl) } AttnLayout::Prefill { lp, .. } => { @@ -1362,12 +1355,12 @@ impl AttnLayout { fn raw_scores(&self, b: &mut Builder, c: &Config, q: &Val, kslab: &Val) -> Val { let d = c.head_dim; let (nq, nkv, g) = (c.n_q, c.n_kv, c.group()); + let seq = c.context_capacity; match self { AttnLayout::Single { .. } => { let q_r = b.reshape(q, vec![nkv, g, d]); - let scores = - b.dot_general(&q_r, kslab, &[0], &[1], &[2], &[2], vec![nkv, g, MAX_SEQ]); - b.reshape(&scores, vec![nq, MAX_SEQ]) + let scores = b.dot_general(&q_r, kslab, &[0], &[1], &[2], &[2], vec![nkv, g, seq]); + b.reshape(&scores, vec![nq, seq]) } AttnLayout::Ragged { bsz, .. } => { let q_r = b.reshape(q, vec![*bsz, nkv, g, d]); @@ -1378,9 +1371,9 @@ impl AttnLayout { &[0, 2], &[3], &[3], - vec![*bsz, nkv, g, MAX_SEQ], + vec![*bsz, nkv, g, seq], ); - b.reshape(&scores, vec![*bsz, nq, MAX_SEQ]) + b.reshape(&scores, vec![*bsz, nq, seq]) } AttnLayout::Prefill { lp, .. } => { let q4 = b.reshape(q, vec![*lp, nkv, g, d]); @@ -1396,12 +1389,13 @@ impl AttnLayout { /// See [`layer_mask`]. fn add_mask(&self, b: &mut Builder, c: &Config, li: usize, scores: &Val) -> Val { let (nq, nkv, g) = (c.n_q, c.n_kv, c.group()); + let seq = c.context_capacity; match self { AttnLayout::Single { mask, mask_local, .. } => { let m = layer_mask(mask, mask_local, c, li); - let mb = b.broadcast(m, &[1], vec![nq, MAX_SEQ]); + let mb = b.broadcast(m, &[1], vec![nq, seq]); b.add(scores, &mb) } AttnLayout::Ragged { @@ -1411,7 +1405,7 @@ impl AttnLayout { .. } => { let m = layer_mask(mask, mask_local, c, li); - let mb = b.broadcast(m, &[0, 2], vec![*bsz, nq, MAX_SEQ]); + let mb = b.broadcast(m, &[0, 2], vec![*bsz, nq, seq]); b.add(scores, &mb) } AttnLayout::Prefill { @@ -1440,15 +1434,16 @@ impl AttnLayout { fn context(&self, b: &mut Builder, c: &Config, attn: &Val, vslab: &Val) -> Val { let d = c.head_dim; let (nq, nkv, g) = (c.n_q, c.n_kv, c.group()); + let seq = c.context_capacity; match self { AttnLayout::Single { .. } => { - let attn_r = b.reshape(attn, vec![nkv, g, MAX_SEQ]); + let attn_r = b.reshape(attn, vec![nkv, g, seq]); let o = b.dot_general(&attn_r, vslab, &[0], &[1], &[2], &[0], vec![nkv, g, d]); let o = b.reshape(&o, vec![nq, d]); b.reshape(&o, vec![nq * d]) } AttnLayout::Ragged { bsz, .. } => { - let attn_r = b.reshape(attn, vec![*bsz, nkv, g, MAX_SEQ]); + let attn_r = b.reshape(attn, vec![*bsz, nkv, g, seq]); let o = b.dot_general( &attn_r, vslab, @@ -1673,6 +1668,7 @@ pub fn emit_decode_with(c: &Config, sample: bool, precision: Precision) -> Strin let k = emit_consts(&mut b, c); let h = c.hidden; + let seq = c.context_capacity; // RoPE cos/sin vectors span the rotary width (`head_dim` for a full-rope arch, // the smaller `rotary_dim` for partial RoPE; equal for the Llama family). let rot = c.rotary_width(); @@ -1699,23 +1695,23 @@ pub fn emit_decode_with(c: &Config, sample: bool, precision: Precision) -> Strin }; // mask: keys s valid iff s <= cache_len -> additive 0 / -1e30, shape [S] - let ii = b.iota(MAX_SEQ); - let clen_b = b.broadcast(&a.cache_len, &[], vec![MAX_SEQ]); + let ii = b.iota(seq); + let clen_b = b.broadcast(&a.cache_len, &[], vec![seq]); let valid = b.compare("LE", &ii, &clen_b, "SIGNED"); - let zeros_s = b.broadcast(&k.zero, &[], vec![MAX_SEQ]); - let negs_s = b.broadcast(&k.neg_big, &[], vec![MAX_SEQ]); + let zeros_s = b.broadcast(&k.zero, &[], vec![seq]); + let negs_s = b.broadcast(&k.neg_big, &[], vec![seq]); let kmask = b.select(&valid, &zeros_s, &negs_s); // Gemma2 local (sliding-window) mask (issue #495): a local layer additionally // drops keys older than the window, keeping key s iff `cache_len - s < W`. The // global `kmask` already encodes causality, so anding the window in is one more // `select`: within-window -> keep `kmask`, else force -1e30. Built once and // reused by every local layer; emitted only for a config with a window - // (Gemma2), so Llama / Qwen2 stay byte-identical. When `W >= MAX_SEQ` the - // predicate is always true (`cache_len - s <= MAX_SEQ-1 < W`), so it is a value + // (Gemma2), so Llama / Qwen2 stay byte-identical. When `W >= context_capacity` + // the predicate is always true, so it is a value // no-op, which is why short-context Gemma2 output is unchanged. let kmask_local = c.sliding_window.map(|w| { let wc = b.const_i32(w as i32); - let wb = b.broadcast(&wc, &[], vec![MAX_SEQ]); + let wb = b.broadcast(&wc, &[], vec![seq]); let age = b.subtract(&clen_b, &ii); // cache_len - s let within = b.compare("LT", &age, &wb, "SIGNED"); b.select(&within, &kmask, &negs_s) @@ -1756,7 +1752,7 @@ pub fn emit_decode_with(c: &Config, sample: bool, precision: Precision) -> Strin }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![c.n_layers, seq, c.n_kv, c.head_dim]).render(); format!( "module @decode_step {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, @@ -1792,7 +1788,7 @@ struct BatchedArgs { token: Val, // [B] i32 pos: Val, // scalar i32 (shared across the batch) cache_len: Val, // scalar i32 (shared across the batch) - kcache: Val, // [B, L, MAX_SEQ, nkv, d] + kcache: Val, // [B, L, context_capacity, nkv, d] vcache: Val, } @@ -1838,13 +1834,25 @@ fn build_batched_arg_schema( let kcache = take_arg( &mut decls, &mut idx, - Ty::f32(vec![bsz, c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]), + Ty::f32(vec![ + bsz, + c.n_layers, + c.context_capacity, + c.n_kv, + c.head_dim, + ]), "kcache".into(), ); let vcache = take_arg( &mut decls, &mut idx, - Ty::f32(vec![bsz, c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]), + Ty::f32(vec![ + bsz, + c.n_layers, + c.context_capacity, + c.n_kv, + c.head_dim, + ]), "vcache".into(), ); @@ -1910,6 +1918,7 @@ pub fn emit_decode_batched_with( let nq = c.n_q; let nkv = c.n_kv; let g = c.group(); + let seq = c.context_capacity; // --- head: per-row embed gather, shared rope vectors, shared key mask --- let tok_idx = b.reshape(&a.token, vec![bsz, 1]); @@ -1922,11 +1931,11 @@ pub fn emit_decode_batched_with( let sin_vec = b.reshape(&sin_row, vec![d]); // shared key mask [S]: key s valid iff s <= cache_len -> additive 0 / -1e30 - let ii = b.iota(MAX_SEQ); - let clen_b = b.broadcast(&a.cache_len, &[], vec![MAX_SEQ]); + let ii = b.iota(seq); + let clen_b = b.broadcast(&a.cache_len, &[], vec![seq]); let valid = b.compare("LE", &ii, &clen_b, "SIGNED"); - let zeros_s = b.broadcast(&k.zero, &[], vec![MAX_SEQ]); - let negs_s = b.broadcast(&k.neg_big, &[], vec![MAX_SEQ]); + let zeros_s = b.broadcast(&k.zero, &[], vec![seq]); + let negs_s = b.broadcast(&k.neg_big, &[], vec![seq]); let kmask = b.select(&valid, &zeros_s, &negs_s); let mut kcache = a.kcache.clone(); @@ -1973,14 +1982,14 @@ pub fn emit_decode_batched_with( // read this layer's cache slabs [B, S, nkv, d] let kl = b.slice( &kcache, - &[(0, bsz), (li, li + 1), (0, MAX_SEQ), (0, nkv), (0, d)], + &[(0, bsz), (li, li + 1), (0, seq), (0, nkv), (0, d)], ); - let kl = b.reshape(&kl, vec![bsz, MAX_SEQ, nkv, d]); + let kl = b.reshape(&kl, vec![bsz, seq, nkv, d]); let vl = b.slice( &vcache, - &[(0, bsz), (li, li + 1), (0, MAX_SEQ), (0, nkv), (0, d)], + &[(0, bsz), (li, li + 1), (0, seq), (0, nkv), (0, d)], ); - let vl = b.reshape(&vl, vec![bsz, MAX_SEQ, nkv, d]); + let vl = b.reshape(&vl, vec![bsz, seq, nkv, d]); // GQA scores: batch over (B, kv head). q head kv*g+grp attends kv head // kv. Output [B, nkv, g, S] reshapes to [B, nq, S] (head = kv*g+grp). @@ -1992,25 +2001,25 @@ pub fn emit_decode_batched_with( &[0, 2], &[3], &[3], - vec![bsz, nkv, g, MAX_SEQ], + vec![bsz, nkv, g, seq], ); - let scores = b.reshape(&scores, vec![bsz, nq, MAX_SEQ]); - let scale_b = b.broadcast(&k.scale, &[], vec![bsz, nq, MAX_SEQ]); + let scores = b.reshape(&scores, vec![bsz, nq, seq]); + let scale_b = b.broadcast(&k.scale, &[], vec![bsz, nq, seq]); let scores = b.multiply(&scores, &scale_b); - let kmask_b = b.broadcast(&kmask, &[2], vec![bsz, nq, MAX_SEQ]); + let kmask_b = b.broadcast(&kmask, &[2], vec![bsz, nq, seq]); let scores = b.add(&scores, &kmask_b); // softmax over the key axis (dim 2) let m = b.reduce_max(&scores, 2, &k.neg_inf); // [B, nq] - let m_b = b.broadcast(&m, &[0, 1], vec![bsz, nq, MAX_SEQ]); + let m_b = b.broadcast(&m, &[0, 1], vec![bsz, nq, seq]); let sh = b.subtract(&scores, &m_b); let e = b.exponential(&sh); let s = b.reduce_add(&e, 2, &k.zero); // [B, nq] - let s_b = b.broadcast(&s, &[0, 1], vec![bsz, nq, MAX_SEQ]); + let s_b = b.broadcast(&s, &[0, 1], vec![bsz, nq, seq]); let attn = b.divide(&e, &s_b); // [B, nq, S] // context: o[b,h,d] = sum_s attn[b,h,s] * vl[b,s,h/g,d] - let attn_r = b.reshape(&attn, vec![bsz, nkv, g, MAX_SEQ]); + let attn_r = b.reshape(&attn, vec![bsz, nkv, g, seq]); let o = b.dot_general( &attn_r, &vl, @@ -2079,7 +2088,7 @@ pub fn emit_decode_batched_with( }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![bsz, c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![bsz, c.n_layers, seq, c.n_kv, c.head_dim]).render(); format!( "module @decode_step {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, @@ -2115,7 +2124,7 @@ struct RaggedArgs { token: Val, // [B] i32 pos: Val, // [B] i32 (per row) cache_len: Val, // [B] i32 (per row) - kcache: Val, // [B, L, MAX_SEQ, nkv, d] + kcache: Val, // [B, L, context_capacity, nkv, d] vcache: Val, } @@ -2167,13 +2176,25 @@ fn build_ragged_arg_schema(b: &mut Builder, c: &Config, bsz: usize) -> (Vec = (0..bsz).map(|i| b.const_i32(i as i32)).collect(); let h = c.hidden; + let seq = c.context_capacity; // --- head: per-row embed gather, per-row rope gather, per-row key mask --- let tok_idx = b.reshape(&a.token, vec![bsz, 1]); @@ -2231,20 +2253,20 @@ pub fn emit_decode_ragged_with( }; // per-row key mask [B,S]: key s valid for row b iff s <= cache_len[b] - let ii = b.iota(MAX_SEQ); // [S] - let ii_b = b.broadcast(&ii, &[1], vec![bsz, MAX_SEQ]); // entry[b,s] = s - let clen_b = b.broadcast(&a.cache_len, &[0], vec![bsz, MAX_SEQ]); // entry[b,s] = cache_len[b] + let ii = b.iota(seq); // [S] + let ii_b = b.broadcast(&ii, &[1], vec![bsz, seq]); // entry[b,s] = s + let clen_b = b.broadcast(&a.cache_len, &[0], vec![bsz, seq]); // entry[b,s] = cache_len[b] let valid = b.compare("LE", &ii_b, &clen_b, "SIGNED"); - let zeros = b.broadcast(&k.zero, &[], vec![bsz, MAX_SEQ]); - let negs = b.broadcast(&k.neg_big, &[], vec![bsz, MAX_SEQ]); + let zeros = b.broadcast(&k.zero, &[], vec![bsz, seq]); + let negs = b.broadcast(&k.neg_big, &[], vec![bsz, seq]); let kmask = b.select(&valid, &zeros, &negs); // [B, S] // Gemma2 local (sliding-window) mask (issue #495), per row: keep key s for row // b iff `cache_len[b] - s < W`. And it into the causal `kmask` with one more // `select`. Emitted only for a windowed config (Gemma2); reused by every local - // layer. No-op on the value when `W >= MAX_SEQ` (short-context parity). + // layer. No-op on the value when `W >= context_capacity` (short-context parity). let kmask_local = c.sliding_window.map(|w| { let wc = b.const_i32(w as i32); - let wb = b.broadcast(&wc, &[], vec![bsz, MAX_SEQ]); + let wb = b.broadcast(&wc, &[], vec![bsz, seq]); let age = b.subtract(&clen_b, &ii_b); // cache_len[b] - s let within = b.compare("LT", &age, &wb, "SIGNED"); b.select(&within, &kmask, &negs) @@ -2294,7 +2316,7 @@ pub fn emit_decode_ragged_with( }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![bsz, c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![bsz, c.n_layers, seq, c.n_kv, c.head_dim]).render(); format!( "module @decode_step {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, @@ -2471,7 +2493,7 @@ pub fn emit_prefill(c: &Config, sample: bool) -> String { } pub fn emit_prefill_with(c: &Config, sample: bool, precision: Precision) -> String { - let lp = PREFILL_LP; + let lp = c.context_capacity; let mut b = Builder::new().with_precision(precision); let (decls, a) = build_prefill_arg_schema(&mut b, c, lp); let k = emit_consts(&mut b, c); @@ -2528,8 +2550,8 @@ pub fn emit_prefill_with(c: &Config, sample: bool, precision: Precision) -> Stri }; // caches start as zeros; prefill writes the [0:Lp] block and returns them - let mut kcache = b.broadcast(&k.zero, &[], vec![c.n_layers, MAX_SEQ, nkv, d]); - let mut vcache = b.broadcast(&k.zero, &[], vec![c.n_layers, MAX_SEQ, nkv, d]); + let mut kcache = b.broadcast(&k.zero, &[], vec![c.n_layers, lp, nkv, d]); + let mut vcache = b.broadcast(&k.zero, &[], vec![c.n_layers, lp, nkv, d]); for li in 0..c.n_layers { let lw = &a.layers[li]; @@ -2564,7 +2586,7 @@ pub fn emit_prefill_with(c: &Config, sample: bool, precision: Precision) -> Stri }; let sig = render_signature(&decls); - let cache_ty = Ty::f32(vec![c.n_layers, MAX_SEQ, c.n_kv, c.head_dim]).render(); + let cache_ty = Ty::f32(vec![c.n_layers, lp, c.n_kv, c.head_dim]).render(); format!( "module @prefill {{\n func.func public @main({sig}) -> ({out_ty}, {cache_ty}, {cache_ty}) {{\n{body} return {l}, {kc}, {vc} : {out_ty}, {cache_ty}, {cache_ty}\n }}\n}}\n", sig = sig, diff --git a/src/lib/mlxcel-xla/src/iree.rs b/src/lib/mlxcel-xla/src/iree.rs index e9fd3d31..76d58e69 100644 --- a/src/lib/mlxcel-xla/src/iree.rs +++ b/src/lib/mlxcel-xla/src/iree.rs @@ -100,10 +100,6 @@ impl WeightBuf { } } -/// Prefill bucket baked into the emitted `prefill` graph (`tensor<256xi32>`, -/// == MAX_SEQ, so it covers any prompt the 256-slot KV cache holds). -pub const PREFILL_LP: usize = 256; - /// Opaque handle to the C-side execution context. #[repr(C)] struct XlaCtx { @@ -120,6 +116,7 @@ unsafe extern "C" { weight_dtypes: *const c_int, weight_ranks: *const c_int, weight_dims: *const i64, + context_capacity: c_int, ) -> *mut XlaCtx; fn xla_llama_prefill( c: *mut XlaCtx, @@ -256,6 +253,7 @@ fn compile_one( flags: &[&str], cache: &Path, tag: &str, + context_capacity: usize, ) -> Result { let mut h = std::collections::hash_map::DefaultHasher::new(); mlir.hash(&mut h); @@ -265,9 +263,10 @@ fn compile_one( // Include the compiler path so a cuda vmfb is never reused for a cpu build // (or across iree-compile versions) just because the graph text matches. iree_compile.to_string_lossy().hash(&mut h); + context_capacity.hash(&mut h); let key = h.finish(); - let mlir_path = cache.join(format!("{tag}-{key:016x}.mlir")); - let vmfb_path = cache.join(format!("{tag}-{key:016x}.vmfb")); + let mlir_path = cache.join(format!("{tag}-c{context_capacity}-{key:016x}.mlir")); + let vmfb_path = cache.join(format!("{tag}-c{context_capacity}-{key:016x}.vmfb")); if vmfb_path.exists() { return Ok(vmfb_path); } @@ -297,6 +296,7 @@ fn compile_pair( prefill_tag: &str, decode_mlir: &str, decode_tag: &str, + context_capacity: usize, ) -> Result<(PathBuf, PathBuf), String> { let iree_compile = iree_compile_bin()?; if !iree_compile.exists() { @@ -308,8 +308,22 @@ fn compile_pair( let flags = target_flags(device)?; let cache = std::env::temp_dir().join("mlxcel-xla-vmfb"); std::fs::create_dir_all(&cache).map_err(|e| format!("mkdir {}: {e}", cache.display()))?; - let pre = compile_one(&iree_compile, prefill_mlir, flags, &cache, prefill_tag)?; - let dec = compile_one(&iree_compile, decode_mlir, flags, &cache, decode_tag)?; + let pre = compile_one( + &iree_compile, + prefill_mlir, + flags, + &cache, + prefill_tag, + context_capacity, + )?; + let dec = compile_one( + &iree_compile, + decode_mlir, + flags, + &cache, + decode_tag, + context_capacity, + )?; Ok((pre, dec)) } @@ -322,7 +336,14 @@ fn compile_vmfbs(device: &str, cfg: &Config) -> Result<(PathBuf, PathBuf), Strin check_packed_supported(device, quant_in_graph() && cfg.supports_packed_quant())?; let prefill = emit_prefill_with(cfg, true, precision); let decode = emit_decode_with(cfg, true, precision); - compile_pair(device, &prefill, "prefill", &decode, "decode") + compile_pair( + device, + &prefill, + "prefill", + &decode, + "decode", + cfg.context_capacity, + ) } /// Map each needed weight name to the safetensors file that holds it, in the same @@ -623,6 +644,7 @@ fn create_ctx( dtypes.as_ptr(), ranks.as_ptr(), dims.as_ptr(), + cfg.context_capacity as c_int, ) }; // Weights are resident on the device now; free the host copy. @@ -641,20 +663,30 @@ fn create_ctx( /// (the raw context is single-threaded), matching the single-sequence session. pub struct IreeLlama { ctx: *mut XlaCtx, + context_capacity: usize, } impl IreeLlama { /// Prepare execution for a model directory on a HAL `device` /// (`"local-task"` for CPU). Compiles the bundled graphs, uploads the /// weights resident, and readies the prefill / decode calls. - pub fn load(model_dir: &Path, device: &str) -> Result { - let cfg = Config::from_json(model_dir)?; + pub fn load(model_dir: &Path, device: &str, context_capacity: usize) -> Result { + let cfg = Config::from_json(model_dir)?.with_context_capacity(context_capacity)?; let (prefill_vmfb, decode_vmfb) = compile_vmfbs(device, &cfg)?; let ctx = create_ctx(model_dir, &cfg, device, &prefill_vmfb, &decode_vmfb)?; - Ok(Self { ctx }) + Ok(Self { + ctx, + context_capacity: cfg.context_capacity, + }) } - /// Seed the KV cache with `token_ids` (length <= [`PREFILL_LP`]) via the + /// Static sequence capacity compiled into the paired graphs and KV buffers. + #[must_use] + pub fn context_capacity(&self) -> usize { + self.context_capacity + } + + /// Seed the KV cache with `token_ids` (length <= the configured capacity) via the /// bucketed prefill graph. The returned token (the graph's argmax at /// `real_len-1`) is unused by the seed-then-decode drive loop. pub fn prefill_seed(&mut self, token_ids: &[i32]) -> Result<(), String> { @@ -676,27 +708,27 @@ impl IreeLlama { self.prefill_padded(prompt) } - /// Pad `prompt` into the [`PREFILL_LP`] bucket, run the prefill, and return its + /// Pad `prompt` into the configured static bucket, run the prefill, and return its /// first token. Accepts an empty prompt (the seed-then-decode loop prefills a /// zero-length prefix when the prompt is a single token). fn prefill_padded(&mut self, prompt: &[i32]) -> Result { - if prompt.len() > PREFILL_LP { + if prompt.len() > self.context_capacity { return Err(format!( - "the OpenXLA M2 prefill graph is bucketed at {PREFILL_LP} tokens; prompt \ - prefix of {} exceeds it (a larger-bucket graph is a follow-up)", - prompt.len() + "prompt prefix of {} exceeds context_capacity={}", + prompt.len(), + self.context_capacity )); } - let mut tokens = vec![0i32; PREFILL_LP]; + let mut tokens = vec![0i32; self.context_capacity]; tokens[..prompt.len()].copy_from_slice(prompt); - let positions: Vec = (0..PREFILL_LP as c_int).collect(); + let positions: Vec = (0..self.context_capacity as c_int).collect(); let mut out = 0i32; // Safety: buffers outlive the call; the shim stores the returned KV. let rc = unsafe { xla_llama_prefill( self.ctx, tokens.as_ptr(), - PREFILL_LP as c_int, + self.context_capacity as c_int, positions.as_ptr(), prompt.len() as c_int, &mut out, @@ -711,6 +743,12 @@ impl IreeLlama { /// Advance one token at `cache_len` (== position), returning the next token /// id (on-device argmax) and writing the new K/V into the resident cache. pub fn decode(&mut self, token: i32, cache_len: i32) -> Result { + if cache_len < 0 || cache_len as usize >= self.context_capacity { + return Err(format!( + "decode position {cache_len} is outside context_capacity={}", + self.context_capacity + )); + } let mut out = 0i32; // Safety: the shim threads its own resident KV; only scalars cross here. let rc = unsafe { xla_llama_decode(self.ctx, token, cache_len, cache_len, &mut out) }; @@ -747,6 +785,7 @@ pub struct IreeRaggedLlama { /// Vocabulary size (logits per row), from the model config; the readback /// buffers and the per-row logits slice are sized by it. vocab: usize, + context_capacity: usize, } impl IreeRaggedLlama { @@ -755,8 +794,13 @@ impl IreeRaggedLlama { /// Verifies the architecture, compiles the bundled prefill + the ragged decode /// graph for `b_max`, uploads the weights resident, and sizes the batch. /// `b_max` must be one of the bundled graphs ([`RAGGED_B_VALUES`]). - pub fn load(model_dir: &Path, device: &str, b_max: usize) -> Result { - let cfg = Config::from_json(model_dir)?; + pub fn load( + model_dir: &Path, + device: &str, + b_max: usize, + context_capacity: usize, + ) -> Result { + let cfg = Config::from_json(model_dir)?.with_context_capacity(context_capacity)?; if !RAGGED_B_VALUES.contains(&b_max) { return Err(format!( "the OpenXLA serve worker selects B_max from {RAGGED_B_VALUES:?}; \ @@ -775,6 +819,7 @@ impl IreeRaggedLlama { "prefill_logits", &decode_mlir, &format!("decode_ragged_logits_b{b_max}"), + cfg.context_capacity, )?; let ctx = create_ctx(model_dir, &cfg, device, &prefill_vmfb, &decode_vmfb)?; // Safety: ctx is a fresh valid context from create_ctx; free it on error. @@ -787,6 +832,7 @@ impl IreeRaggedLlama { ctx, b_max, vocab: cfg.vocab, + context_capacity: cfg.context_capacity, }) } @@ -803,7 +849,13 @@ impl IreeRaggedLlama { self.vocab } - /// Seed slot `slot` with `prompt` (1..=[`PREFILL_LP`] tokens) and return its + /// Static sequence capacity compiled into the paired graphs and KV buffers. + #[must_use] + pub fn context_capacity(&self) -> usize { + self.context_capacity + } + + /// Seed slot `slot` with `prompt` (up to the configured capacity) and return its /// first-token `[vocab]` LOGITS (#449 M3 Stage 2d). The prompt's KV is written /// device-side into the slot's region of the rank-5 cache; the other slots are /// untouched, so a mid-stream admit does not disturb live sequences. The caller @@ -815,15 +867,16 @@ impl IreeRaggedLlama { if prompt.is_empty() { return Err("prefill_slot_logits requires a non-empty prompt".to_string()); } - if prompt.len() > PREFILL_LP { + if prompt.len() > self.context_capacity { return Err(format!( - "prompt of {} exceeds the {PREFILL_LP}-token prefill bucket", - prompt.len() + "prompt of {} exceeds context_capacity={}", + prompt.len(), + self.context_capacity )); } - let mut tokens = vec![0i32; PREFILL_LP]; + let mut tokens = vec![0i32; self.context_capacity]; tokens[..prompt.len()].copy_from_slice(prompt); - let positions: Vec = (0..PREFILL_LP as c_int).collect(); + let positions: Vec = (0..self.context_capacity as c_int).collect(); let mut logits = vec![0f32; self.vocab]; // Safety: input buffers outlive the call; `logits` has self.vocab elements, // which the shim fills; the shim also writes the slot's KV device-side. @@ -832,7 +885,7 @@ impl IreeRaggedLlama { self.ctx, slot as c_int, tokens.as_ptr(), - PREFILL_LP as c_int, + self.context_capacity as c_int, positions.as_ptr(), prompt.len() as c_int, self.vocab as c_int, @@ -863,6 +916,17 @@ impl IreeRaggedLlama { self.b_max )); } + for row in 0..self.b_max { + if pos[row] != cache_len[row] + || pos[row] < 0 + || pos[row] as usize >= self.context_capacity + { + return Err(format!( + "decode row {row} has pos={} and cache_len={} outside context_capacity={} or not equal", + pos[row], cache_len[row], self.context_capacity + )); + } + } let mut logits = vec![0f32; self.b_max * self.vocab]; // Safety: the three input slices are length b_max == bsz; `logits` has // b_max*self.vocab elements, which the shim fills; it threads its rank-5 KV. diff --git a/src/lib/mlxcel-xla/src/lib.rs b/src/lib/mlxcel-xla/src/lib.rs index 9a4a43c1..f34a6adb 100644 --- a/src/lib/mlxcel-xla/src/lib.rs +++ b/src/lib/mlxcel-xla/src/lib.rs @@ -49,6 +49,8 @@ use std::path::{Path, PathBuf}; use mlxcel_core::session::{InferenceSession, SessionCapabilities}; +mod context; + #[cfg(feature = "iree")] mod iree; @@ -99,7 +101,11 @@ mod emitter; mod validation; #[cfg(feature = "iree")] -pub use batch::{EngineEvent, FinishReason, XlaBatchEngine, XlaReferenceEngine}; +pub use batch::{EngineEvent, FinishReason, XlaAdmissionError, XlaBatchEngine, XlaReferenceEngine}; +pub use context::{ + CONTEXT_CAPACITY_ENV, ContextCapacityError, DEFAULT_CONTEXT_CAPACITY, + context_capacity_from_env, validate_request_capacity, +}; #[cfg(feature = "iree")] pub use sampler::SampleParams; @@ -149,6 +155,7 @@ fn read_eos(_model_path: &Path) -> Vec { pub struct XlaInferenceSession { model_path: PathBuf, num_layers: usize, + context_capacity: usize, eos_token_ids: Vec, #[cfg(feature = "iree")] engine: iree::IreeLlama, @@ -192,15 +199,27 @@ impl XlaInferenceSession { /// unsupported model architecture, a missing IREE distribution, an /// `iree-compile` failure, or a weight-loading failure. pub fn load(model_path: &Path, num_layers: usize) -> Result { + let context_capacity = context_capacity_from_env()?; + Self::load_with_context_capacity(model_path, num_layers, context_capacity) + } + + /// Prepare a session with an explicitly selected static context capacity. + pub fn load_with_context_capacity( + model_path: &Path, + num_layers: usize, + context_capacity: usize, + ) -> Result { + let context_capacity = context::validate_context_capacity_value(context_capacity)?; let eos_token_ids = read_eos(model_path); #[cfg(feature = "iree")] { let device = std::env::var("MLXCEL_XLA_DEVICE").unwrap_or_else(|_| default_device().to_string()); - let engine = iree::IreeLlama::load(model_path, &device)?; + let engine = iree::IreeLlama::load(model_path, &device, context_capacity)?; Ok(Self { model_path: model_path.to_path_buf(), num_layers, + context_capacity, eos_token_ids, engine, cache_len: 0, @@ -211,6 +230,7 @@ impl XlaInferenceSession { Ok(Self { model_path: model_path.to_path_buf(), num_layers, + context_capacity, eos_token_ids, }) } @@ -228,6 +248,12 @@ impl XlaInferenceSession { &self.model_path } + /// Static sequence capacity compiled into this session's graph and KV cache. + #[must_use] + pub fn context_capacity(&self) -> usize { + self.context_capacity + } + /// The model's EOS token ids (from `generation_config.json`), so the drive /// loop and the seam can stop on them. #[must_use] @@ -274,6 +300,8 @@ impl XlaInferenceSession { if prompt_tokens.is_empty() { return Err("XLA generation requires a non-empty prompt".to_string()); } + validate_request_capacity(prompt_tokens.len(), max_new_tokens, self.context_capacity) + .map_err(|err| err.to_string())?; // Seed KV with all but the last prompt token; the last token is the first // input to decode_step, which returns the first generated token. This // keeps decode_step's "given the previously emitted token" contract exact @@ -302,6 +330,13 @@ impl InferenceSession for XlaInferenceSession { #[cfg(feature = "iree")] fn prefill(&mut self, token_ids: &[i32]) -> Result<(), String> { + if token_ids.len() > self.context_capacity { + return Err(format!( + "prefill length {} exceeds context_capacity={}", + token_ids.len(), + self.context_capacity + )); + } self.engine.prefill_seed(token_ids)?; self.cache_len = token_ids.len() as i32; Ok(()) @@ -314,6 +349,12 @@ impl InferenceSession for XlaInferenceSession { #[cfg(feature = "iree")] fn decode_step(&mut self, token: i32) -> Result { + if self.cache_len < 0 || self.cache_len as usize >= self.context_capacity { + return Err(format!( + "decode position {} is outside context_capacity={}", + self.cache_len, self.context_capacity + )); + } let next = self.engine.decode(token, self.cache_len)?; self.cache_len += 1; Ok(next) @@ -351,6 +392,19 @@ mod tests { let s = session(); assert_eq!(s.num_layers(), 16); assert_eq!(s.model_path(), Path::new("/tmp/xla-model")); + assert_eq!(s.context_capacity(), DEFAULT_CONTEXT_CAPACITY); + } + + #[test] + fn explicit_context_capacity_is_recorded_and_validated() { + let s = + XlaInferenceSession::load_with_context_capacity(Path::new("/tmp/xla-model"), 16, 1024) + .unwrap(); + assert_eq!(s.context_capacity(), 1024); + assert!( + XlaInferenceSession::load_with_context_capacity(Path::new("/tmp/xla-model"), 16, 0,) + .is_err() + ); } #[test] @@ -369,6 +423,17 @@ mod tests { ); } + #[test] + fn greedy_rejects_context_overflow_before_the_unwired_backend() { + let mut s = + XlaInferenceSession::load_with_context_capacity(Path::new("/tmp/xla-model"), 16, 8) + .unwrap(); + let err = s.generate_greedy(&[1, 2, 3, 4, 5], 4, &[2]).unwrap_err(); + assert!(err.contains("effective_prompt_len=5")); + assert!(err.contains("max_new_tokens=4")); + assert!(err.contains("context_capacity=8")); + } + #[test] fn empty_prompt_is_rejected() { let mut s = session(); diff --git a/src/server/model_worker.rs b/src/server/model_worker.rs index c9161b08..b4221443 100644 --- a/src/server/model_worker.rs +++ b/src/server/model_worker.rs @@ -983,9 +983,16 @@ pub(crate) fn spawn_xla_model_worker( run_core_thread_or_abort("model-worker-xla", move || { let device = std::env::var("MLXCEL_XLA_DEVICE") .unwrap_or_else(|_| mlxcel_xla::default_device().to_string()); + let context_capacity = match mlxcel_xla::context_capacity_from_env() { + Ok(value) => value, + Err(err) => { + tracing::error!("Failed to configure the OpenXLA engine: {err}"); + return; + } + }; tracing::info!( "Model worker thread starting (OpenXLA continuous batching, B_max={b_max}, \ - device={device}), loading model..." + context_capacity={context_capacity}, device={device}), loading model..." ); let load_start = Instant::now(); @@ -996,7 +1003,12 @@ pub(crate) fn spawn_xla_model_worker( return; } }; - let engine = match mlxcel_xla::XlaBatchEngine::load(&model_path, b_max, &device) { + let engine = match mlxcel_xla::XlaBatchEngine::load_with_context_capacity( + &model_path, + b_max, + &device, + context_capacity, + ) { Ok(engine) => engine, Err(err) => { tracing::error!("Failed to load the OpenXLA engine: {err}"); @@ -1007,7 +1019,7 @@ pub(crate) fn spawn_xla_model_worker( tracing::info!( worker_model_id = %worker_model_id, load_seconds = load_elapsed.as_secs_f64(), - "OpenXLA model {worker_model_id} loaded in {:.3}s (B_max={b_max}, device={device})", + "OpenXLA model {worker_model_id} loaded in {:.3}s (B_max={b_max}, context_capacity={context_capacity}, device={device})", load_elapsed.as_secs_f64(), ); loaded.store(true, Ordering::Release);