car-inference 0.53.0

//! Local in-process inference backends — the local mirror of the remote
//! [`ProtocolHandler`](crate::protocol::ProtocolHandler) abstraction.
//!
//! Remote models dispatch cleanly through `ProtocolHandler` + `handler_for`.
//! Local (in-process) models historically did not: the engine hand-dispatched
//! a single hardcoded `Qwen3Model` via ad-hoc `cfg` + tag matching. This module
//! introduces the missing seam so each architecture is one isolated backend
//! impl instead of another arm in `generate_tracked_inner`.
//!
//! Two layers, on purpose:
//!
//! - [`TextDecoder`] — the token-level primitives (encode / forward / decode /
//!   eos / context / cache). The engine's shared decode loop (`drive_generation`)
//!   runs sampling, stop detection, TTFT timing, and the Metal panic-catch +
//!   cache-invalidation over `&mut dyn TextDecoder`, so that engine-state-coupled
//!   logic lives in ONE place and every backend reuses it.
//! - [`LocalInferenceBackend`] — the higher-level surface (capability contract,
//!   prompt rendering, tool-call parsing). A `LocalInferenceBackend` IS a
//!   `TextDecoder` (supertrait); the engine upcasts to drive the loop.
//!
//! Adding an architecture = implement both traits for a new backend struct and
//! add one arm to [`local_backend_for`]. No engine-dispatch surgery.

use crate::schema::{ModelCapability, QuantScheme, Quantization};
use crate::tasks::generate::{parse_tool_calls, render_chat_prompt, GenerateRequest, ToolCall};
use crate::InferenceError;

/// Token-level primitives every in-process text backend exposes.
///
/// `forward` returns the next-token logits as `Vec<f32>` (the unifying shape the
/// shared MLX sampler `sample_from_logits` consumes). Backends whose native
/// forward returns something else (Candle returns a `Tensor`) adapt here.
pub trait TextDecoder: Send {
    /// Encode text to token ids (with the tokenizer's special tokens).
    fn encode(&self, text: &str) -> Result<Vec<u32>, InferenceError>;

    /// Decode token ids back to text.
    fn decode(&self, tokens: &[u32]) -> Result<String, InferenceError>;

    /// One prefill/decode pass; returns the final position's logits.
    fn forward(&mut self, tokens: &[u32], pos: usize) -> Result<Vec<f32>, InferenceError>;

    /// Every token id that terminates generation — the model's `eos` plus any
    /// chat-turn-ending control tokens. The shared loop stops on any of these,
    /// so it never needs to know a specific architecture's eos convention.
    fn eos_ids(&self) -> Vec<u32>;

    /// Context window in tokens (for the prompt-truncation guard).
    fn context_length(&self) -> usize;

    /// Reset the KV cache between independent generations.
    fn clear_kv_cache(&mut self);

    /// Prepare the KV cache to prefill `prompt_tokens`, reusing any already-cached
    /// matching prefix (prompt/prefix caching). Returns the offset to begin
    /// prefilling from: `prompt_tokens[..offset]` are already in the cache (their
    /// KV is identical because KV for a fixed token prefix at fixed positions is
    /// deterministic), so the caller only prefills `prompt_tokens[offset..]` at
    /// position `offset`.
    ///
    /// The default clears the cache and returns 0 — a full re-prefill, no reuse.
    /// Backends opt into cross-call reuse by overriding this (and tracking which
    /// tokens their cache represents). This is the big win for multi-turn agent
    /// loops, where each turn re-sends the whole growing conversation.
    fn begin_prompt(&mut self, prompt_tokens: &[u32]) -> usize {
        let _ = prompt_tokens;
        self.clear_kv_cache();
        0
    }
}

/// Result of one engine-driven decode pass.
pub struct LocalGeneration {
    pub text: String,
    pub ttft_ms: Option<u64>,
    /// `"stop"` (hit an eos id or a stop sequence), `"length"` (hit max_tokens),
    /// or `"local_decode_timeout"` (hit the wall-clock decode ceiling, car#851
    /// — spelled out rather than a bare `"timeout"`, which could collide with a
    /// remote provider's raw `finish_reason`). A ceiling-stopped pass still
    /// carries whatever text it managed to generate.
    pub stop_reason: Option<String>,
    /// Prompt tokens fed to the model (post context-window truncation), so the
    /// in-process path reports real `TokenUsage` like the remote providers do —
    /// without it, decode-throughput (tokens/sec) is unmeasurable (it was always
    /// 0 for local models).
    pub prompt_tokens: usize,
    /// Tokens the model generated this pass.
    pub completion_tokens: usize,
}

/// Outcome of the shared decode loop. Distinguishes a normal failure (the
/// backend is still usable) from one where a panic crossed the compute/FFI
/// boundary and left the backend in an indeterminate state — the caller, which
/// owns the cache lock, must then evict it.
pub enum DriveError {
    /// Normal failure (encode/decode/sample). Backend remains usable.
    Recoverable(InferenceError),
    /// A panic was caught mid-forward. The caller MUST drop its guard and
    /// invalidate the backend from the cache before returning.
    BackendCorrupted(InferenceError),
}

impl DriveError {
    pub fn into_inner(self) -> InferenceError {
        match self {
            DriveError::Recoverable(e) | DriveError::BackendCorrupted(e) => e,
        }
    }
}

/// Every `config.json` `model_type` an in-process Rust MLX backend can load.
///
/// This list is short, and it is meant to stay short. The `mlx-rs` crate ships
/// ops and nn layers but no models, so each entry here is 1-2k lines of
/// hand-written Rust that only exists after someone ports the architecture.
/// (Its workspace has an `mlx-lm` crate, but at v0.0.1 it covers fewer
/// architectures than this list does.)
/// Upstream publishes new `model_type` strings faster than that can be
/// sustained, so anything absent here routes to the external `vllm-mlx`
/// runtime instead — see the local-models rule in CLAUDE.md before adding to
/// it.
pub const NATIVE_MLX_MODEL_TYPES: &[&str] = &[
    // backend::mlx — the dedicated `MlxBackend` path (dense + MoE).
    "qwen3",
    "qwen3_moe",
    // backend::mlx_gemma4 — dispatched by `local_backend_for`.
    "gemma4_unified",
    "gemma4_unified_text",
];

/// Whether an in-process Rust backend can load `model_type` **on this build**.
///
/// Platform-gated on purpose: the native backends are MLX, so off Apple Silicon
/// the honest answer is always "no" regardless of the architecture, and a caller
/// choosing between the native and external source must not be told otherwise.
pub fn has_native_backend(model_type: &str) -> bool {
    #[cfg(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx)))]
    {
        let want = model_type.to_ascii_lowercase();
        NATIVE_MLX_MODEL_TYPES.iter().any(|t| *t == want)
    }
    #[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
    {
        let _ = model_type;
        false
    }
}

/// `config.json` `model_type` values the GGUF path implements.
///
/// Short on purpose. GGUF is not how a machine runs the best model it can: the
/// answer is MLX on Apple Silicon and CUDA elsewhere, both reading safetensors.
/// This path exists for the latency-critical hot set — see the local-models
/// rule in CLAUDE.md.
pub const GGUF_MODEL_TYPES: &[&str] = &["qwen3", "qwen3_moe"];

/// Whether a GGUF checkpoint of this architecture has a backend **on this
/// build**.
///
/// False on Apple Silicon for everything: `backend::candle` is compiled only
/// when the MLX path is absent (see `backend/mod.rs`), so a Mac has no GGUF
/// loader at all and a GGUF row there names a backend that does not exist.
/// Elsewhere it is the Candle GGUF loader, on CUDA, for the architectures it
/// implements.
pub fn gguf_backend_serves(model_type: &str) -> bool {
    #[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
    {
        GGUF_MODEL_TYPES.contains(&model_type)
    }
    #[cfg(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx)))]
    {
        let _ = model_type;
        false
    }
}

/// Whether the native MLX loader can decode this weight layout.
///
/// Platform-independent on purpose: it is a statement about the *format*, so it
/// is testable on every runner rather than only where `has_native_backend` can
/// return true. [`native_backend_serves`] combines it with the architecture.
///
/// The vocabulary is MLX's own — `affine`, `mxfp4`, `mxfp8`, `nvfp4` — and the
/// arms below mirror `backend::mlx`'s `build_qlinear`: affine goes through
/// `quantized_matmul`, `mxfp8` is dequantized to dense at load, and everything
/// else is refused there rather than decoded wrong.
pub fn quantization_is_decodable(quantization: Option<&Quantization>) -> bool {
    let Some(quantization) = quantization else {
        // Nothing declared: a full-precision checkpoint on the dense path.
        return true;
    };
    match quantization.scheme {
        // The affine `quantized_matmul` path, and unquantized weights.
        QuantScheme::AffineGroupInt | QuantScheme::Unquantized => true,
        // Only mxfp8 has a decoder; `mxfp4` is refused at load.
        QuantScheme::BlockScaledFloat => quantization.bits == Some(8),
        // GGUF layouts are not safetensors and never reach this loader.
        QuantScheme::KQuantMixed | QuantScheme::RtnBlock => false,
        // A layout that named itself and was not recognized: `nvfp4` today, and
        // the HuggingFace `quant_method` families (awq, gptq, bitsandbytes).
        //
        // Note this is *stricter* than the loader, deliberately. `build_qlinear`
        // takes the affine arm whenever a layer carries both `.scales` and
        // `.biases`, whatever `mode` says — but MLX emits biases only for
        // affine, so no checkpoint it produces reaches that arm with an
        // unrecognized mode. The `quant_method` families are the ones that
        // matter: the native loader has no path for them and does not refuse
        // them either, it builds a dense layer from a packed weight and emits
        // garbage (car-releases#61). Refusing here is what makes that loud.
        QuantScheme::Unknown => false,
    }
}

/// Whether the in-process Rust backend can serve a checkpoint of this
/// `model_type` **in this quantization**, on this build.
///
/// [`has_native_backend`] answers only the architecture half. The native MLX
/// loader also decodes exactly two weight layouts — see
/// [`quantization_is_decodable`].
///
/// The refusal it mirrors happens while reading weights, which is *after* the
/// download. Consulting the quantization at admission turns a wasted
/// multi-gigabyte fetch and a load-time error into a routing decision: the
/// checkpoint goes to the CAR-managed external runtime, which does serve it.
/// Same rule and same reason as
/// [`external_flux::native_backend_serves`](crate::backend::external_flux) on
/// the image path — the model decides the backend.
pub fn native_backend_serves(model_type: &str, quantization: Option<&Quantization>) -> bool {
    has_native_backend(model_type) && quantization_is_decodable(quantization)
}

/// What a backend claims *without* loading weights — consulted by the registry
/// gate (does a backend exist for this `model_type`?) and dispatch.
///
/// Intended as the single source of truth, and not yet one: the local-directory
/// scan in `registry::synthesize_local_schema` still keys on its own
/// `KNOWN_LLM_TYPES` list, which admits 21 architectures to `ModelSource::Mlx`
/// where native loaders exist for four.
pub struct BackendDescriptor {
    pub backend_name: &'static str,
    /// `config.json` `model_type` strings this backend services.
    pub model_types: &'static [&'static str],
}

/// The local in-process mirror of [`ProtocolHandler`](crate::protocol::ProtocolHandler).
/// One impl per architecture family; loaded lazily and cached in the engine.
pub trait LocalInferenceBackend: TextDecoder {
    /// Stable id for tracing / unsupported-mode messages (e.g. `"native-mlx-qwen3"`).
    fn backend_name(&self) -> &'static str;

    /// The *execution* contract: what THIS loaded checkpoint can actually
    /// service, independent of the registry's routing claim. Mirrors today's
    /// `MlxBackend::supports_capability`.
    fn supports_capability(&self, cap: ModelCapability) -> bool;

    /// Render a request to the model's wire prompt string. The default is the
    /// hardcoded Qwen3 chat format; backends with their own template (gemma) or
    /// the data-driven `ChatTemplate` path override this.
    fn render_prompt(&self, req: &GenerateRequest) -> Result<String, InferenceError> {
        Ok(render_chat_prompt(req))
    }

    /// Extract tool calls from generated text. The default understands the Qwen
    /// Hermes `<tool_call>{json}</tool_call>` convention; architectures with a
    /// different convention (gemma's `<|tool_call>…<tool_call|>`) override.
    fn parse_tool_calls(&self, text: &str) -> (String, Vec<ToolCall>) {
        parse_tool_calls(text)
    }
}

// ── MLX (Apple Silicon) ──────────────────────────────────────────────────────

#[cfg(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx)))]
mod mlx_impl {
    use super::*;
    use crate::backend::MlxBackend;

    impl TextDecoder for MlxBackend {
        fn encode(&self, text: &str) -> Result<Vec<u32>, InferenceError> {
            MlxBackend::encode(self, text)
        }

        fn decode(&self, tokens: &[u32]) -> Result<String, InferenceError> {
            MlxBackend::decode(self, tokens)
        }

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

        fn eos_ids(&self) -> Vec<u32> {
            // Qwen3: the config eos plus the chat-turn-ender `<|im_end|>`.
            let mut ids = Vec::new();
            if let Some(e) = self.eos_token_id() {
                ids.push(e);
            }
            if let Some(e) = self.token_id("<|im_end|>") {
                if !ids.contains(&e) {
                    ids.push(e);
                }
            }
            ids
        }

        fn context_length(&self) -> usize {
            MlxBackend::context_length(self)
        }

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

    impl LocalInferenceBackend for MlxBackend {
        fn backend_name(&self) -> &'static str {
            "native-mlx-qwen3"
        }

        fn supports_capability(&self, cap: ModelCapability) -> bool {
            MlxBackend::supports_capability(self, cap)
        }
        // render_prompt / parse_tool_calls: Qwen3 defaults are correct.
    }

    // ── Gemma 4 (gemma4_unified text) ────────────────────────────────────────

    use crate::backend::mlx_gemma4::{parse_gemma4_tool_calls, Gemma4Backend};

    impl TextDecoder for Gemma4Backend {
        fn encode(&self, text: &str) -> Result<Vec<u32>, InferenceError> {
            Gemma4Backend::encode(self, text)
        }

        fn decode(&self, tokens: &[u32]) -> Result<String, InferenceError> {
            Gemma4Backend::decode(self, tokens)
        }

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

        fn eos_ids(&self) -> Vec<u32> {
            Gemma4Backend::eos_token_ids(self)
        }

        fn context_length(&self) -> usize {
            Gemma4Backend::context_length(self)
        }

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

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

    impl LocalInferenceBackend for Gemma4Backend {
        fn backend_name(&self) -> &'static str {
            "native-mlx-gemma4"
        }

        fn supports_capability(&self, cap: ModelCapability) -> bool {
            use ModelCapability as C;
            match cap {
                // Text-only for now (vision/audio towers not loaded).
                C::Generate
                | C::ToolUse
                | C::MultiToolCall
                | C::Reasoning
                | C::Summarize
                | C::Code
                | C::Classify => true,
                C::Rerank
                | C::Embed
                | C::Grounding
                | C::Vision
                | C::VideoUnderstanding
                | C::AudioUnderstanding
                | C::SpeechToText
                | C::TextToSpeech
                | C::ImageGeneration
                | C::VideoGeneration => false,
            }
        }

        fn render_prompt(&self, req: &GenerateRequest) -> Result<String, InferenceError> {
            match self.chat_template() {
                Some(t) => t.render_request(req),
                None => Ok(render_chat_prompt(req)),
            }
        }

        fn parse_tool_calls(&self, text: &str) -> (String, Vec<ToolCall>) {
            parse_gemma4_tool_calls(text)
        }
    }
}

// Non-macOS (Candle) deliberately does NOT implement these traits: the Candle
// backend is already a single generic GGUF loader, so it has no multi-arch
// dispatch problem to solve. This abstraction targets the macOS MLX path, where
// each architecture (Qwen3, Gemma 4, …) is a distinct hand-written backend that
// would otherwise each need its own arm in the engine's dispatch.

// ── Dispatch (macOS MLX) ─────────────────────────────────────────────────────

/// Read a local model's `config.json` `model_type` (lowercased). This is the
/// authoritative architecture signal — the same field the registry gates on.
#[cfg(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx)))]
fn read_model_type(model_dir: &std::path::Path) -> Result<String, InferenceError> {
    let cfg_path = model_dir.join("config.json");
    let raw = std::fs::read_to_string(&cfg_path).map_err(|e| {
        InferenceError::InferenceFailed(format!("read {}: {e}", cfg_path.display()))
    })?;
    let cfg: serde_json::Value = serde_json::from_str(&raw).map_err(|e| {
        InferenceError::InferenceFailed(format!("parse {}: {e}", cfg_path.display()))
    })?;
    Ok(cfg
        .get("model_type")
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_ascii_lowercase())
}

/// The local in-process mirror of [`handler_for`](crate::protocol::handler_for):
/// construct the loaded backend for a model directory, dispatched on its
/// `config.json` `model_type`.
///
/// Qwen3 is intentionally absent — it keeps the dedicated `MlxBackend` path on
/// the engine (which also backs streaming / tokenize / embeddings). This
/// dispatch is for the *additional* architectures that path doesn't serve.
/// Architecture arms are added as their backends land (Gemma 4: B4).
#[cfg(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx)))]
pub fn local_backend_for(
    model_dir: &std::path::Path,
) -> Result<Box<dyn LocalInferenceBackend>, InferenceError> {
    let model_type = read_model_type(model_dir)?;
    match model_type.as_str() {
        "gemma4_unified" | "gemma4_unified_text" => Ok(Box::new(
            crate::backend::mlx_gemma4::Gemma4Backend::load(model_dir)?,
        )),
        other => Err(InferenceError::InferenceFailed(format!(
            "no in-process MLX backend for model_type '{other}' ({}); Qwen3 uses the \
             dedicated native path and other architectures route to vLLM-MLX",
            model_dir.display()
        ))),
    }
}

#[cfg(test)]
mod native_admission_tests {
    use super::*;

    /// The layout half of the gate, mirroring `backend::mlx`'s `build_qlinear`
    /// arms. Platform-independent so it actually executes on CI — the earlier
    /// version of this test was gated to Apple Silicon, where no CI runner
    /// evaluates it, and its ungated companion asserted only `false`s that
    /// `has_native_backend` already produced off-Mac.
    #[test]
    fn decodable_layouts_mirror_the_loader() {
        let decodable = [
            // Affine group quant → `quantized_matmul`.
            Quantization::from_mlx_config(Some(4), Some(64), None).unwrap(),
            Quantization::from_mlx_config(Some(4), Some(64), Some("affine")).unwrap(),
            // mxfp8 → dequantized to dense at load.
            Quantization::from_mlx_config(Some(8), Some(32), Some("mxfp8")).unwrap(),
            // Full precision → the dense path.
            Quantization::parse("bf16"),
        ];
        for q in &decodable {
            assert!(
                quantization_is_decodable(Some(q)),
                "loader decodes {q:?}, gate must admit it"
            );
        }

        let refused = [
            // Microscaling variants with no decoder.
            Quantization::from_mlx_config(Some(4), Some(32), Some("mxfp4")).unwrap(),
            Quantization::from_mlx_config(Some(4), Some(16), Some("nvfp4")).unwrap(),
            // GGUF layouts never reach the safetensors loader.
            Quantization::parse("Q4_K_M"),
            Quantization::parse("Q8_0"),
        ];
        for q in &refused {
            assert!(
                !quantization_is_decodable(Some(q)),
                "loader cannot decode {q:?}, gate must refuse it"
            );
        }

        // No block at all is a full-precision checkpoint, not an unknown one.
        assert!(quantization_is_decodable(None));
    }

    /// mxfp8 is admitted on its width, not on the `mxfp` prefix.
    #[test]
    fn only_the_eight_bit_microscaling_format_is_decodable() {
        for bits in [2u8, 4, 6, 16] {
            let q = Quantization::from_mlx_config(Some(bits), Some(32), Some("mxfp8")).unwrap();
            assert_eq!(q.scheme, QuantScheme::BlockScaledFloat);
            assert!(
                !quantization_is_decodable(Some(&q)),
                "only 8-bit microscaling has a decoder, not {bits}-bit"
            );
        }
    }

    /// The gate is the conjunction: neither half can rescue the other.
    #[test]
    fn architecture_and_layout_must_both_pass() {
        let affine = Quantization::from_mlx_config(Some(4), Some(64), None).unwrap();
        let mxfp4 = Quantization::from_mlx_config(Some(4), Some(32), Some("mxfp4")).unwrap();

        // An unsupported architecture is refused whatever the layout — and on
        // a non-Apple build, every architecture is unsupported.
        assert!(!native_backend_serves("qwen3_5_moe", Some(&affine)));
        assert!(!native_backend_serves(
            "some_architecture_from_next_month",
            None
        ));

        // A supported architecture still needs a decodable layout. Off Apple
        // Silicon `has_native_backend` is false for everything, so assert
        // against it rather than hardcoding a platform answer.
        let supported = has_native_backend("qwen3");
        assert_eq!(native_backend_serves("qwen3", Some(&affine)), supported);
        assert_eq!(native_backend_serves("qwen3", None), supported);
        assert!(
            !native_backend_serves("qwen3", Some(&mxfp4)),
            "an undecodable layout is refused on every platform"
        );
    }
}