onnxruntime-ep-mlx 0.31.1

//! Unified compiled fast path (`mlx_compile`) — ONE parameterised core for every compiled strategy.
//!
//! The eager translator (`engine::TranslationContext::execute`) rebuilds and dispatches EVERY node
//! of a claimed subgraph as separate unfused MLX primitive launches on every Compute call. For a
//! decoder that is ~393 kernels per token; for a CNN/audio graph it is hundreds per inference —
//! dominating runtime. This module traces the WHOLE claimed subgraph translation into an
//! `mlx_closure` over its dynamic (non-constant) ctx inputs ONCE, compiles it with `mlx_compile`
//! (kernel fusion), caches the compiled closure on the plan, and on each subsequent call just
//! applies the compiled closure to the freshly-wrapped inputs.
//!
//! Historically this lived in two modules — a decode-only shapeless path and a general static-shape
//! path. They are now a single [`crate::engine::CompiledSubgraph`] core parameterised by a
//! [`crate::engine::CompiledConfig`] (`shape_mode` + the opt-in `kv_alias` / `rope_as_data` /
//! `delta_copyout` / `contiguous_outputs` features). The strategies are just configurations:
//!   * **decode**  = `{ Shapeless,  kv_alias, rope_as_data, delta_copyout }` — shapeless so a growing
//!     KV length never retraces; RoPE position + valid-past fed as DATA; delta KV copy-out.
//!   * **general** = `{ ShapeKeyed, contiguous_outputs }` — retraces on a shape change; no KV aliasing.
//!   * **prefill** = `{ ShapeKeyed, kv_alias, rope_as_data, delta_copyout }` (Phase 2) — the SAME
//!     decoder subgraph as decode but with a variable query length S>1, so it is shape-keyed.
//!
//! Two shapeless-compile tricks carried by `rope_as_data` (ported verbatim from the C++ blueprint):
//!   * **RoPE rotate-half via a `[hd,hd]` matmul** (only when `rotary_dim == head_dim`) so the
//!     compiled graph carries no Slice (which shapeless compile cannot shape-infer).
//!   * **Pre-sliced cos/sin rows fed as synthetic closure inputs**, so the position offset is DATA
//!     (not baked) and the compiled graph never slices a cos/sin cache at a runtime position.
//!
//! Ineligible plans and trace/apply failures fall back to the eager translator. Device evaluation
//! failures abort Run: published outputs may already alias inputs, making eager replay unsafe.

use std::collections::HashSet;
use std::os::raw::c_void;
use std::panic::AssertUnwindSafe;

use crate::engine::{
    DeltaWrite, DynInput, MlxError, NodeDesc, OutRef, Plan, ShapeMode, Slot, Src, SynthRope,
    TracePayload, TranslationContext, attention_cross_kv, attention_past_key, copy_out_raw_delta,
    dim_i32, is_separate_qkv_attention, mlx_dtype_from_onnx, publish_device_output,
    read_ctx_input_metadata, read_ctx_input_raw, read_ctx_input_value, rope_row_key,
};
use crate::mlx::{self, Array, Closure, VectorArray};
use crate::registry::CompileShapeSafety;
use crate::sys::mlx as mlxsys;
use crate::sys::ort;

/// Decide whether a compiled fast path is allowed for this plan. The caller passes whether that
/// particular route has an unsupported graph-structure dependency. Static control flow is allowed
/// on the general route and specialized below.
pub fn compile_enabled(has_control_flow: bool) -> bool {
    let killed = std::env::var_os("ONNXRUNTIME_EP_MLX_NO_COMPILE")
        .map(|v| v != "0" && !v.is_empty())
        .unwrap_or(false);
    !has_control_flow && !killed
}

fn node_compile_shape_safety(node: &NodeDesc) -> CompileShapeSafety {
    node.compile_shape_safety
}

fn compile_shape_mode_supported(nodes: &[NodeDesc], shapeless: bool) -> bool {
    nodes.iter().all(|node| {
        (if shapeless {
            node_compile_shape_safety(node).allows_shapeless()
        } else {
            node_compile_shape_safety(node).allows_shape_keyed()
        }) && node
            .subgraphs
            .iter()
            .all(|subgraph| compile_shape_mode_supported(&subgraph.nodes, shapeless))
    })
}

/// Decode has no host-value specialization key. Shape/Size literals and host-read parameters
/// must therefore be proven safe independently of MLX's primitive output-shape inference.
pub fn decode_enabled(has_control_flow: bool, nodes: &[NodeDesc]) -> bool {
    compile_enabled(has_control_flow)
        && compile_shape_mode_supported(nodes, true)
        && !nodes.iter().any(reads_data_dependent_parameter)
}

/// Op types whose subgraphs keep their existing (eager / compiled-decode) routes because they are
/// incompatible with a single static-shape fused trace:
///   * cached attention ops carry KV-cache aliasing + delta copy-out semantics not modelled by the
///     general config (they take the decode/prefill configs instead); cache-free standard Attention
///     can use ordinary shape-keyed capture;
///   * host-computed ops GPU-eval their DYNAMIC input DATA mid-translate (`contiguous_eval`) and/or
///     emit a data-dependent output shape — both illegal inside an `mlx_compile` trace (the
///     placeholder has no data), so a subgraph containing one is never general-compiled.
fn is_general_compile_unsafe(n: &NodeDesc) -> bool {
    let cache_free_attention = n.domain.is_empty()
        && n.op_type == "Attention"
        && n.inputs
            .iter()
            .skip(4)
            .all(|input| input.source == Src::Absent)
        && n.outputs
            .iter()
            .skip(1)
            .all(|output| output.name.is_empty());
    (is_separate_qkv_attention(n) && !cache_free_attention)
        || matches!(
            n.op_type.as_str(),
            "GroupQueryAttention"
                | "PagedAttention"
                | "SparseAttention"
                | "Det"
                | "ImageDecoder"
                | "NonMaxSuppression"
                | "NonZero"
                | "TfIdfVectorizer"
                | "Unique"
        )
        || n.subgraphs
            .iter()
            .any(|sg| sg.nodes.iter().any(is_general_compile_unsafe))
}

fn host_parameter_indices(n: &NodeDesc) -> &'static [usize] {
    if n.domain == "pkg.onnax" && n.op_type == "SumToShape" {
        return &[1];
    }
    match n.op_type.as_str() {
        "Reshape" | "Expand" | "AffineGrid" | "Squeeze" | "Unsqueeze" => &[1],
        "ConstantOfShape" => &[0],
        "Slice" => &[1, 2, 3, 4],
        "Range" => &[0, 1, 2],
        "DFT" => &[1, 2],
        "STFT" => &[1, 3],
        "HannWindow" | "HammingWindow" | "BlackmanWindow" => &[0],
        "MelWeightMatrix" => &[0, 1, 2, 3, 4],
        "MaxUnpool" => &[1, 2],
        "TensorScatter" => &[2],
        "Col2Im" => &[1, 2],
        "CumSum" | "CumProd" => &[1],
        "RNN" | "GRU" | "LSTM" => &[4],
        "TopK" | "Split" => &[1],
        "Pad" => &[1, 3],
        "ReduceSum" | "ReduceMean" | "ReduceMax" | "ReduceMin" | "ReduceProd"
        | "ReduceSumSquare" | "ReduceL1" | "ReduceL2" | "ReduceLogSum" | "ReduceLogSumExp" => &[1],
        _ => &[],
    }
}

fn guarded_boundary_parameter(n: &NodeDesc, r: &crate::engine::TensorRef) -> bool {
    if n.domain == "pkg.onnax" && n.op_type == "SumToShape" {
        return r.source == Src::CtxInput;
    }
    matches!(
        n.op_type.as_str(),
        "Reshape"
            | "Expand"
            | "AffineGrid"
            | "Slice"
            | "ConstantOfShape"
            | "DFT"
            | "STFT"
            | "HannWindow"
            | "HammingWindow"
            | "BlackmanWindow"
            | "RNN"
            | "GRU"
            | "LSTM"
            | "TopK"
            | "Split"
            | "Pad"
            | "Squeeze"
            | "Unsqueeze"
    ) && r.source == Src::CtxInput
}

pub fn boundary_shape_inputs(nodes: &[NodeDesc]) -> Vec<usize> {
    let mut inputs = Vec::new();
    for node in nodes {
        for &i in host_parameter_indices(node) {
            if let Some(r) = node.inputs.get(i)
                && !r.constant
                && !r.shape_const
                && guarded_boundary_parameter(node, r)
                && !inputs.contains(&r.ctx_index)
            {
                inputs.push(r.ctx_index);
            }
        }
        for subgraph in &node.subgraphs {
            for index in boundary_shape_inputs(&subgraph.nodes) {
                if !inputs.contains(&index) {
                    inputs.push(index);
                }
            }
        }
    }
    inputs
}

fn boundary_shape_specialization_key(
    indices: &[usize],
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
) -> Result<String, MlxError> {
    let mut key = String::new();
    for &index in indices {
        let (data, shape, dtype) = read_ctx_input_raw(api, kctx, index)?;
        let count = shape.iter().try_fold(1usize, |count, &dim| {
            usize::try_from(dim)
                .ok()
                .and_then(|dim| count.checked_mul(dim))
        });
        let count = count
            .filter(|&count| count > 0 && count <= 64)
            .ok_or_else(|| "boundary shape specialization requires 1..=64 integers".to_string())?;
        if data.is_null() {
            return Err("boundary shape specialization input has null data".to_string());
        }
        let values: Vec<i64> = match dtype {
            ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 => unsafe {
                std::slice::from_raw_parts(data as *const i64, count).to_vec()
            },
            ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 => unsafe {
                std::slice::from_raw_parts(data as *const i32, count)
                    .iter()
                    .map(|&v| i64::from(v))
                    .collect()
            },
            _ => return Err("boundary shape specialization requires int32/int64".to_string()),
        };
        key.push_str(&format!("P{index}:{dtype}:{shape:?}={values:?};"));
    }
    Ok(key)
}

/// Shape-keyed MLX caches do not observe input values. Only initializers, proven shape-constant
/// intermediates, or explicitly value-guarded general-route boundary parameters are safe.
fn reads_data_dependent_parameter(n: &NodeDesc) -> bool {
    reads_unguarded_parameter(n, false)
}

fn reads_unguarded_parameter(n: &NodeDesc, guard_boundary_values: bool) -> bool {
    if n.subgraphs.iter().any(|sg| {
        sg.nodes
            .iter()
            .any(|n| reads_unguarded_parameter(n, guard_boundary_values))
    }) {
        return true;
    }
    host_parameter_indices(n).iter().any(|&i| {
        n.inputs.get(i).is_some_and(|tr| {
            !matches!(tr.source, Src::Initializer | Src::Absent)
                && !tr.constant
                && !tr.shape_const
                && !(guard_boundary_values && guarded_boundary_parameter(n, tr))
        })
    })
}

/// General capture shares NO_COMPILE and also honors NO_GENERAL_COMPILE for eager A/B validation.
pub fn general_enabled(has_control_flow: bool, nodes: &[NodeDesc]) -> bool {
    if !compile_enabled(has_control_flow) {
        return false;
    }
    if std::env::var_os("ONNXRUNTIME_EP_MLX_NO_GENERAL_COMPILE")
        .map(|v| v != "0" && !v.is_empty())
        .unwrap_or(false)
    {
        return false;
    }
    !nodes
        .iter()
        .any(|n| is_general_compile_unsafe(n) || reads_unguarded_parameter(n, true))
        && compile_shape_mode_supported(nodes, false)
}

fn host_scalar_i64(
    r: &crate::engine::TensorRef,
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
) -> Result<i64, MlxError> {
    let data = match r.source {
        Src::Initializer => r
            .init
            .as_ref()
            .map(|init| init.data)
            .ok_or_else(|| format!("control-flow initializer {} has no data", r.name))?,
        Src::CtxInput => read_ctx_input_raw(api, kctx, r.ctx_index)?.0,
        _ => {
            return Err(format!(
                "control-flow input {} is not host-readable",
                r.name
            ));
        }
    };
    if data.is_null() {
        return Err(format!("control-flow input {} has null data", r.name));
    }
    Ok(unsafe { *(data as *const i64) })
}

fn host_scalar_bool(
    r: &crate::engine::TensorRef,
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
) -> Result<bool, MlxError> {
    let data = match r.source {
        Src::Initializer => r
            .init
            .as_ref()
            .map(|init| init.data)
            .ok_or_else(|| format!("control-flow initializer {} has no data", r.name))?,
        Src::CtxInput => read_ctx_input_raw(api, kctx, r.ctx_index)?.0,
        _ => {
            return Err(format!(
                "control-flow input {} is not host-readable",
                r.name
            ));
        }
    };
    if data.is_null() {
        return Err(format!("control-flow input {} has null data", r.name));
    }
    Ok(unsafe { *(data as *const u8) } != 0)
}

/// Values that change the graph emitted by the eager Rust translator. MLX's shape-keyed cache only
/// observes tensor shapes/dtypes, so specialize the general closure explicitly when these scalar
/// values change. Scan needs no data key: its unroll count is the scan input's shape and therefore
/// already participates in MLX's native shape key.
pub fn control_flow_specialization_key(
    nodes: &[NodeDesc],
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
) -> Result<String, MlxError> {
    fn append(
        nodes: &[NodeDesc],
        api: *const ort::OrtApi,
        kctx: *mut ort::OrtKernelContext,
        out: &mut String,
        ordinal: &mut usize,
    ) -> Result<(), MlxError> {
        for node in nodes {
            match node.op_type.as_str() {
                "If" => {
                    let cond_ref = node
                        .inputs
                        .first()
                        .ok_or_else(|| "control-flow If has no condition input".to_string())?;
                    let cond = host_scalar_bool(cond_ref, api, kctx)?;
                    let current = *ordinal;
                    *ordinal += 1;
                    out.push_str(&format!("I{current}={};", u8::from(cond)));
                    let branch = if cond { "then_branch" } else { "else_branch" };
                    if let Some(sg) = node.subgraphs.iter().find(|sg| sg.attr_name == branch) {
                        append(&sg.nodes, api, kctx, out, ordinal)?;
                    }
                }
                "Loop" => {
                    let trip_ref = node
                        .inputs
                        .first()
                        .ok_or_else(|| "control-flow Loop has no trip-count input".to_string())?;
                    let cond_ref = node
                        .inputs
                        .get(1)
                        .ok_or_else(|| "control-flow Loop has no condition input".to_string())?;
                    let trip = host_scalar_i64(trip_ref, api, kctx)?;
                    let cond = host_scalar_bool(cond_ref, api, kctx)?;
                    let current = *ordinal;
                    *ordinal += 1;
                    out.push_str(&format!("L{current}={trip},{};", u8::from(cond)));
                    if cond
                        && trip > 0
                        && let Some(sg) = node.subgraphs.iter().find(|sg| sg.attr_name == "body")
                    {
                        append(&sg.nodes, api, kctx, out, ordinal)?;
                    }
                }
                "Scan" => {
                    if let Some(sg) = node.subgraphs.iter().find(|sg| sg.attr_name == "body") {
                        append(&sg.nodes, api, kctx, out, ordinal)?;
                    }
                }
                _ => {}
            }
        }
        Ok(())
    }

    let mut key = String::new();
    let mut ordinal = 0;
    append(nodes, api, kctx, &mut key, &mut ordinal)?;
    Ok(key)
}

/// Decide whether the compiled-PREFILL fast path (Phase 2) is allowed for this plan. Prefill is the
/// same decoder subgraph as decode but at query length S>1, so it needs a `GroupQueryAttention` node
/// (the RoPE/KV machinery build declines otherwise) and no control-flow.
///
/// DEFAULT ON. The shape-keyed prefill trace statically narrows attention to the valid prefix
/// `[0, valid_past+S)` (S and valid_past are fixed per shape key), so it no longer runs SDPA over the
/// full KV capacity — the compute-bound waste that made it a ~15-20% TTFT regression. With the
/// narrowing it is byte-identical to eager AND beats eager TTFT, so it is enabled by default. The
/// kill-switch `ONNXRUNTIME_EP_MLX_NO_PREFILL_COMPILE=1` forces eager prefill (A/B numerical validation / safety)
/// without touching the decode path. Also honours the global compile kill-switch. The build itself
/// falls back to eager for any non-decoder / partial-rotary shape.
pub fn prefill_enabled(has_control_flow: bool, nodes: &[NodeDesc]) -> bool {
    if !compile_enabled(has_control_flow) {
        return false;
    }
    if std::env::var_os("ONNXRUNTIME_EP_MLX_NO_PREFILL_COMPILE")
        .map(|v| v != "0" && !v.is_empty())
        .unwrap_or(false)
    {
        return false;
    }
    compile_shape_mode_supported(nodes, false)
        && !nodes.iter().any(reads_data_dependent_parameter)
        && nodes.iter().any(|n| n.op_type == "GroupQueryAttention")
}

/// Query sequence length S = trailing dim of the `input_ids` dynamic ctx input (decode => 1, prefill
/// => prompt length). Scans the plan nodes directly so it works before the compiled closure's
/// `dyn_inputs` list is built. Returns `None` if the input is not found / has no shape.
pub fn detect_seq_len(
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    plan: &Plan,
) -> Option<i32> {
    for node in &plan.nodes {
        for inp in &node.inputs {
            if inp.source == Src::CtxInput
                && !inp.constant
                && inp.name == "input_ids"
                && let Ok((shape, _t)) = read_ctx_input_metadata(api, kctx, inp.ctx_index)
            {
                return shape.last().map(|&d| d as i32);
            }
        }
    }
    if plan
        .nodes
        .iter()
        .any(|node| is_separate_qkv_attention(node) || node.op_type == "GroupQueryAttention")
    {
        // Decoder partitions may consume embeddings directly, or a dynamic position-id island may
        // split `input_ids` away from the main partition. The attention partition still receives a
        // rank-3 token-state boundary tensor [B,S,H], whose middle dimension is the query length.
        let mut seen = HashSet::new();
        for inp in plan.nodes.iter().flat_map(|n| n.inputs.iter()) {
            if inp.source == Src::CtxInput
                && !inp.constant
                && seen.insert(inp.ctx_index)
                && let Ok((shape, _t)) = read_ctx_input_metadata(api, kctx, inp.ctx_index)
                && shape.len() == 3
                && shape[1] > 0
            {
                return Some(shape[1] as i32);
            }
        }
    }
    None
}

pub fn detect_batch_size(
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    plan: &Plan,
) -> Option<i32> {
    let mut seen = HashSet::new();
    for inp in plan.nodes.iter().flat_map(|node| node.inputs.iter()) {
        if inp.source == Src::CtxInput
            && !inp.constant
            && seen.insert(inp.ctx_index)
            && let Ok((shape, _dtype)) = read_ctx_input_metadata(api, kctx, inp.ctx_index)
            && inp.name == "input_ids"
        {
            return shape.first().map(|&dim| dim as i32);
        }
    }
    None
}

fn kv_sequence_length(shape: &[i64]) -> Option<i64> {
    match shape.len() {
        3 => shape.get(1).copied(),
        4 => shape.get(2).copied(),
        _ => None,
    }
}

/// Sequence length of the first self-attention past-K cache. A zero-length self cache marks a new
/// generation and invalidates cached cross K/V for both contrib MHA and native ONNX Attention.
pub fn detect_attention_self_past_len(
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    plan: &Plan,
) -> Option<i64> {
    for node in &plan.nodes {
        let Some(past) = attention_past_key(node) else {
            continue;
        };
        if past.source != Src::CtxInput {
            continue;
        }
        let (shape, _dtype) = read_ctx_input_metadata(api, kctx, past.ctx_index).ok()?;
        return shape.get(2).copied();
    }
    None
}

/// Stable identity for the active attention generation: the first long cross-attention K boundary
/// buffer. Interleaved generators have different buffers, so a key change invalidates private MLX
/// input reuse instead of accidentally sharing one generator's cache with another.
pub fn detect_attention_generation_key(
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    plan: &Plan,
) -> Option<usize> {
    for node in &plan.nodes {
        let Some((key, _value)) = attention_cross_kv(node) else {
            continue;
        };
        let (data, shape, _dtype) = read_ctx_input_raw(api, kctx, key.ctx_index).ok()?;
        if kv_sequence_length(&shape).is_some_and(|length| length > 1) {
            return Some(data as usize);
        }
    }
    None
}

/// Detect whether this session drives a fixed-capacity SHARED KV buffer (present aliased onto past at
/// a runtime-owned max length) as opposed to the growing past/present contract. Reads live ctx once
/// at compiled-closure build time: `cap` = a GQA past-KV cache's seq-axis (2) length, and `total` =
/// valid keys after this step (= `total_sequence_length[0]`, or the `attention_mask` width when that
/// scalar is computed in-subgraph). Shared iff `cap > total - S` (the buffer holds more than the
/// valid past). Returns `(shared, mask_ctx_index, total_seq_ctx_index)`. `None` when the shape is not
/// a recognizable decoder GQA.
fn detect_shared_kv(
    plan: &Plan,
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
) -> Option<(bool, i32, i32)> {
    let s = detect_seq_len(api, kctx, plan).unwrap_or(1);
    let mask_ctx_index = plan
        .nodes
        .iter()
        .flat_map(|n| n.inputs.iter())
        .find(|inp| inp.source == Src::CtxInput && inp.name == "attention_mask")
        .map(|inp| inp.ctx_index as i32)
        .unwrap_or(-1);
    for node in &plan.nodes {
        if node.op_type != "GroupQueryAttention" || node.inputs.len() < 7 {
            continue;
        }
        // Capacity from the past-KV cache (input[3], a ctx input) seq axis.
        let past_k = &node.inputs[3];
        if past_k.source != Src::CtxInput {
            continue;
        }
        let cap = match read_ctx_input_raw(api, kctx, past_k.ctx_index) {
            Ok((_d, shape, _t)) => *shape.get(2)? as i32,
            Err(_) => continue,
        };
        // Valid keys after this step. Prefer total_sequence_length (input[6]) when it arrives as a
        // ctx input; otherwise recover it from the attention_mask width (the in-subgraph scalar is
        // Cast(Gather(Shape(attention_mask),1))).
        let ts = &node.inputs[6];
        let mut total_seq_ctx_index = -1;
        let total = if ts.source == Src::CtxInput && !ts.constant {
            total_seq_ctx_index = ts.ctx_index as i32;
            match read_ctx_input_raw(api, kctx, ts.ctx_index) {
                Ok((data, _shape, _t)) if !data.is_null() => unsafe { *(data as *const i32) },
                _ => continue,
            }
        } else if mask_ctx_index >= 0 {
            match read_ctx_input_raw(api, kctx, mask_ctx_index as usize) {
                Ok((_d, shape, _t)) => *shape.get(1)? as i32,
                Err(_) => continue,
            }
        } else {
            continue;
        };
        if s <= 0 || total < s || total - s > cap {
            return None;
        }
        return Some((cap > total - s, mask_ctx_index, total_seq_ctx_index));
    }
    None
}

/// A contiguous unit-stride slice `[start, stop)` -> owning [`Array`].
fn mk_slice(
    a: mlxsys::mlx_array,
    start: &[i32],
    stop: &[i32],
    stream: mlxsys::mlx_stream,
) -> Result<Array, MlxError> {
    let stride = vec![1i32; start.len()];
    let mut res = unsafe { mlxsys::mlx_array_new() };
    let rc = unsafe {
        mlxsys::mlx_slice(
            &mut res,
            a,
            start.as_ptr(),
            start.len(),
            stop.as_ptr(),
            stop.len(),
            stride.as_ptr(),
            stride.len(),
            stream,
        )
    };
    if rc != 0 {
        unsafe { mlxsys::mlx_array_free(res) };
        return Err("mlx_slice failed".to_string());
    }
    Ok(Array::from_raw(res))
}

/// Concatenate two arrays along `axis` -> owning [`Array`].
fn mk_concat(
    a: mlxsys::mlx_array,
    b: mlxsys::mlx_array,
    axis: i32,
    stream: mlxsys::mlx_stream,
) -> Result<Array, MlxError> {
    let mut v = VectorArray::new();
    v.append(a);
    v.append(b);
    let mut res = unsafe { mlxsys::mlx_array_new() };
    let rc = unsafe { mlxsys::mlx_concatenate_axis(&mut res, v.as_raw(), axis, stream) };
    if rc != 0 {
        unsafe { mlxsys::mlx_array_free(res) };
        return Err("mlx_concatenate_axis failed".to_string());
    }
    Ok(Array::from_raw(res))
}

/// Shape (as `i32`) of a borrowed raw array handle.
fn shape_of(a: mlxsys::mlx_array) -> Vec<i32> {
    let nd = unsafe { mlxsys::mlx_array_ndim(a) };
    let sh = unsafe { mlxsys::mlx_array_shape(a) };
    (0..nd).map(|i| unsafe { *sh.add(i) }).collect()
}

/// The compiled fast path for the given `slot`. Returns:
///   * `Ok(true)`  — the call was handled by the compiled closure (outputs already copied out).
///   * `Ok(false)` — the plan is not compile-eligible (caller must fall back to the eager path).
///   * `Err(_)`    — a hard failure copying results out.
///
/// The plan is accessed through `plan_ptr` (raw) so that NO Rust `&mut Plan` is alive across
/// `mlx_closure_apply` — the trace thunk (invoked synchronously on the first apply / a retrace)
/// needs its own `&mut Plan`, so only one mutable access is ever live at a time (single-threaded,
/// reentrant via FFI, exactly like the C++ EP).
pub fn try_compiled(
    plan_ptr: *mut Plan,
    slot: Slot,
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    stream: mlxsys::mlx_stream,
) -> Result<bool, MlxError> {
    let mut cfg = slot.get(unsafe { &*plan_ptr }).config;
    if !slot.get(unsafe { &*plan_ptr }).enabled {
        return Ok(false);
    }
    {
        let nodes: Vec<_> = unsafe { &*plan_ptr }
            .nodes
            .iter()
            .filter(|n| crate::ops::vision::needs_roi_validation(n))
            .cloned()
            .collect();
        if !nodes.is_empty() {
            let mut ctx = TranslationContext::new(unsafe { &mut *plan_ptr }, api, kctx, stream);
            if !crate::ops::vision::validate_compiled_rois(&mut ctx, &nodes)? {
                return Ok(false);
            }
        }
    }
    {
        let nodes: Vec<_> = unsafe { &*plan_ptr }
            .nodes
            .iter()
            .filter(|n| crate::ops::attention::needs_rotary_position_validation(n))
            .cloned()
            .collect();
        if !nodes.is_empty() {
            let mut ctx = TranslationContext::new(unsafe { &mut *plan_ptr }, api, kctx, stream);
            if !crate::ops::attention::validate_compiled_rotary_positions(&mut ctx, &nodes)? {
                return Ok(false);
            }
        }
    }
    let general = &unsafe { &*plan_ptr }.general;
    let host_specialization_token = if slot == Slot::General
        && (general.control_flow_specialization_enabled
            || !general.boundary_shape_inputs.is_empty())
    {
        let mut key = if general.control_flow_specialization_enabled {
            match control_flow_specialization_key(&unsafe { &*plan_ptr }.nodes, api, kctx) {
                Ok(key) => key,
                Err(_) => return Ok(false),
            }
        } else {
            String::new()
        };
        match boundary_shape_specialization_key(&general.boundary_shape_inputs, api, kctx) {
            Ok(parameters) => key.push_str(&parameters),
            Err(error) => {
                log::debug!("boundary shape specialization declined ({error}); using eager");
                return Ok(false);
            }
        }
        if key.is_empty() {
            None
        } else {
            let general = slot.get_mut(unsafe { &mut *plan_ptr });
            if let Some(&token) = general.host_specializations.get(&key) {
                general.host_specialization_key = Some(key);
                Some(token)
            } else if general.host_specializations.len() >= 16 {
                log::debug!("host-value specialization limit reached; using eager for new values");
                return Ok(false);
            } else {
                let token = general.host_specializations.len() as i32 + 1;
                general.host_specializations.insert(key.clone(), token);
                general.host_specialization_key = Some(key);
                Some(token)
            }
        }
    } else {
        None
    };
    {
        let plan = unsafe { &*plan_ptr };
        let has_gqa = plan
            .nodes
            .iter()
            .any(|n| n.op_type == "GroupQueryAttention");
        let has_separate_attention = plan.nodes.iter().any(is_separate_qkv_attention);
        let has_fused_gqa_rope = plan.nodes.iter().any(|n| {
            n.op_type == "GroupQueryAttention" && n.ints.get("do_rotary").copied().unwrap_or(0) != 0
        });
        if has_gqa && !has_fused_gqa_rope {
            // Qwen exports standalone RotaryEmbedding nodes and sets GQA do_rotary=0. Keep the
            // compiled KV-cache machinery, but do not require or synthesize GQA cos/sin inputs.
            cfg.rope_as_data = false;
            slot.get_mut(unsafe { &mut *plan_ptr }).config = cfg;
        }
        if matches!(cfg.shape_mode, ShapeMode::Shapeless) && has_separate_attention && !has_gqa {
            // Separate-QKV Attention decode has no RoPE inputs and uses growing (non-aliased) KV
            // outputs. Its current-token Q/K/V shapes stay fixed at S=1 while concat carries the
            // dynamic past extent inside MLX, so shapeless compile can replay it without GQA's
            // synthetic RoPE or shared-buffer machinery.
            cfg.rope_as_data = false;
            cfg.kv_alias = false;
            cfg.delta_copyout = false;
            cfg.contiguous_outputs = true;
            slot.get_mut(unsafe { &mut *plan_ptr }).config = cfg;
        }
    }
    // One-time discovery + compile. `build_closure` returns `false` for a TRANSIENT decline (the
    // constant cos/sin cache is not resident yet — it is only populated during the first eager
    // translation, which for prefill is the very first Compute call); in that case we leave
    // `attempted` unset so a later call retries once the cache is warm. Every other outcome (compiled,
    // or a permanent ineligibility / compile failure) is terminal and marks `attempted`.
    if !slot.get(unsafe { &*plan_ptr }).attempted {
        // Timing attribution: the one-time trace + `mlx_compile` (compile) phase.
        let terminal = {
            let _phase = crate::trace::tracer().phase("compile");
            build_closure(plan_ptr, slot, api, kctx, stream)
        };
        if terminal {
            slot.get_mut(unsafe { &mut *plan_ptr }).attempted = true;
        }
    }
    if !slot.get(unsafe { &*plan_ptr }).valid {
        return Ok(false);
    }
    if cfg.kv_alias {
        let Some((shared, _, _)) = detect_shared_kv(unsafe { &*plan_ptr }, api, kctx) else {
            return Ok(false);
        };
        if shared != slot.get(unsafe { &*plan_ptr }).shared_kv {
            // Equal-shaped live lengths may switch growing versus fixed-capacity cache semantics.
            // Reuse the closure only for the cache contract it traced.
            return Ok(false);
        }
    }

    // Skip shape-keyed tracing for empty or native-unrepresentable input geometry on EVERY call.
    // Empty traces can abort inside MLX; dimensions beyond i32 cannot be wrapped as MLX arrays.
    // Metadata-only eager consumers need neither operation.
    if cfg.shape_mode == ShapeMode::ShapeKeyed {
        let plan = unsafe { &*plan_ptr };
        for di in &slot.get(plan).dyn_inputs {
            let (shape, _dtype) = read_ctx_input_metadata(api, kctx, di.ctx_index)?;
            if shape.contains(&0) || shape.iter().any(|&dim| dim > i32::MAX as i64) {
                return Ok(false);
            }
        }
    }

    // Gather the dynamic ctx inputs in closure order. Cross-attention K/V are immutable for one
    // generation and very large, so separate-QKV attention decoders copy them into MLX once and
    // reuse them for both contrib MHA and native ONNX Attention.
    let mut arena: Vec<Array> = Vec::new();
    let mut borrowed_inputs = false;
    let mut input = VectorArray::new();
    {
        let dyn_inputs = slot.get(unsafe { &*plan_ptr }).dyn_inputs.clone();
        let has_gqa = unsafe { &*plan_ptr }
            .nodes
            .iter()
            .any(|n| n.op_type == "GroupQueryAttention");
        let stable_cross_indices = if unsafe { &*plan_ptr }.stable_cross_cache_enabled {
            unsafe { &*plan_ptr }.stable_cross_ctx_indices.clone()
        } else {
            Vec::new()
        };
        for di in &dyn_inputs {
            let value = read_ctx_input_value(api, kctx, di.ctx_index)?;
            if let Some(array) = crate::device::array(unsafe { &*api }, value)? {
                input.append(array.as_raw());
                arena.push(array);
                continue;
            }
            if stable_cross_indices.contains(&di.ctx_index)
                && let Some(arr) = slot
                    .get(unsafe { &*plan_ptr })
                    .stable_cross_inputs
                    .get(&di.ctx_index)
            {
                input.append(arr.as_raw());
                continue;
            }
            let (data, shape, dtype) = read_ctx_input_raw(api, kctx, di.ctx_index)?;
            if crate::engine::packed4_signed(dtype).is_some() {
                return Ok(false);
            }
            // Empty boundaries can become read-only K/V through intermediate elementwise ops.
            // Never replay the decoder trace's total-S append or delta writeback for that call.
            if has_gqa && shape.contains(&0) {
                return Ok(false);
            }
            let ishape: Vec<i32> = shape
                .iter()
                .map(|&d| dim_i32(d))
                .collect::<Result<_, _>>()?;
            if stable_cross_indices.contains(&di.ctx_index)
                && kv_sequence_length(&shape).is_some_and(|length| length > 1)
            {
                let arr = Array::from_data(data, &ishape, mlx_dtype_from_onnx(dtype));
                input.append(arr.as_raw());
                slot.get_mut(unsafe { &mut *plan_ptr })
                    .stable_cross_inputs
                    .insert(di.ctx_index, arr);
                continue;
            }
            // Zero-copy wrap of the live ORT input buffer: MLX borrows it (no-op deallocator) for
            // this apply only. `arena` keeps the wrapper alive until after eval + copy-out, and the
            // ORT tensor is valid for the whole Compute call, so the borrow never dangles. In
            // shared-buffer mode this is the per-token past K/V full [B,kv,cap,hd] buffer whose
            // memcpy dominated decode. (ORT KV buffers are 16 KB page-aligned, so Metal takes the
            // true no-copy path; small unaligned inputs fall back to MLX's internal copy.)
            let arr = Array::from_data_managed(data, &ishape, mlx_dtype_from_onnx(dtype));
            borrowed_inputs = true;
            input.append(arr.as_raw());
            arena.push(arr);
        }
        if let Some(token) = host_specialization_token {
            let zeros = vec![0u8; token as usize];
            let arr = Array::from_data(
                zeros.as_ptr() as *const std::os::raw::c_void,
                &[token],
                mlxsys::mlx_dtype__MLX_UINT8,
            );
            input.append(arr.as_raw());
            arena.push(arr);
        }
    }

    // ---- decode/prefill prologue: position + valid_past fed as DATA ------------------------------
    // `past` = RoPE start = position of the new query rows, and `s` = number of new rows (1 for
    // decode, S for prefill). Both the synth cos/sin rows and the KV delta below use them.
    let mut past: i32 = 0;
    let mut s_new: i32 = 0;
    if cfg.rope_as_data || cfg.kv_alias {
        let s = detect_seq_len(api, kctx, unsafe { &*plan_ptr }).unwrap_or(1);
        s_new = s;
        // Current RoPE start. In the growing contract this is the past-KV sequence length (past_k's
        // seq axis). In the shared-buffer contract past_k's seq axis is the fixed capacity, so the
        // true position is the VALID past = attention_mask width - S. Fed as DATA via the rows below.
        past = {
            let plan = unsafe { &*plan_ptr };
            let c = slot.get(plan);
            if c.shared_kv && c.total_seq_ctx_index >= 0 {
                let (data, _shape, _t) =
                    read_ctx_input_raw(api, kctx, c.total_seq_ctx_index as usize)?;
                if data.is_null() {
                    return Ok(false);
                }
                (unsafe { *(data as *const i32) }) - s
            } else if c.shared_kv && c.mask_ctx_index >= 0 {
                let (shape, _t) = read_ctx_input_metadata(api, kctx, c.mask_ctx_index as usize)?;
                (*shape.get(1).unwrap_or(&0) as i32) - s
            } else {
                if c.rope_past_ctx_index < 0 {
                    return Ok(false);
                }
                let idx = c.rope_past_ctx_index as usize;
                let axis = c.rope_past_axis as usize;
                let (shape, _t) = read_ctx_input_metadata(api, kctx, idx)?;
                if axis < shape.len() {
                    shape[axis] as i32
                } else {
                    0
                }
            }
        };

        if cfg.rope_as_data {
            // Pre-slice each fused-GQA RoPE cos/sin cache outside the compiled graph.
            let synth = slot.get(unsafe { &*plan_ptr }).synth_ropes.clone();
            for sr in &synth {
                let cache_raw = match unsafe { &*plan_ptr }.cache.get(&sr.cache_name) {
                    Some(a) => a.as_raw(),
                    None => return Ok(false),
                };
                let half = shape_of(cache_raw).get(1).copied().unwrap_or(0);
                if half == 0 {
                    return Ok(false);
                }
                let row = mk_slice(cache_raw, &[past, 0], &[past + s, half], stream)?;
                let full = mk_concat(row.as_raw(), row.as_raw(), 1, stream)?;
                input.append(full.as_raw());
                arena.push(row);
                arena.push(full);
            }
        }

        // Shared-buffer contract: append the live `valid_past` (= `past`) as a [1] int32 closure
        // input AFTER the synth RoPE rows. The compiled GQA op reads it to place the in-place K/V
        // write and build the causal mask — as pure per-step data, so shapeless compile never freezes
        // the growing offset. Kept in the arena for the apply's life.
        if slot.get(unsafe { &*plan_ptr }).shared_kv {
            let vp = [past];
            let arr = Array::from_data(
                vp.as_ptr() as *const std::os::raw::c_void,
                &[1],
                mlxsys::mlx_dtype__MLX_INT32,
            );
            input.append(arr.as_raw());
            arena.push(arr);
        }
    }

    // Refresh the live kernel context on the trace payload (only read during a (re)trace).
    unsafe {
        if let Some(p) = slot.get_mut(&mut *plan_ptr).payload.as_mut() {
            p.kctx = kctx;
        }
    }

    // Take the compiled closure OUT of the plan so no `&mut Plan` field is borrowed across apply
    // (the trace thunk mutates the plan through `plan_ptr`).
    let closure = match slot.get_mut(unsafe { &mut *plan_ptr }).closure.take() {
        Some(c) => c,
        None => return Ok(false),
    };
    let apply_res = closure.apply(&input);
    slot.get_mut(unsafe { &mut *plan_ptr }).closure = Some(closure);

    let outs = match apply_res {
        Ok(v) => v,
        Err(e) => {
            if crate::trace::tracer().active() {
                log::debug!("compiled closure apply failed ({e}); falling back to eager");
            }
            slot.get_mut(unsafe { &mut *plan_ptr }).valid = false; // disable, fall back to eager
            return Ok(false);
        }
    };
    // Timing attribution: the synchronous GPU-inclusive eval phase of the compiled closure.
    let tr = crate::trace::tracer();
    let eval_t0 = if tr.active() {
        Some(std::time::Instant::now())
    } else {
        None
    };
    let ext_len = slot.get(unsafe { &*plan_ptr }).ext_outputs.len();
    if outs.size() != ext_len {
        return Err("MLX compiled closure returned an unexpected output count".into());
    }
    let mut device_outputs = !borrowed_inputs && ext_len > 0;
    if device_outputs {
        for i in 0..ext_len {
            let output = &slot.get(unsafe { &*plan_ptr }).ext_outputs[i];
            device_outputs &= publish_device_output(api, kctx, output, &outs.get(i))?;
        }
    }
    // A captured boundary must finish before its guard stops Metal capture.
    let capture = tr.begin_gpu_capture();
    let eval_result = if device_outputs && capture.is_none() {
        let result = unsafe { mlxsys::mlx_async_eval(outs.as_raw()) };
        if result == 0 {
            Ok(())
        } else {
            Err("MLX compiled device evaluation failed; inspect GPU availability".into())
        }
    } else {
        let _eval_region = tr.eval_region();
        mlx::eval(&outs)
    };
    if let Err(error) = eval_result {
        slot.get_mut(unsafe { &mut *plan_ptr }).valid = false;
        if crate::device::enabled() {
            return Err(format!(
                "MLX compiled evaluation failed: {error}; device state may have changed, so reinitialize bound state tensors before retrying with sufficient GPU memory"
            ));
        }
        return Ok(false);
    }
    if let Some(t0) = eval_t0 {
        tr.record_phase("eval", t0.elapsed());
    }

    let ext_len = slot.get(unsafe { &*plan_ptr }).ext_outputs.len();
    if outs.size() != ext_len {
        slot.get_mut(unsafe { &mut *plan_ptr }).valid = false;
        return Ok(false);
    }

    // Shared-buffer KV `present` outputs alias `past` in ORT memory, so their per-call copy-out is a
    // delta write of only the `S` new rows at axis-2 offset `past` instead of the whole
    // [B,kv,cap,hd] buffer (the O(1) copy-out win). The write is gated on the present buffer actually
    // matching the recorded `past` pointer, so a non-aliasing runtime safely takes the full copy.
    // Empty set (growing / non-delta paths) => every output takes the full copy.
    let kv_present: Vec<(String, usize)> = if cfg.delta_copyout {
        slot.get(unsafe { &*plan_ptr }).kv_present_names.clone()
    } else {
        Vec::new()
    };
    for i in 0..ext_len {
        let a = outs.get(i);
        // Clone the small OutRef so we don't hold a plan borrow across the copy-out FFI.
        let o: OutRef = slot.get(unsafe { &*plan_ptr }).ext_outputs[i].clone();
        let delta = kv_present
            .iter()
            .find(|(name, _)| *name == o.name)
            .and_then(|(_, past_ctx)| {
                read_ctx_input_raw(api, kctx, *past_ctx)
                    .ok()
                    .map(|(p, _, _)| DeltaWrite {
                        axis: 2,
                        offset: past as i64,
                        count: s_new as i64,
                        alias_ptr: p as usize,
                    })
            });
        copy_out_raw_delta(api, kctx, &o, a.as_raw(), delta)?;
    }
    // `arena` (transient per-call inputs) drops here — after eval + copy-out have consumed them.
    Ok(true)
}

/// One-time discovery + compile of the closure for `slot`. Populates the slot's `dyn_inputs`
/// (ordered dynamic ctx inputs = closure inputs) and `ext_outputs` (boundary outputs in append
/// order); for `rope_as_data` slots it also discovers `synth_ropes`, the RoPE-start source, and the
/// shared-KV contract. Compiles the closure (shapeless iff `shape_mode == Shapeless`). Leaves
/// `valid = false` (=> caller falls back to eager) if the plan is not eligible or the compile fails.
///
/// Returns whether the outcome is TERMINAL: `true` after a successful compile or a permanent
/// ineligibility, `false` for a transient decline (the constant cos/sin cache is not resident yet)
/// so the caller retries on a later call once the constants have been warmed by an eager translation.
fn build_closure(
    plan_ptr: *mut Plan,
    slot: Slot,
    api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    stream: mlxsys::mlx_stream,
) -> bool {
    let cfg = slot.get(unsafe { &*plan_ptr }).config;
    if !compile_shape_mode_supported(
        &unsafe { &*plan_ptr }.nodes,
        cfg.shape_mode == ShapeMode::Shapeless,
    ) {
        return true;
    }

    let mut dyn_inputs: Vec<DynInput> = Vec::new();
    let mut ext_outputs: Vec<OutRef> = Vec::new();
    let mut synth_ropes: Vec<SynthRope> = Vec::new();
    let mut rope_past: i32 = -1;
    {
        let plan = unsafe { &*plan_ptr };
        // Ordered, de-duplicated dynamic (non-constant) ctx inputs = the closure input vector.
        let mut seen: HashSet<String> = HashSet::new();
        fn collect_body_inputs(
            nodes: &[NodeDesc],
            seen: &mut HashSet<String>,
            dyn_inputs: &mut Vec<DynInput>,
        ) {
            for node in nodes {
                for (i, inp) in node.inputs.iter().enumerate() {
                    if node.op_type == "CastLike" && i == 1 {
                        // Only the statically declared dtype is used, never target data or shape.
                        continue;
                    }
                    if inp.source == Src::CtxInput && !inp.constant && seen.insert(inp.name.clone())
                    {
                        dyn_inputs.push(DynInput {
                            name: inp.name.clone(),
                            ctx_index: inp.ctx_index,
                        });
                    }
                }
                if matches!(node.op_type.as_str(), "If" | "Loop" | "Scan") {
                    for sg in &node.subgraphs {
                        collect_body_inputs(&sg.nodes, seen, dyn_inputs);
                    }
                }
            }
        }
        for node in &plan.nodes {
            for (i, inp) in node.inputs.iter().enumerate() {
                if node.op_type == "CastLike" && i == 1 {
                    // Other consumers still collect this input if they need its live data.
                    continue;
                }
                if inp.source == Src::CtxInput && !inp.constant && seen.insert(inp.name.clone()) {
                    dyn_inputs.push(DynInput {
                        name: inp.name.clone(),
                        ctx_index: inp.ctx_index,
                    });
                    if (cfg.rope_as_data || cfg.kv_alias)
                        && rope_past < 0
                        && inp.name.contains(".key")
                    {
                        rope_past = inp.ctx_index as i32;
                    }
                }
            }
            if rope_past < 0 {
                rope_past = plan
                    .nodes
                    .iter()
                    .find(|n| n.op_type == "GroupQueryAttention" && n.inputs.len() > 3)
                    .map(|n| &n.inputs[3])
                    .filter(|inp| inp.source == Src::CtxInput && !inp.constant)
                    .map(|inp| inp.ctx_index as i32)
                    .unwrap_or(-1);
            }
            if slot == Slot::General && matches!(node.op_type.as_str(), "If" | "Loop" | "Scan") {
                for sg in &node.subgraphs {
                    collect_body_inputs(&sg.nodes, &mut seen, &mut dyn_inputs);
                }
            }
        }
        // External boundary outputs, in stable node/output order.
        for node in &plan.nodes {
            for o in &node.outputs {
                if o.external {
                    ext_outputs.push(o.clone());
                }
            }
        }
        // Distinct RoPE cos/sin caches (GQA inputs[7]/[8] when do_rotary). Their per-step rows are
        // fed as synthetic closure inputs so the compiled graph never slices a cache at runtime.
        if cfg.rope_as_data {
            let mut synth_seen: HashSet<String> = HashSet::new();
            for node in &plan.nodes {
                if node.op_type != "GroupQueryAttention" || node.inputs.len() < 9 {
                    continue;
                }
                let do_rotary = node.ints.get("do_rotary").copied().unwrap_or(0) != 0;
                if !do_rotary {
                    continue;
                }
                for idx in [7usize, 8usize] {
                    let nm = &node.inputs[idx].name;
                    if synth_seen.insert(nm.clone()) {
                        synth_ropes.push(SynthRope {
                            key: rope_row_key(nm),
                            cache_name: node.inputs[idx].cache_key(),
                        });
                    }
                }
            }
        }
    }

    // Need at least one dynamic input (else the graph is fully constant — cheap to leave eager) and
    // at least one boundary output. `rope_as_data` slots additionally need the decoder RoPE shape.
    if dyn_inputs.is_empty() || ext_outputs.is_empty() {
        if crate::trace::tracer().active() {
            log::debug!(
                "compiled closure declined: {} dynamic inputs, {} external outputs",
                dyn_inputs.len(),
                ext_outputs.len()
            );
        }
        return true;
    }
    if cfg.rope_as_data && (rope_past < 0 || synth_ropes.is_empty()) {
        return true; // not the expected decoder shape; stay on the eager path
    }

    let mut shared_kv = false;
    let mut mask_ctx_index = -1;
    let mut total_seq_ctx_index = -1;
    if cfg.rope_as_data {
        // The compiled RoPE uses a [hd,hd] rotate-half matmul, which requires rotary_dim == head_dim
        // (rot == hd). Validate from the live ctx (head dim = last axis of the past-KV cache) vs the
        // cos cache width (half); fall back to eager for a partial-rotary head.
        let hd = match read_ctx_input_raw(api, kctx, rope_past as usize) {
            Ok((_d, shape, _t)) => shape.last().copied().unwrap_or(0) as i32,
            Err(_) => 0,
        };
        // The cos/sin cache is a constant only loaded into `plan.cache` during an eager translation.
        // On the very first Compute call (prefill) it is not resident yet — a TRANSIENT condition, so
        // return non-terminal and retry once the constants are warm (the apply prologue needs them
        // too). A resident cache with the wrong width is a PERMANENT decline (partial-rotary head).
        let (half, cache_resident) = {
            let plan = unsafe { &*plan_ptr };
            match plan.cache.get(&synth_ropes[0].cache_name) {
                Some(a) => {
                    let sh = a.shape();
                    (if sh.len() >= 2 { sh[1] as i32 } else { 0 }, true)
                }
                None => (0, false),
            }
        };
        if !cache_resident {
            if crate::trace::tracer().active() {
                log::debug!("compiled closure deferred: RoPE cache is not resident");
            }
            return false; // transient — constants not warmed yet
        }
        if hd == 0 || half == 0 || 2 * half != hd {
            if crate::trace::tracer().active() {
                log::debug!(
                    "compiled closure declined: unsupported RoPE width (head={hd}, half={half})"
                );
            }
            return true; // partial-rotary head not supported by the compiled path
        }
    }
    if cfg.kv_alias {
        // Detect a fixed-capacity shared KV buffer once from live ctx. In that contract GQA writes
        // the new K/V in place at the (data-dependent) valid-past offset and emits present at the
        // buffer's full capacity — the trace handles both via a data start index + a static-shape
        // additive mask, so the compiled fast path still applies.
        let detected =
            detect_shared_kv(unsafe { &*plan_ptr }, api, kctx).unwrap_or((false, -1, -1));
        shared_kv = detected.0;
        mask_ctx_index = detected.1;
        total_seq_ctx_index = detected.2;
    }

    // Publish the discovered schema + a stable trace payload onto the slot.
    unsafe {
        let c = slot.get_mut(&mut *plan_ptr);
        c.config = cfg;
        c.shared_kv = shared_kv;
        c.mask_ctx_index = mask_ctx_index;
        c.total_seq_ctx_index = total_seq_ctx_index;
        c.dyn_inputs = dyn_inputs;
        c.ext_outputs = ext_outputs;
        c.synth_ropes = synth_ropes;
        c.rope_past_ctx_index = rope_past;
        c.rope_past_axis = 2;
        c.payload = Some(Box::new(TracePayload {
            plan: plan_ptr,
            ort_api: api,
            kctx,
            stream,
            slot,
        }));
    }

    let payload_ptr: *mut c_void = unsafe {
        slot.get_mut(&mut *plan_ptr)
            .payload
            .as_mut()
            .unwrap()
            .as_mut() as *mut TracePayload as *mut c_void
    };
    // Shapeless (decode) so the growing KV length never triggers a recompile; shape-keyed (general /
    // prefill) so a changed input shape safely retraces (re-invokes the thunk) rather than
    // miscomputes. The trace thunk runs lazily on the first `apply`.
    let shapeless = matches!(cfg.shape_mode, ShapeMode::Shapeless);
    let base = Closure::new_func_payload(trace_thunk, payload_ptr);
    match Closure::compile(&base, shapeless) {
        Ok(compiled) => {
            let c = slot.get_mut(unsafe { &mut *plan_ptr });
            c.closure = Some(compiled);
            c.valid = true;
        }
        Err(e) => {
            if crate::trace::tracer().active() {
                log::debug!("compiled closure creation failed ({e}); falling back to eager");
            }
        }
    }
    true
}

/// `mlx_closure` trace thunk (payload = [`TracePayload`]): seed each dynamic ctx input (+ pre-sliced
/// RoPE row / valid_past for `rope_as_data`) from the closure input vector, translate the whole
/// subgraph, and return the cast external boundary outputs. Invoked lazily by mlx on the first
/// `apply` and again after any input shape change (shape-keyed slots). Never unwinds across the FFI
/// boundary.
extern "C" fn trace_thunk(
    out: *mut mlxsys::mlx_vector_array,
    input: mlxsys::mlx_vector_array,
    payload: *mut c_void,
) -> std::os::raw::c_int {
    let result = std::panic::catch_unwind(AssertUnwindSafe(|| trace_body(out, input, payload)));
    match result {
        Ok(Ok(())) => 0,
        Ok(Err(e)) => {
            // Correctness fallback (compiled trace failed → eager). Env-gated so a traced-off run
            // stays silent; the eager path still produces the correct result.
            if crate::trace::tracer().active() {
                log::debug!("compiled trace failed ({e}); falling back to eager");
            }
            1
        }
        Err(payload) => {
            // A panic inside the trace thunk is unexpected (a bug), unlike the benign trace-failed
            // decline above — never keep it silent, even with tracing off. The eager path below
            // still produces a correct result.
            let detail = crate::panic_payload_message(&payload);
            log::warn!(
                "compiled trace panicked ({detail}); falling back to the eager path (result still correct)"
            );
            1
        }
    }
}

fn trace_body(
    out: *mut mlxsys::mlx_vector_array,
    input: mlxsys::mlx_vector_array,
    payload: *mut c_void,
) -> Result<(), MlxError> {
    let pl = unsafe { &mut *(payload as *mut TracePayload) };
    let plan_ptr = pl.plan;
    let api = pl.ort_api;
    let kctx = pl.kctx;
    let stream = pl.stream;
    let slot = pl.slot;

    // Snapshot the closure schema (clone the small metadata) so no plan borrow overlaps translation.
    let (cfg, dyn_inputs, synth_ropes, ext_outputs, shared_kv) = {
        let c = slot.get(unsafe { &*plan_ptr });
        (
            c.config,
            c.dyn_inputs.clone(),
            c.synth_ropes.clone(),
            c.ext_outputs.clone(),
            c.shared_kv,
        )
    };

    // Shape-keyed (prefill) shared-buffer traces know `valid_past` statically for this shape key (the
    // attention_mask width = valid_past + S is part of the key). Recover it here from the live ctx so
    // the compiled attention can be narrowed to the valid prefix `[0, valid_past+S)` instead of masked
    // over the full KV capacity. Only enable the narrow path when valid_past is reliably known (mask
    // width available and non-negative); otherwise leave `narrow` off so the trace keeps the safe
    // full-cap masked path. Shapeless (decode) traces feed valid_past as data and never narrow.
    let mut narrow = false;
    let mut static_valid_past = 0i32;
    if matches!(cfg.shape_mode, ShapeMode::ShapeKeyed) && cfg.rope_as_data && shared_kv {
        let s = detect_seq_len(api, kctx, unsafe { &*plan_ptr }).unwrap_or(1);
        let mask_idx = slot.get(unsafe { &*plan_ptr }).mask_ctx_index;
        if mask_idx >= 0
            && let Ok((_d, shape, _t)) = read_ctx_input_raw(api, kctx, mask_idx as usize)
        {
            let vp = (*shape.get(1).unwrap_or(&0) as i32) - s;
            if vp >= 0 {
                narrow = true;
                static_valid_past = vp;
            }
        }
    }

    let (res_raw, arena, kv_present) = {
        let plan = unsafe { &mut *plan_ptr };
        let native_attention_decode = plan.native_attention_decode;
        let mut tc = TranslationContext::new(plan, api, kctx, stream);
        // Independent of the branch below: record ShapeKeyed-ness alone (general/prefill vs decode),
        // for opt-in fast kernels that must never run during (Shapeless) decode.
        tc.set_shape_keyed_compile(matches!(cfg.shape_mode, ShapeMode::ShapeKeyed));
        if (matches!(cfg.shape_mode, ShapeMode::Shapeless) && native_attention_decode)
            || cfg.rope_as_data
            || cfg.kv_alias
        {
            // RoPE uses the pre-sliced cos/sin ROW placeholders + a matmul rotate-half, so the graph
            // carries no dynamic Slice (which shapeless `mlx_compile` cannot shape-infer).
            tc.set_compiled_trace(shared_kv, narrow, static_valid_past);
        } else {
            // General trace: dynamic inputs are shapeless placeholders with no host data, so any
            // mid-graph host eval must fail the trace cleanly (=> eager) rather than eval a tracer.
            tc.set_general_trace();
        }

        // Seed the dynamic ctx input placeholders (closure inputs [0..ndyn)).
        let ndyn = dyn_inputs.len();
        for (i, di) in dyn_inputs.iter().enumerate() {
            let mut a = unsafe { mlxsys::mlx_array_new() };
            unsafe { mlxsys::mlx_vector_array_get(&mut a, input, i) };
            let raw = tc.keep(Array::from_raw(a));
            tc.seed(di.name.clone(), raw);
        }
        if cfg.rope_as_data {
            // Seed the pre-sliced RoPE cos/sin row placeholders (closure inputs [ndyn..ndyn+nsynth)).
            for (j, sr) in synth_ropes.iter().enumerate() {
                let mut a = unsafe { mlxsys::mlx_array_new() };
                unsafe { mlxsys::mlx_vector_array_get(&mut a, input, ndyn + j) };
                let raw = tc.keep(Array::from_raw(a));
                tc.seed(sr.key.clone(), raw);
            }
        }
        // Shared-buffer: seed the live `valid_past` scalar after any synthetic RoPE rows.
        if shared_kv {
            let mut a = unsafe { mlxsys::mlx_array_new() };
            unsafe { mlxsys::mlx_vector_array_get(&mut a, input, ndyn + synth_ropes.len()) };
            let raw = tc.keep(Array::from_raw(a));
            tc.seed(crate::engine::GQA_VALID_PAST_KEY.to_string(), raw);
        }

        let res = tc.run_trace(&ext_outputs, cfg.contiguous_outputs)?;
        let kv_present = tc.take_compiled_kv_present();
        let arena = tc.take_arena();
        (res.into_raw(), arena, kv_present)
        // `tc` dropped here (its arena already taken) — releases the plan borrow.
    };

    // Hand the trace's transient handles + discovered KV-present set to the slot so the compiled
    // graph (walked after this thunk returns) keeps its inputs alive; they are freed once with the
    // plan.
    unsafe {
        let c = slot.get_mut(&mut *plan_ptr);
        c.trace_transient.extend(arena);
        if cfg.delta_copyout {
            c.kv_present_names = kv_present;
        }
        *out = res_raw;
    }
    Ok(())
}

#[cfg(test)]
mod shape_inference_tests {
    use super::*;
    use crate::engine::TensorRef;

    fn input(name: &str, source: Src) -> TensorRef {
        TensorRef {
            name: name.to_string(),
            source,
            ctx_index: 0,
            constant: false,
            shape_const: false,
            init: None,
        }
    }

    fn cumulative(op_type: &str, axis_source: Src, reason: &'static str) -> NodeDesc {
        let mut node = NodeDesc::new(
            op_type.to_string(),
            String::new(),
            14,
            CompileShapeSafety::ShapeKeyedOnly { reason },
        );
        node.inputs = vec![input("data", Src::CtxInput), input("axis", axis_source)];
        node
    }

    #[test]
    fn general_capture_allows_only_cache_free_standard_attention() {
        let mut node = NodeDesc::new(
            "Attention".to_string(),
            String::new(),
            23,
            CompileShapeSafety::Shapeless,
        );
        node.inputs = ["q", "k", "v", "mask"]
            .into_iter()
            .map(|name| input(name, Src::CtxInput))
            .collect();
        assert!(!is_general_compile_unsafe(&node));
        node.inputs.push(input("", Src::Absent));
        node.inputs.push(input("", Src::Absent));
        assert!(!is_general_compile_unsafe(&node));
        node.inputs[4] = input("past_key", Src::CtxInput);
        assert!(is_general_compile_unsafe(&node));
        node.inputs[4] = input("", Src::Absent);
        node.inputs[5] = input("past_value", Src::CtxInput);
        assert!(is_general_compile_unsafe(&node));
        node.inputs[5] = input("", Src::Absent);
        node.outputs = vec![
            OutRef {
                name: "output".to_string(),
                external: true,
                ctx_index: 0,
                otype: ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
            },
            OutRef {
                name: "present_key".to_string(),
                external: true,
                ctx_index: 1,
                otype: ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
            },
        ];
        assert!(is_general_compile_unsafe(&node));
        node.outputs.pop();
        assert!(!is_general_compile_unsafe(&node));
        node.op_type = "MultiHeadAttention".to_string();
        node.domain = "com.microsoft".to_string();
        assert!(is_general_compile_unsafe(&node));
        node.op_type = "Attention".to_string();
        node.domain.clear();
        node.compile_shape_safety = CompileShapeSafety::EagerOnly {
            reason: "composed attention form",
        };
        assert!(!compile_shape_mode_supported(&[node], false));
    }

    #[test]
    fn missing_output_shape_emitters_require_shape_specialized_compile() {
        for node in [
            NodeDesc::new(
                "CumSum".to_string(),
                String::new(),
                14,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SCAN_SHAPE_REASON,
                },
            ),
            NodeDesc::new(
                "CumProd".to_string(),
                String::new(),
                26,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SCAN_SHAPE_REASON,
                },
            ),
            NodeDesc::new(
                "Slice".to_string(),
                String::new(),
                14,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SLICE_SHAPE_REASON,
                },
            ),
            NodeDesc::new(
                "LinearAttention".to_string(),
                "com.microsoft".to_string(),
                1,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: "emits MLX Slice, Pad, and Scan",
                },
            ),
            NodeDesc::new(
                "QMoE".to_string(),
                "com.microsoft".to_string(),
                1,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SLICE_SHAPE_REASON,
                },
            ),
        ] {
            assert!(!node_compile_shape_safety(&node).allows_shapeless());
        }
        assert!(
            node_compile_shape_safety(&NodeDesc::new(
                "Add".to_string(),
                String::new(),
                14,
                CompileShapeSafety::Shapeless,
            ))
            .allows_shapeless()
        );
        assert!(compile_shape_mode_supported(
            &[NodeDesc::new(
                "CumSum".to_string(),
                String::new(),
                14,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SCAN_SHAPE_REASON,
                },
            )],
            false
        ));
        assert!(!compile_shape_mode_supported(
            &[NodeDesc::new(
                "CumSum".to_string(),
                String::new(),
                14,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SCAN_SHAPE_REASON,
                },
            )],
            true
        ));
        let eager_only_gqa = NodeDesc::new(
            "GroupQueryAttention".to_string(),
            "com.microsoft".to_string(),
            1,
            CompileShapeSafety::EagerOnly {
                reason: "attention-bias form",
            },
        );
        assert!(!compile_shape_mode_supported(
            std::slice::from_ref(&eager_only_gqa),
            false
        ));
        assert!(!compile_shape_mode_supported(&[eager_only_gqa], true));
    }

    #[test]
    fn constant_of_shape_guards_the_first_input() {
        let mut node = NodeDesc::new(
            "ConstantOfShape".to_string(),
            String::new(),
            23,
            CompileShapeSafety::ShapeKeyedOnly {
                reason: "runtime fill shape requires shape/value guards",
            },
        );
        let mut shape = input("shape", Src::CtxInput);
        shape.ctx_index = 3;
        node.inputs = vec![shape];
        assert!(reads_data_dependent_parameter(&node));
        assert!(!reads_unguarded_parameter(&node, true));
        assert_eq!(boundary_shape_inputs(std::slice::from_ref(&node)), vec![3]);
        assert!(!compile_shape_mode_supported(
            std::slice::from_ref(&node),
            true
        ));
        node.inputs[0].source = Src::Intermediate;
        assert!(reads_unguarded_parameter(&node, true));
        node.inputs[0].shape_const = true;
        assert!(!reads_unguarded_parameter(&node, true));
    }

    #[test]
    fn affine_grid_size_values_guard_general_capture_only() {
        let mut node = NodeDesc::new(
            "AffineGrid".to_string(),
            String::new(),
            20,
            CompileShapeSafety::ShapeKeyedOnly {
                reason: "host-generated homogeneous grid",
            },
        );
        let mut size = input("size", Src::CtxInput);
        size.ctx_index = 1;
        node.inputs = vec![input("theta", Src::CtxInput), size];
        assert!(reads_data_dependent_parameter(&node));
        assert!(!reads_unguarded_parameter(&node, true));
        assert_eq!(boundary_shape_inputs(std::slice::from_ref(&node)), vec![1]);
        assert!(!compile_shape_mode_supported(
            std::slice::from_ref(&node),
            true
        ));
        node.inputs[1].source = Src::Intermediate;
        assert!(reads_unguarded_parameter(&node, true));
        assert!(boundary_shape_inputs(std::slice::from_ref(&node)).is_empty());
        node.inputs[1].shape_const = true;
        assert!(!reads_data_dependent_parameter(&node));
        node.inputs[1].shape_const = false;
        node.inputs[1].source = Src::Initializer;
        assert!(!reads_data_dependent_parameter(&node));
    }

    #[test]
    fn max_unpool_host_reads_require_constant_indices_and_output_shape() {
        let mut node = NodeDesc::new(
            "MaxUnpool".to_string(),
            String::new(),
            22,
            CompileShapeSafety::ShapeKeyedOnly {
                reason: "host-validated indices and output shape",
            },
        );
        node.inputs = vec![
            input("x", Src::CtxInput),
            input("indices", Src::Initializer),
            input("shape", Src::Initializer),
        ];
        assert!(!reads_unguarded_parameter(&node, true));
        for index in [1, 2] {
            for source in [Src::CtxInput, Src::Intermediate] {
                node.inputs[index].source = source;
                assert!(reads_unguarded_parameter(&node, true));
                assert!(reads_data_dependent_parameter(&node));
            }

            node.inputs[index].source = Src::Initializer;
        }
    }

    #[test]
    fn col2im_host_geometry_requires_constant_controls() {
        let mut node = NodeDesc::new(
            "Col2Im".to_string(),
            String::new(),
            18,
            CompileShapeSafety::ShapeKeyedOnly {
                reason: "host-read image and block shapes",
            },
        );
        node.inputs = vec![
            input("x", Src::CtxInput),
            input("image", Src::Initializer),
            input("block", Src::Initializer),
        ];
        assert!(!reads_unguarded_parameter(&node, true));
        for index in [1, 2] {
            for source in [Src::CtxInput, Src::Intermediate] {
                node.inputs[index].source = source;
                assert!(reads_unguarded_parameter(&node, true));
                assert!(reads_data_dependent_parameter(&node));
            }
            node.inputs[index].source = Src::Initializer;
        }
    }

    #[test]
    fn tensor_scatter_never_captures_live_write_indices() {
        let mut node = NodeDesc::new(
            "TensorScatter".to_string(),
            String::new(),
            24,
            CompileShapeSafety::ShapeKeyedOnly {
                reason: crate::registry::MLX_SLICE_SHAPE_REASON,
            },
        );
        node.inputs = vec![
            input("past", Src::CtxInput),
            input("update", Src::CtxInput),
            input("indices", Src::Initializer),
        ];
        assert!(!reads_data_dependent_parameter(&node));
        for source in [Src::CtxInput, Src::Intermediate] {
            node.inputs[2].source = source;
            assert!(reads_data_dependent_parameter(&node));
            assert!(reads_unguarded_parameter(&node, true));
            assert!(boundary_shape_inputs(std::slice::from_ref(&node)).is_empty());
            assert!(!general_enabled(false, std::slice::from_ref(&node)));
        }
        node.inputs[2].source = Src::Absent;
        assert!(!reads_data_dependent_parameter(&node));
    }

    #[test]
    fn reduction_axes_are_guarded_for_all_compile_routes() {
        for op in [
            "ReduceSum",
            "ReduceMean",
            "ReduceMax",
            "ReduceMin",
            "ReduceProd",
            "ReduceSumSquare",
            "ReduceL1",
            "ReduceL2",
            "ReduceLogSum",
            "ReduceLogSumExp",
        ] {
            for source in [Src::CtxInput, Src::Intermediate] {
                let mut node = cumulative(op, source, "runtime reduction axes");
                assert!(reads_data_dependent_parameter(&node));
                assert!(reads_unguarded_parameter(&node, true));
                assert!(!compile_shape_mode_supported(
                    std::slice::from_ref(&node),
                    true
                ));
                assert!(compile_shape_mode_supported(
                    std::slice::from_ref(&node),
                    false
                ));
                if source == Src::Intermediate {
                    node.inputs[1].shape_const = true;
                    assert!(!reads_unguarded_parameter(&node, true));
                }
            }
            let node = cumulative(op, Src::Initializer, "runtime reduction axes");
            assert!(!reads_data_dependent_parameter(&node));
        }
    }

    #[test]
    fn signal_parameters_require_value_keys_or_eager_execution() {
        for (op, parameters) in [
            ("DFT", vec![1, 2]),
            ("STFT", vec![1, 3]),
            ("HannWindow", vec![0]),
            ("HammingWindow", vec![0]),
            ("BlackmanWindow", vec![0]),
        ] {
            let mut node = NodeDesc::new(
                op.to_string(),
                String::new(),
                20,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: "host-read signal controls",
                },
            );
            node.inputs = (0..4)
                .map(|i| {
                    let mut r = input(&format!("input{i}"), Src::CtxInput);
                    r.ctx_index = i;
                    r
                })
                .collect();
            assert!(reads_data_dependent_parameter(&node));
            assert!(!reads_unguarded_parameter(&node, true));
            assert_eq!(
                boundary_shape_inputs(std::slice::from_ref(&node)),
                parameters
            );
            for &i in &parameters {
                node.inputs[i].source = Src::Intermediate;
            }
            assert!(reads_unguarded_parameter(&node, true));
            for &i in &parameters {
                node.inputs[i].source = Src::Initializer;
            }
            assert!(!reads_data_dependent_parameter(&node));
        }
        let mut mel = NodeDesc::new(
            "MelWeightMatrix".to_string(),
            String::new(),
            17,
            CompileShapeSafety::ShapeKeyedOnly {
                reason: "host-computed weights",
            },
        );
        mel.inputs = (0..5)
            .map(|i| input(&format!("input{i}"), Src::CtxInput))
            .collect();
        assert!(reads_unguarded_parameter(&mel, true));
        assert!(boundary_shape_inputs(&[mel]).is_empty());
    }

    #[test]
    fn runtime_cumsum_axis_is_not_a_shape_key() {
        assert!(
            reads_data_dependent_parameter(&cumulative(
                "CumSum",
                Src::CtxInput,
                crate::registry::MLX_SCAN_SHAPE_REASON,
            )),
            "a runtime axis value must not be baked into a shape-keyed closure"
        );
        assert!(
            !reads_data_dependent_parameter(&cumulative(
                "CumSum",
                Src::Initializer,
                crate::registry::MLX_SCAN_SHAPE_REASON,
            )),
            "an initializer axis is compile-time data"
        );
    }

    #[test]
    fn runtime_cumprod_axis_is_not_a_shape_key() {
        assert!(
            reads_data_dependent_parameter(&cumulative(
                "CumProd",
                Src::CtxInput,
                crate::registry::MLX_SCAN_SHAPE_REASON,
            )),
            "a runtime axis value must not be baked into a shape-keyed closure"
        );
        assert!(
            !reads_data_dependent_parameter(&cumulative(
                "CumProd",
                Src::Initializer,
                crate::registry::MLX_SCAN_SHAPE_REASON,
            )),
            "an initializer axis is compile-time data"
        );
    }

    #[test]
    fn direct_runtime_shape_parameter_is_value_guarded() {
        let mut reshape = NodeDesc::new(
            "Reshape".to_string(),
            String::new(),
            14,
            CompileShapeSafety::Shapeless,
        );
        reshape.inputs = vec![input("data", Src::CtxInput), input("shape", Src::CtxInput)];
        assert!(reads_data_dependent_parameter(&reshape));
        assert!(!reads_unguarded_parameter(&reshape, true));
        assert_eq!(
            boundary_shape_inputs(std::slice::from_ref(&reshape)),
            vec![0]
        );

        reshape.inputs[1].source = Src::Intermediate;
        assert!(reads_data_dependent_parameter(&reshape));
        reshape.inputs[1].shape_const = true;
        assert!(
            !reads_data_dependent_parameter(&reshape),
            "Shape/Size-derived parameters change with the shape-keyed cache key"
        );
    }

    #[test]
    fn specialized_routes_do_not_exempt_general_route_boundary_guards() {
        for op in [
            "Reshape",
            "Expand",
            "Slice",
            "TopK",
            "Split",
            "Pad",
            "Squeeze",
            "Unsqueeze",
        ] {
            let mut parameterized = NodeDesc::new(
                op.to_string(),
                String::new(),
                14,
                CompileShapeSafety::Shapeless,
            );
            parameterized.inputs = vec![
                input("data", Src::CtxInput),
                input("parameter", Src::CtxInput),
                input("end", Src::CtxInput),
            ];
            assert!(reads_data_dependent_parameter(&parameterized));
            assert!(!reads_unguarded_parameter(&parameterized, true));
            assert!(!decode_enabled(false, std::slice::from_ref(&parameterized)));
            let gqa = NodeDesc::new(
                "GroupQueryAttention".to_string(),
                "com.microsoft".to_string(),
                1,
                CompileShapeSafety::Shapeless,
            );
            assert!(!prefill_enabled(false, &[gqa, parameterized]));
        }
    }

    #[test]
    fn recurrent_lengths_require_value_guards_for_validation() {
        for op in ["RNN", "GRU", "LSTM"] {
            let mut node = NodeDesc::new(
                op.to_string(),
                String::new(),
                22,
                CompileShapeSafety::ShapeKeyedOnly {
                    reason: crate::registry::MLX_SLICE_SHAPE_REASON,
                },
            );
            node.inputs = vec![crate::engine::TensorRef::absent(); 5];
            node.inputs[4] = input("lengths", Src::CtxInput);
            node.inputs[4].ctx_index = 4;
            assert!(reads_data_dependent_parameter(&node));
            assert!(!reads_unguarded_parameter(&node, true));
            assert_eq!(boundary_shape_inputs(std::slice::from_ref(&node)), vec![4]);
            node.inputs[4].source = Src::Intermediate;
            assert!(reads_unguarded_parameter(&node, true));
            node.inputs[4].shape_const = true;
            assert!(!reads_unguarded_parameter(&node, true));
        }
    }
}