car-inference 0.56.1

//! Text generation through `mlx-swift-lm`, Apple's own Swift MLX stack.
//!
//! This is the replacement for the hand-written `backend::mlx` path. Two
//! measured reasons, on the same `Qwen3-0.6B-6bit` checkpoint and machine:
//!
//! | | prefill @2048 | architectures |
//! |---|---|---|
//! | `backend::mlx` (mlx-rs) | 5,831 tok/s | 3, hand-written |
//! | `mlx-swift-lm` | 15,898 tok/s | 60+, upstream |
//!
//! The Rust binding calls itself unofficial, ships no model implementations,
//! and runs ~2.7x slower end to end for reasons still open (see
//! `docs/solutions/mlx-rs-matmul-is-3x-slower-than-python-mlx.md`). The Swift
//! package is maintained in Apple's `ml-explore` org alongside MLX itself,
//! tracks new architectures within weeks, and loads `mlx-community` repos by
//! id — which is exactly where CAR's checkpoints already come from.
//!
//! The division of labour matches the backend it replaces: Swift owns loading
//! and the forward pass, and everything else — sampling, chat templating,
//! routing, policy, the KV-transfer surface — stays in Rust.

use crate::InferenceError;
#[cfg(car_mlxlm_swift_built)]
use std::ffi::CString;
use std::path::Path;

#[cfg(car_mlxlm_swift_built)]
mod ffi {
    use std::ffi::c_char;
    extern "C" {
        pub fn car_mlxlm_load(path: *const c_char) -> i32;
        pub fn car_mlxlm_free(handle: i32);
        pub fn car_mlxlm_clear_cache(handle: i32) -> i32;
        pub fn car_mlxlm_forward(
            handle: i32,
            tokens: *const i32,
            count: i32,
            pos: i32,
            out: *mut f32,
            out_cap: i32,
        ) -> i32;
        pub fn car_mlxlm_begin_prompt(handle: i32, tokens: *const i32, count: i32) -> i32;
        pub fn car_mlxlm_last_error(buf: *mut c_char, cap: i32) -> i32;
        pub fn car_mlxlm_supports(model_type: *const c_char) -> i32;
        pub fn car_mlxlm_embed_load(path: *const c_char) -> i32;
        pub fn car_mlxlm_embed_free(handle: i32);
        pub fn car_mlxlm_embed(
            handle: i32,
            tokens: *const i32,
            count: i32,
            out: *mut f32,
            out_cap: i32,
        ) -> i32;
    }
}

/// Whether this build carries the Swift stack.
///
/// The shim needs a full Swift toolchain and fetches a SwiftPM package, so a
/// build without it is expected rather than exceptional — same posture as the
/// FoundationModels shim, and the reason callers must check before routing.
pub fn is_available() -> bool {
    cfg!(car_mlxlm_swift_built)
}

/// Whether the linked `mlx-swift-lm` registers a loader for this `config.json`
/// `model_type`.
///
/// Asking the library is the point. The predicate this replaces read a
/// hand-maintained Rust constant, which could only ever be as current as the
/// last person to edit it — and after the Swift migration it still named the
/// four architectures the deleted Rust loaders covered while the linked package
/// implemented sixty, so every other architecture was refused admission to a
/// backend that would have served it.
///
/// `false` when the Swift stack is not in this build, which is the honest
/// answer: an unbuilt shim loads nothing, and a caller told otherwise would
/// route to a stub.
///
/// Two limits worth knowing before widening this. `contains` checks that a
/// creator is registered for the key, while `createModel` additionally filters
/// on per-entry `matches(configuration)` predicates and can still throw
/// `unsupportedModelType` — so this over-approximates, and the whole point of
/// asking before a multi-gigabyte download is lost if that gap grows. It is
/// empty today because the only predicate registrations live in `MLXVLM` and
/// this asks `LLMTypeRegistry`; do not widen the query to the VLM registry
/// without handling it. And nothing in CI exercises this call at all — no
/// workflow uses a macOS runner, so `car_mlxlm_swift_built` is never set there.
/// A signature change is a loud Swift compile error on an operator Mac; a
/// semantic change to `contains` would be silent everywhere.
pub fn supports_model_type(model_type: &str) -> bool {
    #[cfg(car_mlxlm_swift_built)]
    {
        let Ok(c) = CString::new(model_type.to_ascii_lowercase()) else {
            // An interior NUL cannot be a `model_type`; it also cannot cross
            // the C ABI, so refuse rather than truncate to a prefix that might
            // spuriously match.
            return false;
        };
        unsafe { ffi::car_mlxlm_supports(c.as_ptr()) == 1 }
    }
    #[cfg(not(car_mlxlm_swift_built))]
    {
        let _ = model_type;
        false
    }
}

#[cfg(car_mlxlm_swift_built)]
fn last_error() -> String {
    let mut buf = vec![0i8; 512];
    let n = unsafe { ffi::car_mlxlm_last_error(buf.as_mut_ptr(), buf.len() as i32) };
    if n <= 0 {
        return "unknown error".to_string();
    }
    let bytes: Vec<u8> = buf[..n as usize].iter().map(|c| *c as u8).collect();
    String::from_utf8_lossy(&bytes).into_owned()
}

/// A loaded model. Dropping it releases the Swift-side session.
///
/// Tokenization and the stop-token set live here rather than in Swift on
/// purpose: the Swift shim loads a stub tokenizer, and CAR already has the
/// real one in Rust. That keeps the C ABI to the one thing Swift is here for —
/// the forward pass — and leaves everything the rest of the engine reasons
/// about (ids, eos, context window) on the Rust side, where `MlxBackend` had
/// it too.
pub struct SwiftLmBackend {
    #[cfg(car_mlxlm_swift_built)]
    handle: i32,
    #[cfg(car_mlxlm_swift_built)]
    vocab_size: usize,
    tokenizer: tokenizers::Tokenizer,
    context_length: usize,
    eos: Vec<u32>,
    /// The checkpoint's own `chat_template.jinja`, when it ships one.
    ///
    /// This is what makes support for 60+ architectures real rather than
    /// nominal. The hand-written backends this replaces hardcoded one prompt
    /// format each — Qwen3's in the trait default, Gemma's in an override —
    /// so a Swift backend that can LOAD any architecture would still have
    /// rendered every one of them as Qwen3 without this.
    chat_template: Option<crate::tasks::chat_template::ChatTemplate>,
    /// Lowercased `model_type` from config.json, for tool-call conventions,
    /// which are per-architecture and not carried in the chat template.
    model_type: String,
    /// Where this checkpoint lives, so the embedder can be loaded from it on
    /// demand without threading the path through every caller.
    #[cfg(car_mlxlm_swift_built)]
    model_dir: std::path::PathBuf,
    /// What this backend was loaded for; decides which handle exists.
    role: BackendRole,
    /// `MLXEmbedders` handle. Set at load time for [`BackendRole::Embedding`],
    /// and `-1` until first use for [`BackendRole::Text`].
    ///
    /// Separate from `handle` because embeddings come from a different Swift
    /// factory with its own registry and a pooling step; the LLM container
    /// exposes logits, which are not an embedding. Lazy because the overwhelming
    /// majority of loaded backends never embed — an embedding checkpoint
    /// declares `capabilities: ["embed"]` and is a different catalog row — so
    /// paying a second weight load at construction would be wrong for every
    /// generation model.
    #[cfg(car_mlxlm_swift_built)]
    embed_handle: std::cell::Cell<i32>,
}

/// Collect every token id that should stop generation, from the places a
/// checkpoint may declare them.
///
/// Pure and injected, because the bug it fixes cannot be caught any other way
/// on a machine that only holds Qwen3 checkpoints. This used to be two Qwen
/// literals — `<|endoftext|>` and `<|im_end|>` — with a comment claiming parity
/// with the deleted `MlxBackend`. It was not parity: that backend also fell
/// back to `</s>`, and `Gemma4Backend` read `config.json`'s `eos_token_id` as
/// scalar **or** array. So every non-Qwen architecture got an EMPTY stop set
/// and decoded to `max_tokens`, emitting a hallucinated next turn — while
/// `has_native_backend` happily admitted it, because `mlx-swift-lm` registers
/// sixty of them.
///
/// Every source is additive and deduped rather than first-wins: a checkpoint
/// commonly declares `eos_token_id` for the sequence end *and* a distinct
/// chat-turn ender that only the template knows about, and stopping on either
/// is correct. `generation_config.json` is consulted first because it is the
/// field HF generation actually reads.
fn derive_eos_ids(
    generation_config: Option<&serde_json::Value>,
    config: Option<&serde_json::Value>,
    tokenizer_config: Option<&serde_json::Value>,
    lookup: impl Fn(&str) -> Option<u32>,
) -> Vec<u32> {
    use serde_json::Value;
    let mut eos: Vec<u32> = Vec::new();
    let push = |id: u32, eos: &mut Vec<u32>| {
        if !eos.contains(&id) {
            eos.push(id);
        }
    };

    // `eos_token_id`: scalar or array, in either file.
    for source in [generation_config, config].into_iter().flatten() {
        match source.get("eos_token_id") {
            Some(Value::Number(n)) => {
                if let Some(id) = n.as_u64().and_then(|v| u32::try_from(v).ok()) {
                    push(id, &mut eos);
                }
            }
            Some(Value::Array(items)) => {
                for id in items
                    .iter()
                    .filter_map(Value::as_u64)
                    .filter_map(|v| u32::try_from(v).ok())
                {
                    push(id, &mut eos);
                }
            }
            _ => {}
        }
    }

    // `tokenizer_config.json` spells it as a token, either a bare string or an
    // `AddedToken` object. Gemma's `<end_of_turn>` arrives this way.
    // Deliberately not `pad_token`: it equals the eos token in some families
    // and is a distinct filler token in others, where stopping on it would
    // truncate output.
    let named = tokenizer_config
        .and_then(|c| c.get("eos_token"))
        .and_then(|v| match v {
            Value::String(s) => Some(s.as_str()),
            Value::Object(o) => o.get("content").and_then(Value::as_str),
            _ => None,
        });
    if let Some(id) = named.and_then(&lookup) {
        push(id, &mut eos);
    }

    // Additive fallback for checkpoints that declare nothing machine-readable.
    // Kept last and kept short: these are guesses, and a guess that collides
    // with a live token would truncate output.
    for token in [
        "<|endoftext|>",
        "<|im_end|>",
        "</s>",
        "<end_of_turn>",
        "<|eot_id|>",
        "<|end|>",
        "<|return|>",
    ] {
        if let Some(id) = lookup(token) {
            push(id, &mut eos);
        }
    }
    eos
}

#[cfg(test)]
mod eos_derivation_tests {
    use super::derive_eos_ids;
    use serde_json::json;

    /// A stand-in vocabulary. Returns an id only for tokens the "tokenizer"
    /// knows, so a lookup for a token from another family yields `None` exactly
    /// as it would on a real checkpoint.
    fn vocab(known: &'static [(&'static str, u32)]) -> impl Fn(&str) -> Option<u32> {
        move |t: &str| known.iter().find(|(name, _)| *name == t).map(|(_, id)| *id)
    }

    /// The regression. Every one of these architectures produced an EMPTY stop
    /// set under the two-Qwen-literal derivation, so `drive_generation` never
    /// broke and decoded to `max_tokens`.
    ///
    /// Synthetic configs on purpose: this machine holds only `qwen3`
    /// checkpoints, which is exactly why the bug shipped. The shapes are the
    /// real ones — Gemma declares an `eos_token_id` array, Llama 3 a scalar
    /// plus an `AddedToken` object, Mistral neither.
    #[test]
    fn non_qwen_architectures_get_a_stop_token() {
        // Gemma 3/4: array of ids, and `<end_of_turn>` is the turn ender.
        let gemma = derive_eos_ids(
            Some(&json!({ "eos_token_id": [1, 106] })),
            Some(&json!({ "eos_token_id": 1 })),
            Some(&json!({ "eos_token": { "content": "<end_of_turn>" } })),
            vocab(&[("<end_of_turn>", 106), ("<eos>", 1)]),
        );
        assert!(gemma.contains(&106), "gemma turn ender missing: {gemma:?}");
        assert!(gemma.contains(&1), "gemma sequence eos missing: {gemma:?}");

        // Llama 3: scalar in generation_config, `<|eot_id|>` as the turn ender.
        let llama = derive_eos_ids(
            Some(&json!({ "eos_token_id": 128009 })),
            Some(&json!({})),
            Some(&json!({ "eos_token": "<|eot_id|>" })),
            vocab(&[("<|eot_id|>", 128009)]),
        );
        assert_eq!(llama, vec![128009], "llama: {llama:?}");

        // Mistral: declares nothing machine-readable; the literal fallback is
        // the only thing standing between it and an unbounded decode.
        let mistral = derive_eos_ids(None, Some(&json!({})), None, vocab(&[("</s>", 2)]));
        assert_eq!(mistral, vec![2], "mistral: {mistral:?}");
    }

    /// Qwen must be unchanged — it is the one family that already worked, and a
    /// fix that moved it would be a regression traded for a regression.
    #[test]
    fn qwen_derivation_is_unchanged() {
        let qwen = derive_eos_ids(
            Some(&json!({ "eos_token_id": 151645 })),
            Some(&json!({ "eos_token_id": 151643 })),
            Some(&json!({ "eos_token": "<|im_end|>" })),
            vocab(&[("<|endoftext|>", 151643), ("<|im_end|>", 151645)]),
        );
        assert!(qwen.contains(&151643) && qwen.contains(&151645), "{qwen:?}");
    }

    /// Sources are additive and deduped: the same id declared in the array, the
    /// scalar, the tokenizer config and the literal fallback yields one entry.
    #[test]
    fn sources_union_without_duplicates() {
        let ids = derive_eos_ids(
            Some(&json!({ "eos_token_id": [7, 7] })),
            Some(&json!({ "eos_token_id": 7 })),
            Some(&json!({ "eos_token": "</s>" })),
            vocab(&[("</s>", 7)]),
        );
        assert_eq!(ids, vec![7], "duplicates leaked: {ids:?}");
    }

    /// And additive really means additive: a turn ender the tokenizer knows is
    /// kept alongside a differently-numbered declared sequence eos, because
    /// stopping on either is correct and a chat model emits the turn ender.
    #[test]
    fn a_turn_ender_is_kept_alongside_the_sequence_eos() {
        let ids = derive_eos_ids(
            Some(&json!({ "eos_token_id": 7 })),
            None,
            None,
            vocab(&[("</s>", 7), ("<end_of_turn>", 9)]),
        );
        assert!(ids.contains(&7) && ids.contains(&9), "{ids:?}");
    }

    /// The negative control. A checkpoint whose tokenizer knows none of the
    /// fallback literals and which declares nothing yields an empty set — the
    /// case `load` warns about. Without this the assertions above would pass
    /// against a function that returned a fixed non-empty list.
    #[test]
    fn nothing_declared_and_nothing_known_yields_empty() {
        let ids = derive_eos_ids(None, None, None, vocab(&[("some_unrelated_token", 3)]));
        assert!(ids.is_empty(), "expected no stop tokens, got {ids:?}");
    }

    /// Malformed declarations must not panic or invent ids.
    #[test]
    fn malformed_declarations_are_ignored() {
        let ids = derive_eos_ids(
            Some(&json!({ "eos_token_id": "not-a-number" })),
            Some(&json!({ "eos_token_id": [ -1, 4.5, null ] })),
            Some(&json!({ "eos_token": { "no_content_field": true } })),
            vocab(&[]),
        );
        assert!(ids.is_empty(), "{ids:?}");
    }
}

/// Which set of weights a backend load should actually bring into memory.
///
/// A checkpoint like `Qwen3-Embedding-0.6B` is loadable two ways: as a causal
/// LM through `LLMModelFactory`, and as an embedder through
/// `EmbedderModelFactory`, which adds the checkpoint's pooling step. They are
/// separate handles over separate copies of the weights.
///
/// Before this existed, the embedding path loaded BOTH — `ensure_text_backend`
/// brought up the LM, then the first `embed_one` lazily brought up the embedder
/// beside it. For an embedding-only catalog row the LM copy was never used, and
/// the admission reservation had been sized for one copy, so the coordinator
/// was told about half the memory actually taken. Choosing at load time fixes
/// the accounting and the waste together.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendRole {
    /// Generation, streaming, tokenize. Loads the LM.
    Text,
    /// Embeddings only. Loads the embedder and never the LM.
    Embedding,
}

use crate::resource_policy::NESTED_TEXT_CONFIGS;

/// A numeric setting from the top level of `config.json`, else from the
/// language model's nested config. One lookup for every such setting, so a
/// new multimodal layout cannot leave one of them unread.
fn text_config_u64(config: &serde_json::Value, key: &str) -> Option<u64> {
    config[key].as_u64().or_else(|| {
        NESTED_TEXT_CONFIGS
            .iter()
            .find_map(|nested| config.get(*nested)?.get(key)?.as_u64())
    })
}

/// The checkpoint's vocabulary size, 0 when it declares none. Reading only the
/// top level made every multimodal model (`gemma4_unified`) fail to load.
fn config_vocab_size(config: &serde_json::Value) -> usize {
    text_config_u64(config, "vocab_size").unwrap_or(0) as usize
}

#[cfg(test)]
mod vocab_size_tests {
    use super::config_vocab_size;
    use serde_json::json;

    #[test]
    fn a_multimodal_checkpoint_declares_its_vocabulary_under_text_config() {
        // Shape of mlx-community/gemma-4-12B-it-4bit's config.json.
        let gemma4 = json!({
            "model_type": "gemma4_unified",
            "text_config": {"model_type": "gemma4_unified_text", "vocab_size": 262144},
            "vision_config": {},
        });
        assert_eq!(config_vocab_size(&gemma4), 262144);
        assert_eq!(config_vocab_size(&json!({"vocab_size": 151936})), 151936);
        // The top level wins when both are present.
        assert_eq!(
            config_vocab_size(&json!({"vocab_size": 7, "text_config": {"vocab_size": 9}})),
            7
        );
        assert_eq!(config_vocab_size(&json!({})), 0);
        // Other multimodal layouts nest it under their own key.
        assert_eq!(
            config_vocab_size(&json!({"llm_config": {"vocab_size": 92553}})),
            92553
        );
        assert_eq!(
            config_vocab_size(&json!({"language_config": {"vocab_size": 102400}})),
            102400
        );
    }
}

impl SwiftLmBackend {
    /// Load a checkpoint directory — the same `mlx-community` layout the
    /// deleted Rust backends read, so installed models need no conversion.
    /// Load for generation. See [`BackendRole`] for why the embedding path
    /// takes [`Self::load_with_role`] instead.
    pub fn load(model_dir: &Path) -> Result<Self, InferenceError> {
        Self::load_with_role(model_dir, BackendRole::Text)
    }

    pub fn load_with_role(model_dir: &Path, role: BackendRole) -> Result<Self, InferenceError> {
        let tokenizer = tokenizers::Tokenizer::from_file(model_dir.join("tokenizer.json"))
            .map_err(|e| InferenceError::TokenizationError(format!("tokenizer.json: {e}")))?;
        let config: serde_json::Value = serde_json::from_str(
            &std::fs::read_to_string(model_dir.join("config.json"))
                .map_err(|e| InferenceError::InferenceFailed(format!("read config.json: {e}")))?,
        )
        .map_err(|e| InferenceError::InferenceFailed(format!("parse config.json: {e}")))?;
        let vocab_size = config_vocab_size(&config);
        if vocab_size == 0 {
            return Err(InferenceError::InferenceFailed(
                "config.json has no vocab_size".to_string(),
            ));
        }
        // Architectures spell the window differently, and a silent 32768 is
        // wrong in both directions: it truncates a long-context model's prompt
        // and overruns a short one. Check the nested text config too — that is
        // where multimodal checkpoints put it.
        let context_length = ["max_position_embeddings", "seq_length", "n_positions"]
            .iter()
            .find_map(|k| text_config_u64(&config, k))
            .map(|v| v as usize)
            .unwrap_or_else(|| {
                tracing::warn!(
                    dir = %model_dir.display(),
                    "config.json declares no context window under any known key; \
                     assuming 32768, which may truncate prompts or overrun the model"
                );
                32768
            });
        let model_type = config["model_type"]
            .as_str()
            .unwrap_or_default()
            .to_ascii_lowercase();
        let chat_template = crate::tasks::chat_template::ChatTemplate::load(model_dir)?;
        let read_json = |name: &str| -> Option<serde_json::Value> {
            serde_json::from_str(&std::fs::read_to_string(model_dir.join(name)).ok()?).ok()
        };
        let eos = derive_eos_ids(
            read_json("generation_config.json").as_ref(),
            Some(&config),
            read_json("tokenizer_config.json").as_ref(),
            |t| tokenizer.token_to_id(t),
        );
        if eos.is_empty() {
            // Not fatal — `max_tokens` still bounds the loop — but every
            // response will run to the cap and trail a hallucinated next turn,
            // so name the model rather than let it look like a slow model.
            tracing::warn!(
                model_type = %model_type,
                dir = %model_dir.display(),
                "no EOS token could be derived; generation will stop only at max_tokens"
            );
        }
        Self::load_inner(
            model_dir,
            role,
            vocab_size,
            tokenizer,
            context_length,
            eos,
            chat_template,
            model_type,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn load_inner(
        model_dir: &Path,
        role: BackendRole,
        vocab_size: usize,
        tokenizer: tokenizers::Tokenizer,
        context_length: usize,
        eos: Vec<u32>,
        chat_template: Option<crate::tasks::chat_template::ChatTemplate>,
        model_type: String,
    ) -> Result<Self, InferenceError> {
        #[cfg(not(car_mlxlm_swift_built))]
        {
            let _ = (
                model_dir,
                role,
                vocab_size,
                tokenizer,
                context_length,
                eos,
                chat_template,
                model_type,
            );
            Err(InferenceError::InferenceFailed(
                "this build has no mlx-swift-lm backend (needs a full Swift toolchain and \
                 CAR_BUILD_SWIFT_LM at build time)"
                    .to_string(),
            ))
        }
        #[cfg(car_mlxlm_swift_built)]
        {
            let path = CString::new(model_dir.to_string_lossy().as_bytes()).map_err(|e| {
                InferenceError::InferenceFailed(format!("model path is not a C string: {e}"))
            })?;
            // One factory, not both. `BackendRole` explains why.
            let (handle, embed_handle) = match role {
                BackendRole::Text => {
                    let h = unsafe { ffi::car_mlxlm_load(path.as_ptr()) };
                    if h < 0 {
                        return Err(InferenceError::InferenceFailed(format!(
                            "mlx-swift-lm failed to load {}: {}",
                            model_dir.display(),
                            last_error()
                        )));
                    }
                    // Embedder stays lazy: a text backend asked to embed is
                    // unusual but reachable, and it should not pay for weights
                    // it will probably never use.
                    (h, -1)
                }
                BackendRole::Embedding => {
                    let h = unsafe { ffi::car_mlxlm_embed_load(path.as_ptr()) };
                    if h < 0 {
                        return Err(InferenceError::InferenceFailed(format!(
                            "mlx-swift-lm could not load {} as an embedding model: {}",
                            model_dir.display(),
                            last_error()
                        )));
                    }
                    // No LM handle at all. Every text entry point checks.
                    (-1, h)
                }
            };
            Ok(Self {
                handle,
                vocab_size,
                tokenizer,
                context_length,
                eos,
                chat_template,
                model_type,
                model_dir: model_dir.to_path_buf(),
                embed_handle: std::cell::Cell::new(embed_handle),
                role,
            })
        }
    }

    /// Whether this backend holds LM weights. False for an embedding-only load.
    pub fn serves_text(&self) -> bool {
        self.role == BackendRole::Text
    }

    fn require_text(&self) -> Result<(), InferenceError> {
        if self.role != BackendRole::Text {
            return Err(InferenceError::UnsupportedMode {
                mode: "text-generation",
                backend: "mlx-swift-lm-embedder",
                reason: "this backend was loaded for embeddings only, so it holds no \
                         language-model weights; route generation to a text model",
            });
        }
        Ok(())
    }

    pub fn forward(&mut self, tokens: &[u32], pos: usize) -> Result<Vec<f32>, InferenceError> {
        self.require_text()?;
        #[cfg(not(car_mlxlm_swift_built))]
        {
            let _ = (tokens, pos);
            Err(InferenceError::InferenceFailed(
                "this build has no mlx-swift-lm backend".to_string(),
            ))
        }
        #[cfg(car_mlxlm_swift_built)]
        {
            // Match mlx-swift-lm's default prefill step size. A single large
            // Qwen3 MoE forward produced unrelated output on a 5,322-token
            // coding conversation, while bounded prefill of the same tokens
            // produced the requested task. Bound every caller (streaming,
            // ordinary generation, and scoring), including a reused suffix.
            const PREFILL_STEP: usize = 512;
            if tokens.len() > PREFILL_STEP {
                let mut logits = Vec::new();
                for (index, chunk) in tokens.chunks(PREFILL_STEP).enumerate() {
                    logits = self.forward(chunk, pos + index * PREFILL_STEP)?;
                }
                return Ok(logits);
            }
            let ids: Vec<i32> = tokens.iter().map(|t| *t as i32).collect();
            let mut out = vec![0f32; self.vocab_size];
            let n = unsafe {
                ffi::car_mlxlm_forward(
                    self.handle,
                    ids.as_ptr(),
                    ids.len() as i32,
                    pos as i32,
                    out.as_mut_ptr(),
                    out.len() as i32,
                )
            };
            if n < 0 {
                return Err(InferenceError::InferenceFailed(format!(
                    "mlx-swift-lm forward failed: {}",
                    last_error()
                )));
            }
            out.truncate(n as usize);
            Ok(out)
        }
    }

    /// Reuse the cached KV of the longest shared prefix; returns the position
    /// to prefill from. See `MlxBackend::begin_prompt`.
    pub fn begin_prompt(&mut self, tokens: &[u32]) -> usize {
        if self.role != BackendRole::Text {
            return 0;
        }
        #[cfg(not(car_mlxlm_swift_built))]
        {
            let _ = tokens;
            0
        }
        #[cfg(car_mlxlm_swift_built)]
        {
            let ids: Vec<i32> = tokens.iter().map(|t| *t as i32).collect();
            let n =
                unsafe { ffi::car_mlxlm_begin_prompt(self.handle, ids.as_ptr(), ids.len() as i32) };
            n.max(0) as usize
        }
    }

    /// Tokenize WITHOUT the tokenizer's special tokens.
    ///
    /// Distinct from the `TextDecoder::encode` path, which adds them. Callers
    /// that splice sequences together need the raw ids.
    pub fn tokenize_raw(&self, text: &str) -> Result<Vec<u32>, InferenceError> {
        self.tokenizer
            .encode(text, false)
            .map(|e| e.get_ids().to_vec())
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))
    }

    /// Detokenize without skipping special tokens, so callers see exactly what
    /// is in the sequence.
    pub fn detokenize_raw(&self, tokens: &[u32]) -> Result<String, InferenceError> {
        self.tokenizer
            .decode(tokens, false)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))
    }

    pub fn encode(&self, text: &str) -> Result<Vec<u32>, InferenceError> {
        self.tokenizer
            .encode(text, true)
            .map(|e| e.get_ids().to_vec())
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))
    }

    pub fn decode(&self, tokens: &[u32]) -> Result<String, InferenceError> {
        self.tokenizer
            .decode(tokens, true)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))
    }

    pub fn eos_token_id(&self) -> Option<u32> {
        self.eos.first().copied()
    }

    pub fn token_id(&self, token: &str) -> Option<u32> {
        self.tokenizer.token_to_id(token)
    }

    pub fn context_length(&self) -> usize {
        self.context_length
    }

    /// What this loaded checkpoint can actually service.
    pub fn supports_capability(&self, cap: crate::schema::ModelCapability) -> bool {
        <Self as crate::backend::local::LocalInferenceBackend>::supports_capability(self, cap)
    }

    /// Embeddings are not served by this backend.
    ///
    /// `MlxBackend` pooled the last-token hidden state itself. mlx-swift-lm
    /// exposes logits through the C ABI, not hidden states, and Apple ships a
    /// separate `MLXEmbedders` package for this — wiring that is its own
    /// change rather than something to fake here. Returning an error keeps the
    /// caller's fallback path honest instead of handing back a pooled logit
    /// vector that is not an embedding.
    pub fn embed_one(&mut self, text: &str) -> Result<Vec<f32>, InferenceError> {
        #[cfg(not(car_mlxlm_swift_built))]
        {
            let _ = text;
            Err(InferenceError::InferenceFailed(
                "this build has no mlx-swift-lm backend, so it serves no embeddings".to_string(),
            ))
        }
        #[cfg(car_mlxlm_swift_built)]
        {
            let handle = self.ensure_embedder()?;
            let ids = self.tokenize_raw(text)?;
            if ids.is_empty() {
                // An empty input is not an error, and a zero vector is the
                // conventional answer; every caller here feeds it to cosine
                // similarity, which treats it as maximally dissimilar.
                return Ok(vec![0.0; self.embed_dimensions()?]);
            }
            let tokens: Vec<i32> = ids.iter().map(|&t| t as i32).collect();
            // Sized from the model's own hidden width, not a guess: the Swift
            // side refuses rather than truncating if it disagrees.
            let mut out = vec![0.0f32; self.embed_dimensions()?];
            let n = unsafe {
                ffi::car_mlxlm_embed(
                    handle,
                    tokens.as_ptr(),
                    tokens.len() as i32,
                    out.as_mut_ptr(),
                    out.len() as i32,
                )
            };
            if n < 0 {
                return Err(InferenceError::InferenceFailed(format!(
                    "mlx-swift-lm embedding failed: {}",
                    last_error()
                )));
            }
            out.truncate(n as usize);
            Ok(out)
        }
    }

    /// Qwen3-Embedding's documented query form. The instruction is part of the
    /// text the model sees, not a separate channel — same as the deleted
    /// backend did, so stored document vectors stay comparable.
    pub fn embed_query(
        &mut self,
        text: &str,
        instruction: &str,
    ) -> Result<Vec<f32>, InferenceError> {
        self.embed_one(&format!("Instruct: {instruction}\nQuery: {text}"))
    }

    /// The embedding width, read from the checkpoint rather than assumed.
    #[cfg(car_mlxlm_swift_built)]
    fn embed_dimensions(&self) -> Result<usize, InferenceError> {
        let config: serde_json::Value = serde_json::from_str(
            &std::fs::read_to_string(self.model_dir.join("config.json"))
                .map_err(|e| InferenceError::InferenceFailed(format!("read config.json: {e}")))?,
        )
        .map_err(|e| InferenceError::InferenceFailed(format!("parse config.json: {e}")))?;
        config
            .get("hidden_size")
            .and_then(serde_json::Value::as_u64)
            .map(|v| v as usize)
            .ok_or_else(|| {
                InferenceError::InferenceFailed(
                    "config.json declares no hidden_size, so the embedding width is unknown"
                        .to_string(),
                )
            })
    }

    /// Load the embedder on first use and cache the handle.
    #[cfg(car_mlxlm_swift_built)]
    fn ensure_embedder(&self) -> Result<i32, InferenceError> {
        let existing = self.embed_handle.get();
        if existing > 0 {
            return Ok(existing);
        }
        let path = CString::new(self.model_dir.to_string_lossy().as_bytes()).map_err(|e| {
            InferenceError::InferenceFailed(format!("model path is not a C string: {e}"))
        })?;
        let handle = unsafe { ffi::car_mlxlm_embed_load(path.as_ptr()) };
        if handle < 0 {
            return Err(InferenceError::InferenceFailed(format!(
                "mlx-swift-lm could not load {} as an embedding model: {}",
                self.model_dir.display(),
                last_error()
            )));
        }
        self.embed_handle.set(handle);
        Ok(handle)
    }

    /// Reset the KV cache. A no-op on an embedding backend, which has none.
    ///
    /// The guard matters because `handle` is `-1` for
    /// [`BackendRole::Embedding`] and `-1` is not a valid handle. Every other
    /// use of `handle` is already behind `require_text`, `serves_text`, or a
    /// `> 0` check; this was the last one that passed the sentinel across the
    /// C ABI, where it was benign only because Swift's dictionary lookup
    /// misses and the entry point guards. That is defence in depth standing in
    /// for a type — `Option<i32>` would make each of these a compile error
    /// instead of six restatements of one invariant, and is the better shape
    /// if this struct is touched again.
    pub fn clear_kv_cache(&mut self) {
        #[cfg(car_mlxlm_swift_built)]
        {
            if !self.serves_text() {
                return;
            }
            unsafe {
                ffi::car_mlxlm_clear_cache(self.handle);
            }
        }
    }
}

impl Drop for SwiftLmBackend {
    fn drop(&mut self) {
        #[cfg(car_mlxlm_swift_built)]
        unsafe {
            if self.handle > 0 {
                ffi::car_mlxlm_free(self.handle);
            }
            let embed = self.embed_handle.get();
            if embed > 0 {
                ffi::car_mlxlm_embed_free(embed);
            }
        }
    }
}

/// Does prompt-prefix reuse actually fire on a second chat turn?
///
/// A changelog fragment claims 2.4x on a third turn, measured on the deleted
/// Rust backend. The question is whether a real second turn is a strict
/// extension of what the cache holds, and it is answerable here instead of by
/// reasoning about it.
///
/// (An earlier version of this comment said `mlx-swift-lm`'s `KVCache` exposes
/// no truncate. It does — `trim(_:)` and `trimPromptCache(_:numTokens:)` — and
/// `begin_prompt` now uses them.)
#[cfg(all(test, car_mlxlm_swift_built))]
mod swift_lm_prefix_reuse_tests {
    use super::*;

    fn argmax(v: &[f32]) -> usize {
        v.iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap()
    }

    #[test]
    fn batched_prefill_preserves_positions_after_prefix_reuse() {
        let Some(dir) = std::env::var_os("HOME")
            .map(std::path::PathBuf::from)
            .map(|home| home.join(".car/models/Qwen3-0.6B-MLX"))
            .filter(|dir| dir.join("config.json").exists())
        else {
            eprintln!("SKIP: Qwen3-0.6B-MLX is not installed");
            return;
        };
        let mut backend = SwiftLmBackend::load(&dir).unwrap();
        let prefix = backend
            .tokenize_raw("The following is background context: ")
            .unwrap();
        let mut primer = prefix.clone();
        primer.extend(
            backend
                .tokenize_raw("a previous unrelated question.")
                .unwrap(),
        );
        backend.forward(&primer, 0).unwrap();
        let mut target = prefix;
        target.extend(
            backend
                .tokenize_raw(&"Background data. ".repeat(400))
                .unwrap(),
        );
        target.extend(backend.tokenize_raw("The capital of France is").unwrap());
        let reused = backend.begin_prompt(&target);
        assert!(reused > 0 && target.len() - reused > 1024);
        let warm = backend.forward(&target[reused..], reused).unwrap();
        backend.clear_kv_cache();
        let cold = backend.forward(&target, 0).unwrap();
        assert_eq!(warm.len(), cold.len());
        assert!(warm.iter().chain(&cold).all(|value| value.is_finite()));
        assert_eq!(
            argmax(&warm),
            argmax(&cold),
            "batched suffix changed the prediction"
        );
        eprintln!(
            "batched prefix test: {} tokens, {reused} reused, matching top-1",
            target.len()
        );
    }

    #[test]
    fn measures_whether_a_second_turn_reuses_its_prefix() {
        let Some(dir) = std::env::var_os("HOME")
            .map(std::path::PathBuf::from)
            .map(|h| h.join(".car/models/Qwen3-0.6B-MLX"))
            .filter(|d| d.join("config.json").exists())
        else {
            eprintln!("SKIP swift_lm_prefix_reuse_tests: Qwen3-0.6B-MLX not in ~/.car/models");
            return;
        };
        let mut backend = SwiftLmBackend::load(&dir).expect("load");

        // Turn one: prefill a prompt, then decode a few tokens the way
        // `drive_generation` does — each sampled token fed back through
        // `forward`, which is what appends it to the Swift-side cache.
        let turn1: Vec<u32> = backend.tokenize_raw("The capital of France is").unwrap();
        let start = backend.begin_prompt(&turn1);
        assert_eq!(start, 0, "a cold cache must prefill from zero");
        let mut pos = 0usize;
        let logits = backend.forward(&turn1, pos).unwrap();
        pos += turn1.len();
        let mut generated = Vec::new();
        for _ in 0..4 {
            let next = logits
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
                .map(|(i, _)| i as u32)
                .unwrap();
            generated.push(next);
            let _ = backend.forward(&[next], pos).unwrap();
            pos += 1;
        }

        // Turn two, as a chat template renders it: the earlier prompt, the
        // assistant's reply, and turn markers between them.
        let mut turn2 = turn1.clone();
        turn2.extend_from_slice(&generated);
        turn2.extend(backend.tokenize_raw("\nUser: and of Italy?").unwrap());
        let reuse_realistic = backend.begin_prompt(&turn2);

        // And the shape prefix reuse is actually built for: a strict extension
        // of exactly what the cache holds, with nothing inserted.
        let mut backend2 = SwiftLmBackend::load(&dir).expect("load");
        let base: Vec<u32> = backend2.tokenize_raw("The capital of France is").unwrap();
        backend2.begin_prompt(&base);
        backend2.forward(&base, 0).unwrap();
        let mut extended = base.clone();
        extended.extend(backend2.tokenize_raw(" Paris and").unwrap());
        let reuse_strict = backend2.begin_prompt(&extended);

        eprintln!(
            "prefix reuse: realistic second turn reused {reuse_realistic} of {} tokens; \
             strict extension reused {reuse_strict} of {}",
            turn2.len(),
            extended.len()
        );

        assert!(
            reuse_strict > 0,
            "prefix reuse does not fire even on a strict extension: the mechanism is inert"
        );

        // A DIVERGENT prompt — an edited message, a retried branch — must keep
        // the shared prefix by trimming the divergent tail, not discard
        // everything. This used to return 0; `trimPromptCache` is what changed.
        //
        // Correctness first: reusing a trimmed cache has to produce the same
        // logits as a cold prefill of the same tokens, or the speedup is just a
        // faster way to be wrong.
        let mut warm = SwiftLmBackend::load(&dir).expect("load");
        let shared: Vec<u32> = warm.tokenize_raw("The capital city of").unwrap();
        let mut first = shared.clone();
        first.extend(warm.tokenize_raw(" Germany is Berlin").unwrap());
        warm.begin_prompt(&first);
        warm.forward(&first, 0).unwrap();

        let mut edited = shared.clone();
        edited.extend(warm.tokenize_raw(" France is").unwrap());
        let reused = warm.begin_prompt(&edited);
        assert!(
            reused >= shared.len(),
            "a divergent prompt reused {reused} tokens; the shared prefix is {} long",
            shared.len()
        );
        let warm_logits = warm.forward(&edited[reused..], reused).unwrap();

        let mut cold = SwiftLmBackend::load(&dir).expect("load");
        cold.begin_prompt(&edited);
        let cold_logits = cold.forward(&edited, 0).unwrap();

        assert_eq!(warm_logits.len(), cold_logits.len());
        let warm_top = argmax(&warm_logits);
        let cold_top = argmax(&cold_logits);
        assert_eq!(
            warm_top, cold_top,
            "trimmed-cache reuse changed the prediction: warm {warm_top} vs cold {cold_top}"
        );
        eprintln!(
            "divergent prompt: reused {reused} of {} tokens, top-1 unchanged",
            edited.len()
        );

        // The claim under test, and it holds. It was argued in review that a
        // second turn must reuse nothing: the cache ends up holding
        // prompt+generation, `mlx-swift-lm` cannot truncate a KV cache, so a
        // partial match resets to zero. The premise is right and the conclusion
        // does not follow, because a chat template puts the turn markers AFTER
        // the assistant's reply, not between the prompt and it — the assistant
        // header is already in turn one's prompt. So turn two is a strict
        // extension of exactly what the cache holds, and the match is total.
        //
        // Measuring settled in seconds what neither of us could settle by
        // reading the code.
        assert!(
            reuse_realistic >= turn1.len(),
            "a realistic second turn reused only {reuse_realistic} of {} tokens; \
             prefix reuse across turns has regressed",
            turn2.len()
        );
    }
}

/// An embedding-only row must not bring up the language model.
///
/// The bug this locks down: `ensure_text_backend` loaded the LM, then the first
/// `embed_one` lazily loaded the embedder beside it — two copies of the same
/// checkpoint in memory, while the admission reservation had been sized for
/// one. The coordinator was told about half of what was actually taken, which
/// is the kind of accounting error that surfaces as an OOM somewhere else.
#[cfg(all(test, car_mlxlm_swift_built))]
mod swift_lm_role_tests {
    use super::*;

    fn embedding_checkpoint() -> Option<std::path::PathBuf> {
        let dir = std::env::var_os("HOME")
            .map(std::path::PathBuf::from)?
            .join(".car/models/Qwen3-Embedding-0.6B-MLX");
        dir.join("config.json").exists().then_some(dir)
    }

    #[test]
    fn an_embedding_backend_holds_no_language_model() {
        let Some(dir) = embedding_checkpoint() else {
            eprintln!("SKIP swift_lm_role_tests: Qwen3-Embedding-0.6B-MLX not in ~/.car/models");
            return;
        };
        let mut backend =
            SwiftLmBackend::load_with_role(&dir, BackendRole::Embedding).expect("load embedder");

        assert!(!backend.serves_text(), "role should be Embedding");

        // Embedding works, and the vector is real rather than a shape.
        let v = backend.embed_one("a domestic cat on a windowsill").unwrap();
        assert!(v.iter().any(|x| x.abs() > 1e-6), "embedding is all zeros");

        // Generation refuses, loudly and by name, instead of loading weights
        // on demand and quietly doubling memory.
        let err = backend.forward(&[1, 2, 3], 0).unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("embeddings only") || msg.contains("text-generation"),
            "expected a clear refusal, got: {msg}"
        );
        assert_eq!(backend.begin_prompt(&[1, 2, 3]), 0, "no KV cache to reuse");
    }

    /// The positive control: a Text-role load of the SAME checkpoint does serve
    /// generation. Without this the assertions above would pass against a
    /// backend that was simply broken.
    #[test]
    fn a_text_backend_on_the_same_checkpoint_still_generates() {
        let Some(dir) = embedding_checkpoint() else {
            eprintln!("SKIP swift_lm_role_tests: checkpoint absent");
            return;
        };
        let mut backend =
            SwiftLmBackend::load_with_role(&dir, BackendRole::Text).expect("load as text");
        assert!(backend.serves_text());
        let logits = backend.forward(&[1, 2, 3], 0).expect("text role generates");
        assert!(!logits.is_empty(), "no logits from a text-role backend");
    }
}

/// Embeddings, against the checkpoint on disk.
///
/// The migration deleted a working embedding path (the Rust backend pooled the
/// last hidden state by hand) and replaced it with an unconditional error, and
/// nothing caught it: `car-memgine` calls `embed()` with `.ok()` and falls
/// through to keyword seeds, so retrieval silently degraded with no log line
/// and every test still passed. These assert the vector is real, not merely
/// that a call returned.
#[cfg(all(test, car_mlxlm_swift_built))]
mod swift_lm_embed_tests {
    use super::*;

    fn embedding_model() -> Option<std::path::PathBuf> {
        let dir = dirs_home()?.join(".car/models/Qwen3-Embedding-0.6B-MLX");
        dir.join("config.json").exists().then_some(dir)
    }

    fn dirs_home() -> Option<std::path::PathBuf> {
        std::env::var_os("HOME").map(std::path::PathBuf::from)
    }

    #[test]
    fn embeds_and_the_vector_means_something() {
        let Some(dir) = embedding_model() else {
            // Loud, so an always-skip cannot read as a pass.
            eprintln!("SKIP swift_lm_embed_tests: Qwen3-Embedding-0.6B-MLX not in ~/.car/models");
            return;
        };
        let mut backend = SwiftLmBackend::load(&dir).expect("load embedding checkpoint");

        let cat = backend
            .embed_one("a domestic cat sitting on a windowsill")
            .unwrap();
        let kitten = backend.embed_one("a small kitten by the window").unwrap();
        let finance = backend
            .embed_one("quarterly earnings guidance for the fiscal year")
            .unwrap();

        assert_eq!(cat.len(), kitten.len());
        assert_eq!(cat.len(), finance.len());
        assert!(cat.len() >= 512, "unexpected width {}", cat.len());

        // Not all-zero, not NaN — the failure modes a "did it return" test misses.
        assert!(
            cat.iter().any(|v| v.abs() > 1e-6),
            "embedding is all zeros: the pooling step produced nothing"
        );
        assert!(cat.iter().all(|v| v.is_finite()), "embedding has NaN/inf");

        // Normalized, so cosine similarity is a dot product.
        let norm: f32 = cat.iter().map(|v| v * v).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-2, "not L2-normalized: {norm}");

        // The property memgine actually depends on, and the one an error-
        // returning stub cannot satisfy: related text is nearer than unrelated.
        let dot = |a: &[f32], b: &[f32]| a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>();
        let near = dot(&cat, &kitten);
        let far = dot(&cat, &finance);
        assert!(
            near > far,
            "semantic ordering lost: cat~kitten {near:.4} should exceed cat~finance {far:.4}"
        );
    }

    #[test]
    fn the_same_text_embeds_identically() {
        let Some(dir) = embedding_model() else {
            eprintln!("SKIP swift_lm_embed_tests: checkpoint absent");
            return;
        };
        let mut backend = SwiftLmBackend::load(&dir).expect("load");
        let a = backend.embed_one("determinism matters").unwrap();
        let b = backend.embed_one("determinism matters").unwrap();
        assert_eq!(a.len(), b.len());
        for (x, y) in a.iter().zip(&b) {
            assert!((x - y).abs() < 1e-5, "nondeterministic embedding");
        }
    }
}

/// Parity for the Swift backend, against the same independent fixture the
/// `mlx-rs` backend is checked with.
///
/// A replacement backend that is faster and wrong is worthless, and the Gemma 4
/// loader (car#1782) is the standing proof that "it loads and returns logits"
/// is not evidence of correctness. This asserts the Swift path agrees with
/// `mlx_lm` on the same prompts, to the same one-bf16-ulp bound.
#[cfg(all(test, car_mlxlm_swift_built))]
mod swift_lm_parity_tests {
    use super::SwiftLmBackend;
    use std::path::PathBuf;

    const FIXTURE: &str = include_str!("../../tests/fixtures/qwen3-0.6b-reference-logits.json");

    fn bf16_ulp(v: f32) -> f32 {
        if v == 0.0 {
            return f32::MIN_POSITIVE;
        }
        (v.abs().log2().floor() - 7.0).exp2()
    }

    /// The fixture stores prompts, but tokenization lives in the Rust
    /// tokenizer, not the Swift shim (which loads a stub). Tokenize with the
    /// same tokenizer the mlx backend uses so the two paths see identical ids.
    fn tokenize(prompt: &str, dir: &std::path::Path) -> Option<Vec<u32>> {
        let tk = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).ok()?;
        Some(tk.encode(prompt, false).ok()?.get_ids().to_vec())
    }

    /// Swift agrees with the reference on every prompt's top-1 and to within
    /// 3 bf16 ulps on the top-5.
    ///
    /// That bound is wider than the 1 ulp the `mlx-rs` backend is held to, and
    /// it is justified by measurement rather than by convenience. The first
    /// hypothesis was that this build's TWO MLX runtimes (mlx-sys for the
    /// media backends, mlx-swift here) were perturbing it. That was tested and
    /// is FALSE: running mlx-swift alone in a separate process
    /// (`scripts/mlx-swift-bench/SwiftLM`) against the same fixture gives the
    /// same 2-3 ulp spread. mlx-swift is simply a third independent
    /// implementation of the same bf16 arithmetic, and it orders its
    /// operations differently than `mlx_lm` does.
    ///
    /// The bound still discriminates. A real defect is orders of magnitude
    /// larger: the control in the mlx-rs parity test, substituting a wrong
    /// top-1 token, fails on a gap of 16.9 — over 100x this.
    #[test]
    fn swift_backend_matches_the_mlx_lm_reference() {
        let dir = PathBuf::from(std::env::var("HOME").unwrap_or_default())
            .join(".car/models/Qwen3-0.6B-MLX");
        if !dir.join("config.json").exists() {
            eprintln!(
                "SKIPPED swift_backend_matches_the_mlx_lm_reference: {} not installed. \
This check is not running — it protects nothing on this machine.",
                dir.display()
            );
            return;
        }

        let fixture: serde_json::Value = serde_json::from_str(FIXTURE).expect("fixture parses");
        let cases = fixture["cases"].as_array().expect("cases");
        let vocab = 151_936usize;
        let mut backend = SwiftLmBackend::load(&dir).expect("load via mlx-swift-lm");

        for case in cases {
            let prompt = case["prompt"].as_str().expect("prompt");
            let Some(ids) = tokenize(prompt, &dir) else {
                eprintln!("SKIPPED: no tokenizer.json at {}", dir.display());
                return;
            };
            assert_eq!(
                ids.len() as u64,
                case["prompt_tokens"].as_u64().expect("prompt_tokens"),
                "tokenizer disagrees with the reference on {prompt:?}"
            );

            backend.clear_kv_cache();
            let logits = backend.forward(&ids, 0).expect("forward");
            assert_eq!(logits.len(), vocab, "unexpected logit count");

            let expected = case["top5"].as_array().expect("top5");
            let want_top1 = expected[0][0].as_u64().expect("id") as usize;
            let got_top1 = logits
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
                .map(|(i, _)| i)
                .expect("argmax");

            // Same tie rule as the mlx-rs parity test: bf16 logits collide, and
            // which of two equal values an argmax returns is a property of the
            // reduction rather than of the model.
            if got_top1 != want_top1 {
                let gap = (logits[got_top1] - logits[want_top1]).abs();
                assert!(
                    gap <= bf16_ulp(logits[want_top1]),
                    "{prompt:?}: Swift chose {got_top1} ({}) where the reference chose \
{want_top1} ({}) — a gap of {gap}, too large to be a tie",
                    logits[got_top1],
                    logits[want_top1]
                );
            }

            for entry in expected {
                let id = entry[0].as_u64().expect("id") as usize;
                let want = entry[1].as_f64().expect("logit") as f32;
                let ulps = (logits[id] - want).abs() / bf16_ulp(want);
                // 3 ulps, not the 1 the mlx-rs backend is held to. Measured
                // in isolation (see the doc comment) so this is a property of
                // mlx-swift, not of linking it next to another MLX.
                //
                // The bound still discriminates: a real defect is orders of
                // magnitude larger, not one step. The control that proves it
                // is in the mlx-rs parity test, where substituting a wrong
                // top-1 token fails on a gap of 16.9 — 135x this.
                assert!(
                    ulps <= 3.0,
                    "{prompt:?}: token {id} is {} against the reference's {want} — \
{ulps:.1} bf16 ulps, beyond cross-implementation rounding",
                    logits[id]
                );
            }
        }
    }
}

/// The trait the shared decode loop drives, so this backend is a drop-in for
/// `MlxBackend` rather than a parallel path with its own loop.
impl crate::backend::local::TextDecoder for SwiftLmBackend {
    fn encode(&self, text: &str) -> Result<Vec<u32>, InferenceError> {
        self.tokenizer
            .encode(text, true)
            .map(|e| e.get_ids().to_vec())
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))
    }

    fn decode(&self, tokens: &[u32]) -> Result<String, InferenceError> {
        self.tokenizer
            .decode(tokens, true)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))
    }

    fn forward(&mut self, tokens: &[u32], pos: usize) -> Result<Vec<f32>, InferenceError> {
        SwiftLmBackend::forward(self, tokens, pos)
    }

    fn eos_ids(&self) -> Vec<u32> {
        self.eos.clone()
    }

    fn context_length(&self) -> usize {
        self.context_length
    }

    fn clear_kv_cache(&mut self) {
        SwiftLmBackend::clear_kv_cache(self)
    }

    fn begin_prompt(&mut self, prompt_tokens: &[u32]) -> usize {
        SwiftLmBackend::begin_prompt(self, prompt_tokens)
    }
}

/// End-to-end generation through the shared decode loop, not just a forward
/// pass.
///
/// The parity test above checks logits. This checks that the backend is a
/// genuine `TextDecoder` drop-in: the same loop that drives `MlxBackend`
/// tokenizes, prefills, samples, stops on eos, and produces text. A backend
/// that returns correct logits but cannot be driven by the engine is not a
/// replacement for anything.
#[cfg(all(test, car_mlxlm_swift_built))]
mod swift_lm_decode_tests {
    use super::SwiftLmBackend;
    use crate::backend::local::TextDecoder;
    use std::path::PathBuf;

    #[test]
    fn generates_through_the_shared_text_decoder_surface() {
        let dir = PathBuf::from(std::env::var("HOME").unwrap_or_default())
            .join(".car/models/Qwen3-0.6B-MLX");
        if !dir.join("config.json").exists() {
            eprintln!("SKIPPED generates_through_the_shared_text_decoder_surface: no checkpoint");
            return;
        }
        let mut backend = SwiftLmBackend::load(&dir).expect("load");

        // The trait surface the engine drives, exercised in the same order.
        assert!(
            backend.context_length() >= 4096,
            "context window looks wrong"
        );
        assert!(!backend.eos_ids().is_empty(), "no stop tokens");

        let prompt = "The capital of France is";
        let ids = backend.encode(prompt).expect("encode");
        assert!(!ids.is_empty());

        let offset = backend.begin_prompt(&ids);
        assert!(offset <= ids.len());
        let mut logits = backend.forward(&ids[offset..], offset).expect("prefill");

        let eos = backend.eos_ids();
        let mut out = Vec::new();
        for pos in (ids.len()..).take(6) {
            let tok = logits
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
                .map(|(i, _)| i as u32)
                .unwrap_or(0);
            if eos.contains(&tok) {
                break;
            }
            out.push(tok);
            logits = backend.forward(&[tok], pos).expect("decode step");
        }

        let text = backend.decode(&out).expect("decode");
        // Not asserting an exact continuation — that is the parity test's job,
        // and greedy text is chaotic. Asserting the loop produced language at
        // all, which is what caught the Gemma 4 backend returning "-1--".
        assert!(!out.is_empty(), "generated nothing");
        assert!(
            text.chars().any(|c| c.is_alphabetic()),
            "generated no letters: {text:?} — the decode loop is producing junk"
        );
        eprintln!("swift decode loop produced: {text:?}");
    }
}

/// The dispatcher contract, so `local_backend_for` can return this for every
/// architecture instead of the two CAR hand-wrote loaders for.
impl crate::backend::local::LocalInferenceBackend for SwiftLmBackend {
    fn backend_name(&self) -> &'static str {
        "mlx-swift-lm"
    }

    fn supports_capability(&self, cap: crate::schema::ModelCapability) -> bool {
        use crate::schema::ModelCapability as C;
        // The execution contract of a text LM, independent of what the
        // registry claims. Matches what `MlxBackend` reported, because it is
        // the same question about the same class of checkpoint: this is the
        // text tower, so vision/audio/media are not served here even when the
        // checkpoint came from a multimodal repo.
        match cap {
            C::Generate
            | C::ToolUse
            | C::MultiToolCall
            | C::Reasoning
            | C::Summarize
            | C::Code
            | C::Classify
            | C::Embed
            | C::Rerank => true,
            C::Grounding
            | C::Vision
            | C::VideoUnderstanding
            | C::AudioUnderstanding
            | C::SpeechToText
            | C::TextToSpeech
            | C::ImageGeneration
            | C::VideoGeneration => false,
        }
    }

    fn render_prompt(&self, req: &crate::GenerateRequest) -> Result<String, InferenceError> {
        // The checkpoint's own template first. The trait default is Qwen3's
        // format, which would silently mis-render every other architecture
        // this backend can now load.
        match &self.chat_template {
            Some(t) => t.render_request(req),
            None => Ok(crate::tasks::generate::render_chat_prompt(req)),
        }
    }

    fn parse_tool_calls(&self, text: &str) -> (String, Vec<crate::ToolCall>) {
        // Tool-call syntax is per-architecture and is NOT carried in the chat
        // template, so it dispatches on model_type. Gemma emits
        // `<|tool_call>…<tool_call|>`; the shared parser understands Qwen's Hermes
        // `<tool_call>{json}</tool_call>`, which most others follow.
        if self.model_type.starts_with("gemma") {
            crate::tasks::generate::parse_gemma4_tool_calls(text)
        } else {
            crate::tasks::generate::parse_tool_calls(text)
        }
    }
}