diff --git a/Cargo.toml b/Cargo.toml index cd3a49c..28d0105 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ memmap2 = "0.9.11" png = "0.17.16" rfd = "0.15.4" serde = { version = "1.0.228", features = ["derive"] } -serde_json = { version = "1.0.149", features = ["preserve_order"] } +serde_json = { version = "1.0.149", features = ["preserve_order", "raw_value"] } sha2 = "0.11.0" ureq = { version = "3.3.0", default-features = false, features = ["rustls"] } diff --git a/PLAN.md b/PLAN.md index d0acfd1..1196796 100644 --- a/PLAN.md +++ b/PLAN.md @@ -12,12 +12,12 @@ sessions, tools, and turn orchestration. | Area | Status | Available now | Still open | | --- | --- | --- | --- | -| App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, native Application/Edit/Window menus, Preferences panel, Model Manager, and streaming composer | Native app release packaging | +| App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, Chat/Stats segmented workspace, native Application/Edit/Window menus, Preferences panel, Model Manager, and streaming composer | Native app release packaging | | Projects and sessions | Implemented with transcripts and KV resume | Native folder picker, project-name dialog, create/select/delete projects and sessions, durable structured chat history, and per-session KV checkpoints | Rename/archive and model binding | | Persistence | Implemented for local chat | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults, streamed message updates, reopen tests, context/speed state, and local KV rehydration | Downloaded-artifact metadata and session-runtime migrations; HTTP chats deliberately never enter SQLite | -| Model runtime | DeepSeek Flash vertical implemented | Rust-owned mmap/model/session graph, local Metal kernels, full configured sparse context, DS4 sampling, one process-wide UI/HTTP owner, cancellation, idle unloading, prefix continuation, and KV checkpoint save/load | DSpark/SSD/steering and GLM/Pro execution | -| Local endpoint | Chat Completions vertical implemented | Automatic configurable localhost listener (port 4000 by default), verified-on-disk model discovery, transient requests, non-streaming and SSE Chat Completions, reasoning, usage, sampling, stops, cancellation, tools, and exact in-process tool replay | GUI status and enable/disable controls, Responses, legacy Completions, Anthropic Messages, full `ds4_server.c` fixture parity, and cross-restart exact tool replay | -| Agent | Solid local chat implemented | Multi-turn role rendering, structured reasoning disclosure, Markdown, automatic scrolling, context/speed display, durable transcript rehydration, and KV-backed resume | Tools, approvals, compaction, and richer statistics | +| Model runtime | DeepSeek Flash decode and batched prefill implemented | Rust-owned mmap/model/session graph, local Metal kernels, full configured sparse context, exact absolute-position chunk scheduling, cold and resumed layer-major prefill (including unaligned compressor/indexer frontiers), DS4 sampling, one process-wide UI/HTTP owner, cancellation, idle unloading, prefix continuation, KV checkpoint save/load, and lock-free runtime metrics | Remaining KV continuation/rewrite parity, DSpark/SSD/steering, GLM/Pro execution, and sampled per-kernel GPU timing if coarse phase metrics prove insufficient | +| Local endpoint | Exact parity in progress | Configurable localhost listener, verified-on-disk model discovery, transient Chat/Completions/Messages/Responses routes, native JSON/SSE envelopes, exact tool schema/call/result replay, live tool continuation, and differential C/Rust fixtures for tool use, cache accounting, reasoning, cold 6k-token prefill, resumed unaligned prefill, and real Pi tool execution | Remaining `ds4_server.c` parser/tool recovery, KV-store policy, disconnect/cancellation, scheduling/lifecycle, and error/context-limit fixtures; enable/disable controls | +| Agent | Solid local chat implemented | Multi-turn role rendering, structured reasoning disclosure, Markdown, automatic scrolling, context/speed display, durable transcript rehydration, KV-backed resume, and a live Stats dashboard | Tools, approvals, and compaction | | A2UI | Future extension | bDS2 provides the reference structured render-tool contract | Native inline surfaces for local chat after core chat and tools are stable | | Dev brain | Not started | No vault access | Obsidian vault selection, durable access, and agent knowledge/memory tools | | Release | Not started | Development binary builds and tests | `.app` packaging, signing, entitlements, and notarization | @@ -29,8 +29,12 @@ executor, streams generated text, supports cancellation and multi-turn role replay, persists reasoning and answers while they stream, rehydrates selected sessions, and unloads model resources after the configured idle timeout. The current Rust graph supports the full configured context with ratio-4 sparse -attention and durable local KV checkpoints. The shared HTTP Chat Completions -service is available; the remaining protocol routes and agent tools remain open. +attention, fixed-size cold/resumed batch chunks, unaligned compressor/indexer +continuation, and durable local KV checkpoints. The shared HTTP service exposes +every `ds4_server.c` route and passes the initial differential corpus for native +tool use/continuation, cache accounting, reasoning, and Pi execution. Full +reference parity remains in progress until the remaining parser, KV-store, +cancellation, scheduling, and error fixtures are ported. ## Phase 0 — project and session shell @@ -220,6 +224,27 @@ and generation value; show the same effective defaults and validation as the CLI; and construct one engine configuration plus one turn configuration without chat or HTTP code re-parsing preference fields. +### Runtime observability (implemented) + +- Keep instrumentation in the shared model/runtime and HTTP paths so local chat + and endpoint work report through the same counters. +- Publish hot-path state with relaxed atomics only. Generation and prefill never + wait for the UI, allocate telemetry events, write logs, or scan the KV cache. +- Sample snapshots in the UI every 200 ms and retain only a bounded in-memory + history for prefill, decode, request-rate, and KV read/write-rate graphs. +- Show model phase and lifecycle, prefill/decode throughput, context occupancy, + request queue and source, exact memory/disk hits, misses, invalid entries, + prefix reuse, checkpoint storage and I/O, endpoint traffic and latency, + errors, and model mapping details on the Stats tab. +- Derive KV file sizes when checkpoints are already being saved and scan the + cache directory only once at startup. Add sampled Metal command-buffer or + per-kernel timings later only if coarse phase counters cannot diagnose a + measured engine bottleneck. + +Exit criterion: model generation and endpoint traffic visibly update the Stats +dashboard without introducing locks, telemetry I/O, or unbounded history on the +prefill/decode path. + ### Targeted downloads and storage - Preferences only selects the model/runtime configuration. A separate Model @@ -266,37 +291,63 @@ after the configured idle timeout, and pass the matching `../ds4` token-output fixture. Repeat through the local endpoint and verify both paths share the same engine lifecycle. -## Phase 2 — OpenAI-compatible local endpoint +## Phase 2 — exact `ds4_server.c` HTTP parity -Status: **partially implemented**. The shared single-model runtime, -verified-model discovery, configurable localhost port, and the first -`ds4_server.c`-compatible Chat Completions vertical are implemented. The -remaining protocol routes plus GUI status and enable/disable controls remain -open. +Status: **exact parity in progress**. The shared single-model runtime, +verified-model discovery, configurable localhost port, all four generation +routes, native JSON/SSE response envelopes, exact batched prefill, tool replay, +and live continuation exist. Differential C/Rust fixtures cover tools, +reasoning, usage/cache accounting, and real Pi execution. Full parser/tool +recovery, KV-store, cancellation, scheduling, lifecycle, and error fixture +parity remain open, so no route is considered complete yet. + +`../ds4/ds4_server.c` is normative for every externally observable HTTP +behavior. “Compatible” approximations are not acceptable: routes, accepted +fields, defaults, aliases, prompt rendering, status codes, headers, JSON/SSE +shapes and ordering, five-second prefill keepalives, reasoning/tool translation, +usage, finish reasons, errors, cancellation, and cache continuation must match +the C server exactly. Port its fixtures alongside each Rust slice and keep every +previous parity fixture passing. + +This requirement covers the complete communication behavior, not only the HTTP +wire format. The Rust request path must follow the same request parsing and +validation, model aliases, prompt/token construction, reasoning mode, sampling, +stop matching, tool schema/call/result translation, live tool continuation, +canonical replay, KV lookup/rewrite/checkpoint policy, scheduling, disconnect +cancellation, usage accounting, finish decisions, and model lifecycle invoked +by `ds4_server.c`. `ds4.c` is normative for engine operations reached from the +server. A route is complete only when its differential request corpus passes +against the C server; partial route support stays labeled incomplete. 1. **Partially implemented:** add a localhost-only HTTP service whose lifecycle is independent of the engine. It can listen while the model is unloaded and acquires the configured model only for inference. It listens on configurable port 4000 by default; start/stop controls remain open. Never expose a LAN listener without an explicit setting. -2. Implement the DwarfStar compatibility surface from `ds4_server.c` in this +2. **Partially implemented:** expose the complete DwarfStar HTTP surface from + `ds4_server.c` in this order, preserving its request parsing, prompt formatting, reasoning fields, streaming events, errors, usage accounting, stop behavior, tool replay, and cancellation semantics rather than inventing an app-specific protocol: - `GET /v1/models` + - `GET /v1/models/{known-alias}` + - `POST /v1/messages` - `POST /v1/chat/completions` - `POST /v1/responses` - `POST /v1/completions` + - `OPTIONS` with the same CORS behavior when enabled + `GET /v1/models` scans managed artifacts without forcing a model to load and reports only main model files that are fully downloaded and have a matching on-disk verification marker. Missing, partial, unverified, or checksum-invalid artifacts are never advertised. -3. Support streaming SSE, reasoning output, usage accounting, cancellation, - sampling parameters, tool schemas, and tool choice. -4. Port exact sampled tool-call replay and deterministic canonicalization from +3. **Partially implemented:** streaming SSE, reasoning output, usage accounting, + sampling parameters, tool schemas, and automatic/disabled tool choice work; + exact disconnect cancellation and remaining forced-choice errors are open. +4. **Partially implemented:** port exact sampled tool-call replay and deterministic canonicalization from `ds4_server.c` so a client's normalized JSON does not destroy KV-prefix reuse. -5. Route API work through the single process-wide model owner and engine-owned +5. **Partially implemented:** route API work through the single process-wide model owner and engine-owned KV store defined in Phase 1. External requests are always stateless at the application layer: never write their messages to SQLite and never bind a Responses `conversation`, request identifier, or client-supplied session key @@ -304,15 +355,19 @@ open. DS4-compatible prefix lookup may accelerate that replay using opaque KV cache files. Add resident batching only if one-model serialized execution is later proven insufficient. -6. Display endpoint address, model, active requests, token rates, and errors in - the GUI. +6. **Implemented:** display endpoint address, model, active requests, token + rates, cache behavior, and errors in the GUI Stats dashboard. -Exit criterion: the upstream DwarfStar server tests pass against the app for -non-streaming and streaming Chat Completions and Responses calls, including a -multi-turn tool call. Verified installed models are listed without loading one; -missing/unverified models are absent; local and HTTP generations cannot load two -models concurrently; external calls leave no project, session, message, or -transcript rows behind; and compatible KV cache files remain reusable. +Exit criterion: the complete `ds4_server.c` fixture suite and a differential +black-box corpus pass against the Rust app for Models, Anthropic Messages, Chat +Completions, Responses, Completions, and OPTIONS, in streaming and non-streaming +forms where supported. The corpus covers valid requests, aliases/defaults, +context limits, tools and multi-turn continuation, reasoning, cache reuse, +cancellation/disconnects, queueing, errors, usage, and finish reasons. Verified +installed models are listed without loading one; missing/unverified models are +absent; local and HTTP generations cannot load two models concurrently; +external calls leave no project, session, message, or transcript rows behind; +and compatible KV cache files remain reusable. ## Phase 3 — durable agent sessions @@ -524,10 +579,11 @@ tests, and release packaging remain open. 1. Finish the Flash session core as one larger slice: port the ratio-4 indexer and indexed attention for full configured context, then add prefix sync and compatible KV checkpoint save/load. **Implemented.** -2. Put the OpenAI-compatible endpoint on the same engine lifecycle so local - chats and HTTP requests cannot load duplicate models. Keep HTTP conversation - state transient except for opaque KV cache files. **Chat Completions vertical - implemented; remaining reference routes are open.** +2. Complete exact `ds4_server.c` communication parity as one tracked program, + beginning with exact layer-major/chunked/resumed prefill and KV continuation, + then porting every parser/route/state machine and its fixtures. Keep HTTP + conversation state transient except for opaque KV cache files. **In + progress; no route is considered complete until differential tests pass.** 3. Local transcript/KV persistence and the conversation UI are implemented; port the agent tools after the shared endpoint lifecycle is stable. 4. Add the opt-in Dev Brain tool after local file/search boundaries and agent diff --git a/scripts/endpoint_continuation.py b/scripts/endpoint_continuation.py new file mode 100644 index 0000000..ac569a6 --- /dev/null +++ b/scripts/endpoint_continuation.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +import json +import sys +import urllib.request + + +base_url = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:9000" +prompt = "Call echo with text hi. After its result, reply with only DONE." +parameters = { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], +} +chat_tool = { + "type": "function", + "function": { + "name": "echo", + "description": "Echo text", + "parameters": parameters, + }, +} +anthropic_tool = { + "name": "echo", + "description": "Echo text", + "input_schema": parameters, +} +responses_tool = { + "type": "function", + "name": "echo", + "description": "Echo text", + "parameters": parameters, +} + + +def post(path, payload): + request = urllib.request.Request( + base_url + path, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + return json.load(response) + + +first = post( + "/v1/chat/completions", + { + "model": "deepseek-v4-flash", + "messages": [{"role": "user", "content": prompt}], + "tools": [chat_tool], + "reasoning_effort": "none", + "temperature": 0, + "max_tokens": 128, + }, +) +call = first["choices"][0]["message"]["tool_calls"][0] +print("chat-first", json.dumps(first["usage"], separators=(",", ":"))) +second = post( + "/v1/chat/completions", + { + "model": "deepseek-v4-flash", + "messages": [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": "", "tool_calls": [call]}, + { + "role": "tool", + "tool_call_id": call["id"], + "content": "hi", + }, + ], + "tools": [chat_tool], + "reasoning_effort": "none", + "temperature": 0, + "max_tokens": 128, + }, +) +print("chat", json.dumps(second, separators=(",", ":"))) + +first = post( + "/v1/messages", + { + "model": "deepseek-v4-flash", + "messages": [{"role": "user", "content": prompt}], + "tools": [anthropic_tool], + "thinking": {"type": "disabled"}, + "temperature": 0, + "max_tokens": 128, + }, +) +call = next(block for block in first["content"] if block["type"] == "tool_use") +print("anthropic-first", json.dumps(first["usage"], separators=(",", ":"))) +second = post( + "/v1/messages", + { + "model": "deepseek-v4-flash", + "messages": [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": [call]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": call["id"], + "content": "hi", + } + ], + }, + ], + "tools": [anthropic_tool], + "thinking": {"type": "disabled"}, + "temperature": 0, + "max_tokens": 128, + }, +) +print("anthropic", json.dumps(second, separators=(",", ":"))) + +first = post( + "/v1/responses", + { + "model": "deepseek-v4-flash", + "input": prompt, + "tools": [responses_tool], + "reasoning": {"effort": "none"}, + "temperature": 0, + "max_output_tokens": 128, + }, +) +call = next(item for item in first["output"] if item["type"] == "function_call") +print("responses-first", json.dumps(first["usage"], separators=(",", ":"))) +second = post( + "/v1/responses", + { + "model": "deepseek-v4-flash", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": prompt}], + }, + call, + { + "type": "function_call_output", + "call_id": call["call_id"], + "output": "hi", + }, + ], + "tools": [responses_tool], + "reasoning": {"effort": "none"}, + "temperature": 0, + "max_output_tokens": 128, + }, +) +print("responses", json.dumps(second, separators=(",", ":"))) diff --git a/scripts/endpoint_parity.py b/scripts/endpoint_parity.py new file mode 100644 index 0000000..3006ff6 --- /dev/null +++ b/scripts/endpoint_parity.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +import json +import sys +import urllib.request + + +base_url = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:9000" +streaming = "--stream" in sys.argv +parameters = { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], +} +cases = [ + ( + "chat", + "/v1/chat/completions", + { + "model": "deepseek-v4-flash", + "messages": [ + { + "role": "user", + "content": "Call echo with text hi. Do not answer normally.", + } + ], + "tools": [ + { + "type": "function", + "function": { + "name": "echo", + "description": "Echo text", + "parameters": parameters, + }, + } + ], + "reasoning_effort": "none", + "temperature": 0, + "max_tokens": 128, + }, + ), + ( + "anthropic", + "/v1/messages", + { + "model": "deepseek-v4-flash", + "messages": [ + { + "role": "user", + "content": "Call echo with text hi. Do not answer normally.", + } + ], + "tools": [ + { + "name": "echo", + "description": "Echo text", + "input_schema": parameters, + } + ], + "thinking": {"type": "disabled"}, + "temperature": 0, + "max_tokens": 128, + }, + ), + ( + "responses", + "/v1/responses", + { + "model": "deepseek-v4-flash", + "input": "Call echo with text hi. Do not answer normally.", + "tools": [ + { + "type": "function", + "name": "echo", + "description": "Echo text", + "parameters": parameters, + } + ], + "reasoning": {"effort": "none"}, + "temperature": 0, + "max_output_tokens": 128, + }, + ), +] + +for name, path, payload in cases: + if streaming: + payload["stream"] = True + if name == "chat": + payload["stream_options"] = {"include_usage": True} + request = urllib.request.Request( + base_url + path, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + if streaming: + for line in response: + if line.strip(): + print(name, line.decode().rstrip()) + else: + print(name, json.dumps(json.load(response), separators=(",", ":"))) diff --git a/scripts/endpoint_reasoning.py b/scripts/endpoint_reasoning.py new file mode 100644 index 0000000..8d52115 --- /dev/null +++ b/scripts/endpoint_reasoning.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +import json +import sys +import urllib.request + + +base_url = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:9000" +prompt = "Think briefly, then reply with only the number: 17 * 19." +cases = [ + ( + "chat", + "/v1/chat/completions", + { + "model": "deepseek-v4-flash", + "messages": [{"role": "user", "content": prompt}], + "reasoning_effort": "low", + "temperature": 0, + "max_tokens": 96, + }, + ), + ( + "anthropic", + "/v1/messages", + { + "model": "deepseek-v4-flash", + "messages": [{"role": "user", "content": prompt}], + "thinking": {"type": "enabled", "budget_tokens": 64}, + "output_config": {"effort": "low"}, + "temperature": 0, + "max_tokens": 96, + }, + ), + ( + "responses", + "/v1/responses", + { + "model": "deepseek-v4-flash", + "input": prompt, + "reasoning": {"effort": "low", "summary": "auto"}, + "temperature": 0, + "max_output_tokens": 96, + }, + ), +] + +for name, path, payload in cases: + request = urllib.request.Request( + base_url + path, + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + print(name, json.dumps(json.load(response), separators=(",", ":"))) diff --git a/src/app.rs b/src/app.rs index 349907e..1d92a3e 100644 --- a/src/app.rs +++ b/src/app.rs @@ -5,6 +5,7 @@ pub(crate) use view::app_theme; use crate::database::{AppPreferences, Database, ProjectWithSessions, StoredMessage}; #[cfg(target_os = "macos")] use crate::engine::ChatTurn; +use crate::metrics::{Metrics, MetricsSnapshot}; use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; #[cfg(target_os = "macos")] use crate::runtime::{ActiveGeneration, CheckpointTarget, GenerationEvent, GenerationService}; @@ -16,6 +17,7 @@ use crate::settings::{ use iced::widget::{markdown, scrollable}; use iced::{Size, Subscription, Task, keyboard, window}; use rfd::AsyncFileDialog; +use std::collections::VecDeque; use std::fs; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -25,6 +27,7 @@ use std::thread; use std::time::{Duration, Instant}; const APP_ID: &str = "DS4Server.rfc1437.de"; +const METRICS_SAMPLE_INTERVAL: Duration = Duration::from_millis(200); #[derive(Clone)] struct PreferenceDraft { @@ -414,6 +417,13 @@ pub(crate) struct App { pub(super) context_used: u32, pub(super) context_limit: u32, pub(super) tokens_per_second: Option, + pub(super) detail_tab: DetailTab, + metrics: Arc, + pub(super) metrics_snapshot: MetricsSnapshot, + pub(super) metrics_history: VecDeque, + last_http_requests: u64, + last_kv_read_bytes: u64, + last_kv_write_bytes: u64, #[cfg(target_os = "macos")] generation_service: Option, #[cfg(target_os = "macos")] @@ -425,6 +435,22 @@ pub(crate) struct App { error: Option, } +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(super) enum DetailTab { + #[default] + Chat, + Stats, +} + +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct MetricsPoint { + pub(super) decode_tokens_per_second: f32, + pub(super) prefill_tokens_per_second: f32, + pub(super) http_requests_per_second: f32, + pub(super) kv_read_bytes_per_second: f32, + pub(super) kv_write_bytes_per_second: f32, +} + #[derive(Clone, Debug)] pub(super) struct ChatMessage { id: i32, @@ -571,6 +597,9 @@ pub(crate) enum Message { CreateSession, SelectSession(i32, i32), DeleteSession(i32), + ShowChat, + ShowStats, + MetricsTick, } impl App { @@ -584,9 +613,12 @@ impl App { Err(error) => return Self::failed(error, main_window), }; let context_limit = preferences.context_tokens.max(0) as u32; + let metrics = + Arc::new(Metrics::new(&application_support_path().join("kv-cache"))); + let metrics_snapshot = metrics.snapshot(); #[cfg(target_os = "macos")] let (runtime_preferences, generation_service, endpoint, service_error) = - spawn_services(&preferences); + spawn_services(&preferences, Arc::clone(&metrics)); Self { main_window, model_manager_window: None, @@ -613,6 +645,13 @@ impl App { context_used: 0, context_limit, tokens_per_second: None, + detail_tab: DetailTab::Chat, + metrics, + metrics_snapshot, + metrics_history: VecDeque::with_capacity(120), + last_http_requests: 0, + last_kv_read_bytes: 0, + last_kv_write_bytes: 0, #[cfg(target_os = "macos")] generation_service, #[cfg(target_os = "macos")] @@ -644,9 +683,11 @@ impl App { let preference_draft = PreferenceDraft::from_saved(&preferences) .expect("default preferences must use a supported model"); let context_limit = preferences.context_tokens.max(0) as u32; + let metrics = Arc::new(Metrics::new(&application_support_path().join("kv-cache"))); + let metrics_snapshot = metrics.snapshot(); #[cfg(target_os = "macos")] let (runtime_preferences, generation_service, endpoint, service_error) = - spawn_services(&preferences); + spawn_services(&preferences, Arc::clone(&metrics)); #[cfg(target_os = "macos")] let startup_error = service_error; #[cfg(not(target_os = "macos"))] @@ -677,6 +718,13 @@ impl App { context_used: 0, context_limit, tokens_per_second: None, + detail_tab: DetailTab::Chat, + metrics, + metrics_snapshot, + metrics_history: VecDeque::with_capacity(120), + last_http_requests: 0, + last_kv_read_bytes: 0, + last_kv_write_bytes: 0, #[cfg(target_os = "macos")] generation_service, #[cfg(target_os = "macos")] @@ -740,6 +788,9 @@ impl App { self.error = None; } } + Message::ShowChat => self.detail_tab = DetailTab::Chat, + Message::ShowStats => self.detail_tab = DetailTab::Stats, + Message::MetricsTick => self.sample_metrics(), Message::PreferenceModelChanged(model) => { self.preference_draft.model = model; if !model.supports_dspark() { @@ -1153,6 +1204,8 @@ impl App { window::close_requests().map(Message::WindowClosed), window::close_events().map(Message::WindowClosed), ]; + subscriptions + .push(iced::time::every(METRICS_SAMPLE_INTERVAL).map(|_| Message::MetricsTick)); #[cfg(target_os = "macos")] subscriptions.push(iced::time::every(Duration::from_millis(50)).map(|_| { match crate::native_menu::next_event() { @@ -1327,6 +1380,7 @@ impl App { models_path(), application_support_path().join("kv-cache").join("http"), endpoint_port, + Arc::clone(&self.metrics), ) { Ok(endpoint) => Some(endpoint), Err(error) => { @@ -1569,6 +1623,7 @@ impl App { .iter() .map(|message| ChatTurn { user: message.user, + skip_previous_eos: false, reasoning: message.reasoning.clone(), reasoning_complete: message.reasoning_complete, content: message.content.clone(), @@ -1577,6 +1632,7 @@ impl App { #[cfg(target_os = "macos")] messages.push(ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: prompt.clone(), @@ -1729,11 +1785,53 @@ impl App { #[cfg(not(target_os = "macos"))] false } + + fn sample_metrics(&mut self) { + let snapshot = self.metrics.snapshot(); + let http_requests_per_second = snapshot + .http_requests + .saturating_sub(self.last_http_requests) as f32 + / METRICS_SAMPLE_INTERVAL.as_secs_f32(); + let kv_read_bytes_per_second = snapshot + .kv_read_bytes + .saturating_sub(self.last_kv_read_bytes) as f32 + / METRICS_SAMPLE_INTERVAL.as_secs_f32(); + let kv_write_bytes_per_second = snapshot + .kv_write_bytes + .saturating_sub(self.last_kv_write_bytes) + as f32 + / METRICS_SAMPLE_INTERVAL.as_secs_f32(); + self.last_http_requests = snapshot.http_requests; + self.last_kv_read_bytes = snapshot.kv_read_bytes; + self.last_kv_write_bytes = snapshot.kv_write_bytes; + self.metrics_history.push_back(MetricsPoint { + decode_tokens_per_second: if snapshot.phase == crate::metrics::RuntimePhase::Generating + { + snapshot.decode_tokens_per_second + } else { + 0.0 + }, + prefill_tokens_per_second: if snapshot.phase == crate::metrics::RuntimePhase::Prefilling + { + snapshot.prefill_tokens_per_second + } else { + 0.0 + }, + http_requests_per_second, + kv_read_bytes_per_second, + kv_write_bytes_per_second, + }); + if self.metrics_history.len() > 120 { + self.metrics_history.pop_front(); + } + self.metrics_snapshot = snapshot; + } } #[cfg(target_os = "macos")] fn spawn_services( preferences: &AppPreferences, + metrics: Arc, ) -> ( Arc>, Option, @@ -1741,7 +1839,7 @@ fn spawn_services( Option, ) { let runtime_preferences = Arc::new(RwLock::new(preferences.clone())); - let generation = match GenerationService::spawn() { + let generation = match GenerationService::spawn(Arc::clone(&metrics)) { Ok(generation) => generation, Err(error) => return (runtime_preferences, None, None, Some(error)), }; @@ -1751,6 +1849,7 @@ fn spawn_services( models_path(), application_support_path().join("kv-cache").join("http"), u16::try_from(preferences.endpoint_port).unwrap_or(4000), + metrics, ); match endpoint { Ok(endpoint) => (runtime_preferences, Some(generation), Some(endpoint), None), diff --git a/src/app/view.rs b/src/app/view.rs index a3bea63..5566521 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -1,5 +1,6 @@ use super::{ - ActiveDownload, App, Message, ModelDownload, ModelOperation, chat_scroll_id, models_path, + ActiveDownload, App, DetailTab, Message, MetricsPoint, ModelDownload, ModelOperation, + chat_scroll_id, models_path, }; use crate::database::{ProjectWithSessions, Session}; use crate::model::{ @@ -12,6 +13,7 @@ use iced::widget::{ pick_list, progress_bar, row, scrollable, stack, svg, text, text_input, }; use iced::{Alignment, Background, Border, Color, Element, Length, Theme, window}; +use std::collections::VecDeque; use std::path::Path; const ICON_FOLDER: &[u8] = include_bytes!("../../assets/icons/folder.svg"); @@ -191,6 +193,43 @@ impl App { } fn detail(&self) -> Element<'_, Message> { + let chat_active = self.detail_tab == DetailTab::Chat; + let stats_active = self.detail_tab == DetailTab::Stats; + let tabs = container( + row![ + button(text("Chat").size(13)) + .width(92) + .padding([7, 18]) + .on_press(Message::ShowChat) + .style(move |theme, status| segmented_button_style(theme, status, chat_active)), + button(text("Stats").size(13)) + .width(92) + .padding([7, 18]) + .on_press(Message::ShowStats) + .style(move |theme, status| segmented_button_style( + theme, + status, + stats_active + )), + ] + .spacing(2), + ) + .padding(3) + .style(segmented_control_style); + let body = match self.detail_tab { + DetailTab::Chat => self.chat_detail(), + DetailTab::Stats => self.stats_dashboard(), + }; + column![ + container(tabs).center_x(Length::Fill).padding([12, 0]), + body + ] + .width(Length::Fill) + .height(Length::Fill) + .into() + } + + fn chat_detail(&self) -> Element<'_, Message> { let Some(item) = self.selected_project() else { let open_project_content = row![icon(ICON_FOLDER_PLUS, 17), text("Open project…"),] .spacing(8) @@ -396,6 +435,373 @@ impl App { .into() } + fn stats_dashboard(&self) -> Element<'_, Message> { + let stats = &self.metrics_snapshot; + let context_fraction = if stats.context_limit == 0 { + 0.0 + } else { + stats.context_used.min(stats.context_limit) as f32 / stats.context_limit as f32 + }; + let cache_fraction = if stats.last_prompt_tokens == 0 { + 0.0 + } else { + stats.last_cached_tokens as f32 / stats.last_prompt_tokens as f32 + }; + let cache_hit_fraction = if stats.kv_lookups == 0 { + 0.0 + } else { + stats.kv_hits as f32 / stats.kv_lookups as f32 + }; + let endpoint = if stats.server_listening { + format!("Listening · 127.0.0.1:{}", stats.server_port) + } else { + "Stopped".to_owned() + }; + let phase_color = match stats.phase { + crate::metrics::RuntimePhase::Generating => Color::from_rgb8(84, 170, 255), + crate::metrics::RuntimePhase::Prefilling | crate::metrics::RuntimePhase::Loading => { + Color::from_rgb8(240, 180, 70) + } + crate::metrics::RuntimePhase::Ready => Color::from_rgb8(72, 176, 112), + crate::metrics::RuntimePhase::Failed => Color::from_rgb8(220, 80, 86), + crate::metrics::RuntimePhase::Unloaded => muted_text(), + }; + let heading = container( + row![ + column![ + text("Runtime observability").size(24), + text(format!( + "{} · {} · uptime {}", + stats.model, + stats.source.label(), + format_duration(stats.uptime_seconds as f64) + )) + .size(12) + .color(muted_text()), + ] + .spacing(5), + Space::with_width(Length::Fill), + container(text(stats.phase.label()).size(12).color(phase_color)) + .padding([7, 11]) + .style(move |_| status_badge_style(phase_color)), + ] + .align_y(Alignment::Center), + ) + .padding(16) + .style(overview_style); + + let headline = column![ + row![ + metric_card( + "DECODE", + format!("{:.1} tok/s", stats.decode_tokens_per_second), + format!( + "{} completion tokens total", + format_count(stats.completion_tokens) + ), + ), + metric_card( + "PREFILL", + format!("{:.1} tok/s", stats.prefill_tokens_per_second), + format!("{} prompt tokens total", format_count(stats.prompt_tokens)), + ), + ] + .spacing(10), + row![ + metric_card( + "CONTEXT", + format!( + "{} / {}", + format_count(u64::from(stats.context_used)), + format_count(u64::from(stats.context_limit)) + ), + format!("{:.0}% occupied", context_fraction * 100.0), + ), + metric_card( + "WORK", + format!( + "{} active · {} queued", + stats.http_active, stats.queue_depth + ), + format!("{} runtime requests", format_count(stats.runtime_requests)), + ), + ] + .spacing(10), + ] + .spacing(10); + + let throughput = stats_panel( + "MODEL ACTIVITY · LAST 24 SECONDS", + column![ + mini_chart( + &self.metrics_history, + |point| point.decode_tokens_per_second, + Color::from_rgb8(84, 170, 255), + ), + row![ + text("Decode") + .size(12) + .color(Color::from_rgb8(84, 170, 255)), + Space::with_width(Length::Fill), + text(format!("{:.1} tok/s", stats.decode_tokens_per_second)) + .size(12) + .color(muted_text()), + ], + mini_chart( + &self.metrics_history, + |point| point.prefill_tokens_per_second, + Color::from_rgb8(157, 119, 255), + ), + row![ + text("Prefill") + .size(12) + .color(Color::from_rgb8(157, 119, 255)), + Space::with_width(Length::Fill), + text(format!("{:.1} tok/s", stats.prefill_tokens_per_second)) + .size(12) + .color(muted_text()), + ], + ] + .spacing(7) + .into(), + ); + let requests = stats_panel( + "SERVER REQUEST RATE · LAST 24 SECONDS", + column![ + mini_chart( + &self.metrics_history, + |point| point.http_requests_per_second, + Color::from_rgb8(72, 176, 112), + ), + row![ + text(format!("{} requests", format_count(stats.http_requests))).size(12), + Space::with_width(Length::Fill), + text(format!( + "{} errors · {} streaming", + stats.http_errors, stats.http_streaming_requests + )) + .size(12) + .color(muted_text()), + ], + ] + .spacing(7) + .into(), + ); + let latest = self.metrics_history.back().copied().unwrap_or_default(); + let kv_io = stats_panel( + "KV CHECKPOINT I/O · LAST 24 SECONDS", + column![ + mini_chart( + &self.metrics_history, + |point| point.kv_read_bytes_per_second, + Color::from_rgb8(67, 194, 203), + ), + row![ + text(if stats.kv_read_active { + "Read · active" + } else { + "Read" + }) + .size(12) + .color(Color::from_rgb8(67, 194, 203)), + Space::with_width(Length::Fill), + text(format_rate(latest.kv_read_bytes_per_second)) + .size(12) + .color(muted_text()), + ], + mini_chart( + &self.metrics_history, + |point| point.kv_write_bytes_per_second, + Color::from_rgb8(240, 180, 70), + ), + row![ + text(if stats.kv_write_active { + "Write · active" + } else { + "Write" + }) + .size(12) + .color(Color::from_rgb8(240, 180, 70)), + Space::with_width(Length::Fill), + text(format_rate(latest.kv_write_bytes_per_second)) + .size(12) + .color(muted_text()), + ], + ] + .spacing(7) + .into(), + ); + + let model = stats_panel( + "MODEL CORE", + column![ + metric_row("State", stats.phase.label()), + metric_row("Loaded model", stats.model), + metric_row("Mapped weights", format_bytes(stats.model_bytes)), + metric_row("Tensors", format_count(stats.tensor_count)), + metric_row("Vocabulary", format_count(stats.vocabulary_size)), + metric_row("Last load", format_milliseconds(stats.model_load_ms)), + metric_row( + "Lifecycle", + format!( + "{} loads · {} unloads", + stats.model_loads, stats.model_unloads + ), + ), + ] + .spacing(9) + .into(), + ); + let runtime = stats_panel( + "GENERATION", + column![ + metric_row("Last runtime", format_milliseconds(stats.last_runtime_ms)), + metric_row( + "Average runtime", + format_milliseconds(stats.average_runtime_ms) + ), + metric_row("Last prompt", format_count(stats.last_prompt_tokens)), + metric_row("Last reused", format_count(stats.last_cached_tokens)), + metric_row( + "Last completion", + format_count(stats.last_completion_tokens) + ), + metric_row("Cached tokens total", format_count(stats.cached_tokens)), + metric_row( + "Cache reuse", + format!("{:.0}% of last prompt", cache_fraction * 100.0), + ), + metric_row( + "Results", + format!( + "{} completed · {} failed", + stats.completed_requests, stats.failed_requests + ), + ), + ] + .spacing(9) + .into(), + ); + let cache = stats_panel( + "KV CACHE", + column![ + metric_row( + "Total", + format!( + "{} · {} files", + format_bytes(stats.kv_bytes), + stats.kv_files + ) + ), + metric_row( + "Local sessions", + format!( + "{} · {} files", + format_bytes(stats.local_kv_bytes), + stats.local_kv_files + ), + ), + metric_row( + "HTTP transient", + format!( + "{} · {} files", + format_bytes(stats.http_kv_bytes), + stats.http_kv_files + ), + ), + metric_row("Checkpoint writes", format_count(stats.checkpoint_writes)), + metric_row( + "Exact hits", + format!( + "{} · {} memory / {} disk", + stats.kv_hits, stats.kv_memory_hits, stats.kv_disk_hits + ), + ), + metric_row( + "Misses", + format!("{} · {} invalid", stats.kv_misses, stats.kv_invalid), + ), + metric_row("Lookups", format_count(stats.kv_lookups)), + metric_row( + "Exact hit rate", + format!("{:.1}%", cache_hit_fraction * 100.0) + ), + metric_row("Prefix hits", format_count(stats.kv_prefix_hits)), + metric_row( + "Reads", + format!( + "{} · {} · {} errors · last {}", + stats.kv_read_operations, + format_bytes(stats.kv_read_bytes), + stats.kv_read_errors, + format_milliseconds(stats.last_kv_read_ms), + ), + ), + metric_row( + "Writes", + format!( + "{} · {} · {} errors · last {}", + stats.kv_write_operations, + format_bytes(stats.kv_write_bytes), + stats.kv_write_errors, + format_milliseconds(stats.last_kv_write_ms), + ), + ), + progress_bar(0.0..=1.0, cache_fraction.min(1.0)).height(4), + ] + .spacing(9) + .into(), + ); + let server = stats_panel( + "LOCAL SERVER", + column![ + metric_row("Endpoint", endpoint), + metric_row("Active", format_count(u64::from(stats.http_active))), + metric_row("Completed", format_count(stats.http_completed)), + metric_row("Chat completions", format_count(stats.http_chat_requests)), + metric_row("Model queries", format_count(stats.http_model_requests)), + metric_row( + "Runtime sources", + format!( + "{} local · {} HTTP", + stats.local_requests, stats.endpoint_generations + ), + ), + metric_row("Received", format_bytes(stats.http_bytes_received)), + metric_row("Last latency", format_milliseconds(stats.last_http_ms)), + metric_row( + "Average latency", + format_milliseconds(stats.average_http_ms) + ), + ] + .spacing(9) + .into(), + ); + + scrollable( + container( + column![ + heading, + headline, + throughput, + kv_io, + requests, + row![model, runtime].spacing(10), + row![cache, server].spacing(10), + text("Counters are published by the runtime with relaxed atomics and sampled by the UI every 200 ms.") + .size(11) + .color(muted_text()), + ] + .spacing(12), + ) + .padding(24) + .max_width(960) + .center_x(Length::Fill), + ) + .height(Length::Fill) + .into() + } + fn preferences_panel(&self) -> Element<'_, Message> { let dspark_toggle: Option Message> = self .preference_draft @@ -1218,6 +1624,101 @@ fn action_button<'a>(content: impl Into>) -> Button<'a, Mes button(content).padding([8, 14]).style(action_button_style) } +fn metric_card(label: &'static str, value: String, detail: String) -> Element<'static, Message> { + container( + column![ + text(label).size(10).color(muted_text()), + text(value).size(20), + text(detail).size(11).color(muted_text()), + ] + .spacing(5), + ) + .padding(14) + .width(Length::FillPortion(1)) + .style(preference_group_style) + .into() +} + +fn metric_row(label: &'static str, value: impl ToString) -> Element<'static, Message> { + row![ + text(label).size(12).color(muted_text()), + Space::with_width(Length::Fill), + text(value.to_string()).size(12), + ] + .spacing(12) + .align_y(Alignment::Center) + .into() +} + +fn stats_panel<'a>(title: &'static str, content: Element<'a, Message>) -> Element<'a, Message> { + container(column![text(title).size(10).color(muted_text()), content].spacing(11)) + .padding(14) + .width(Length::FillPortion(1)) + .style(preference_group_style) + .into() +} + +fn mini_chart( + history: &VecDeque, + value: fn(&MetricsPoint) -> f32, + color: Color, +) -> Element<'static, Message> { + let values = history + .iter() + .rev() + .take(120) + .map(value) + .collect::>(); + let maximum = values.iter().copied().fold(1.0_f32, f32::max); + let mut bars = row![].spacing(1).height(54).align_y(Alignment::End); + for value in values.into_iter().rev() { + let height = if value > 0.0 { + (value / maximum * 52.0).max(2.0) + } else { + 1.0 + }; + bars = bars.push( + container(Space::new(Length::Fill, height)) + .width(Length::FillPortion(1)) + .style(move |_| chart_bar_style(color)), + ); + } + if history.is_empty() { + bars = bars.push( + container(text("Waiting for samples…").size(11).color(muted_text())) + .center_x(Length::Fill) + .center_y(Length::Fill), + ); + } + container(bars) + .height(58) + .width(Length::Fill) + .padding([3, 0]) + .into() +} + +fn format_count(value: u64) -> String { + if value >= 1_000_000 { + format!("{:.1}M", value as f64 / 1_000_000.0) + } else if value >= 1_000 { + format!("{:.1}K", value as f64 / 1_000.0) + } else { + value.to_string() + } +} + +fn format_milliseconds(milliseconds: u64) -> String { + if milliseconds >= 1_000 { + format!("{:.2}s", milliseconds as f64 / 1_000.0) + } else { + format!("{milliseconds}ms") + } +} + +fn format_rate(bytes_per_second: f32) -> String { + format!("{}/s", format_bytes(bytes_per_second.max(0.0) as u64)) +} + fn danger_button<'a>(content: impl Into>) -> Button<'a, Message> { button(content).padding([8, 14]).style(danger_button_style) } @@ -1250,6 +1751,53 @@ fn danger_button_style(theme: &Theme, status: button::Status) -> button::Style { style } +fn segmented_control_style(_: &Theme) -> container::Style { + container::Style { + background: Some(Background::Color(Color::from_rgb8(24, 24, 26))), + border: Border { + color: Color::from_rgb8(57, 57, 60), + width: 1.0, + radius: 18.0.into(), + }, + ..container::Style::default() + } +} + +fn segmented_button_style(_: &Theme, status: button::Status, selected: bool) -> button::Style { + let background = if selected { + Some(Background::Color(Color::from_rgb8(55, 55, 58))) + } else if status == button::Status::Hovered { + Some(Background::Color(Color::from_rgb8(39, 39, 42))) + } else { + None + }; + button::Style { + background, + text_color: if selected { Color::WHITE } else { muted_text() }, + border: Border { + radius: 14.0.into(), + ..Border::default() + }, + ..button::Style::default() + } +} + +fn status_badge_style(color: Color) -> container::Style { + container::Style { + background: Some(Background::Color(color.scale_alpha(0.12))), + border: Border { + color: color.scale_alpha(0.45), + width: 1.0, + radius: 12.0.into(), + }, + ..container::Style::default() + } +} + +fn chart_bar_style(color: Color) -> container::Style { + container::Style::default().background(color.scale_alpha(0.82)) +} + fn overview_style(_: &Theme) -> container::Style { container::Style { background: Some(Background::Color(Color::from_rgb8(31, 31, 33))), diff --git a/src/engine.rs b/src/engine.rs index 973aaed..7d036bd 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -3,12 +3,16 @@ mod gguf; mod metal; mod tokenizer; +#[cfg(target_os = "macos")] +use crate::metrics::{KvLookup, Metrics}; use crate::model::ModelChoice; use crate::settings::TurnSettings; use crate::settings::{EngineSettings, ReasoningMode}; use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value}; use sha2::{Digest, Sha256}; use std::path::{Path, PathBuf}; +#[cfg(target_os = "macos")] +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; #[cfg(target_os = "macos")] use std::time::Instant; @@ -256,8 +260,14 @@ impl Model { .encode_conversation(system, messages, reasoning) } - fn render_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec { - self.tokenizer.encode_continuation(prompt, reasoning) + fn render_continuation( + &self, + prompt: &str, + reasoning: ReasoningMode, + skip_previous_eos: bool, + ) -> Vec { + self.tokenizer + .encode_continuation(prompt, reasoning, skip_previous_eos) } pub(crate) fn token_bytes(&self, token: i32) -> Option> { @@ -295,11 +305,13 @@ impl Model { pub(crate) struct Generator { executor: metal::Executor, checkpoint: Option, + metrics: Arc, } #[derive(Clone)] pub(crate) struct ChatTurn { pub(crate) user: bool, + pub(crate) skip_previous_eos: bool, pub(crate) reasoning: Option, pub(crate) reasoning_complete: bool, pub(crate) content: String, @@ -311,11 +323,13 @@ pub(crate) struct GenerationOutput { pub(crate) cached_tokens: u32, pub(crate) completion_tokens: u32, pub(crate) finish_reason: &'static str, + pub(crate) previous_checkpoint_bytes: Option, + pub(crate) checkpoint_bytes: u64, } #[cfg(target_os = "macos")] impl Generator { - pub(crate) fn open(settings: &EngineSettings) -> Result { + pub(crate) fn open(settings: &EngineSettings, metrics: Arc) -> Result { if settings.model != ModelChoice::DeepSeekV4Flash { return Err("local generation currently supports DeepSeek V4 Flash only".into()); } @@ -330,13 +344,19 @@ impl Generator { model, settings.context_tokens.max(1) as u32, settings.execution.quality, + settings.execution.prefill_chunk, )?; Ok(Self { executor, checkpoint: None, + metrics, }) } + pub(crate) fn summary(&self) -> ModelSummary { + self.executor.model().summary() + } + pub(crate) fn generate( &mut self, checkpoint: &Path, @@ -346,12 +366,20 @@ impl Generator { mut emit: impl FnMut(bool, String), mut progress: impl FnMut(u32, u32, Option), ) -> Result { - self.select_checkpoint(checkpoint)?; - let (output, prompt_complete) = + let history = messages + .split_last() + .map_or(messages, |(_, history)| history); + self.select_checkpoint( + checkpoint, + conversation_tag(&settings.system_prompt, settings.reasoning_mode, history), + )?; + let (mut output, prompt_complete) = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; let mut completed = messages.to_vec(); completed.push(output.message.clone()); - self.executor.save_checkpoint( + output.previous_checkpoint_bytes = + std::fs::metadata(checkpoint).ok().map(|item| item.len()); + self.save_checkpoint( checkpoint, if prompt_complete { conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed) @@ -359,6 +387,9 @@ impl Generator { [0; 32] }, )?; + output.checkpoint_bytes = std::fs::metadata(checkpoint) + .map(|item| item.len()) + .unwrap_or(0); Ok(output) } @@ -378,15 +409,17 @@ impl Generator { conversation_tag(&settings.system_prompt, settings.reasoning_mode, history); let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag))); if self.executor.checkpoint_tag() != history_tag { - self.select_checkpoint(&checkpoint)?; + self.select_checkpoint(&checkpoint, history_tag)?; if self.executor.checkpoint_tag() != history_tag { self.executor.reset()?; self.checkpoint = None; let _ = std::fs::remove_file(&checkpoint); } + } else { + self.metrics.kv_lookup(KvLookup::MemoryHit); } - let (output, prompt_complete) = + let (mut output, prompt_complete) = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; let mut completed = messages.to_vec(); completed.push(output.message.clone()); @@ -399,7 +432,10 @@ impl Generator { ) })?; let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag))); - self.executor.save_checkpoint( + output.previous_checkpoint_bytes = std::fs::metadata(&completed_checkpoint) + .ok() + .map(|item| item.len()); + self.save_checkpoint( &completed_checkpoint, if prompt_complete { completed_tag @@ -407,23 +443,62 @@ impl Generator { [0; 32] }, )?; + output.checkpoint_bytes = std::fs::metadata(&completed_checkpoint) + .map(|item| item.len()) + .unwrap_or(0); self.checkpoint = Some(completed_checkpoint); Ok(output) } - fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> { + fn select_checkpoint( + &mut self, + checkpoint: &Path, + expected_tag: [u8; 32], + ) -> Result<(), String> { if self.checkpoint.as_deref() == Some(checkpoint) { + self.metrics + .kv_lookup(if self.executor.checkpoint_tag() == expected_tag { + KvLookup::MemoryHit + } else { + KvLookup::Miss + }); return Ok(()); } self.executor.reset()?; - if self.executor.load_checkpoint(checkpoint).is_err() { - self.executor.reset()?; - let _ = std::fs::remove_file(checkpoint); - } + self.metrics.kv_read_started(); + let started = Instant::now(); + let loaded = self + .executor + .load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes)); + self.metrics + .kv_read_finished(started.elapsed(), loaded.is_err()); + let lookup = match loaded { + Ok(true) if self.executor.checkpoint_tag() == expected_tag => KvLookup::DiskHit, + Ok(true) | Ok(false) => KvLookup::Miss, + Err(_) => { + self.executor.reset()?; + let _ = std::fs::remove_file(checkpoint); + KvLookup::Invalid + } + }; + self.metrics.kv_lookup(lookup); self.checkpoint = Some(checkpoint.to_owned()); Ok(()) } + fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> { + self.metrics.kv_write_started(); + let started = Instant::now(); + let result = self + .executor + .save_checkpoint(checkpoint, tag, &mut |bytes| { + self.metrics.kv_write_bytes(bytes); + }); + self.metrics + .kv_write_finished(started.elapsed(), result.is_err()); + result + } + fn generate_inner( &mut self, messages: &[ChatTurn], @@ -443,11 +518,11 @@ impl Generator { ) => { let mut tokens = self.executor.tokens().to_vec(); - tokens.extend( - self.executor - .model() - .render_continuation(&latest.content, settings.reasoning_mode), - ); + tokens.extend(self.executor.model().render_continuation( + &latest.content, + settings.reasoning_mode, + latest.skip_previous_eos, + )); tokens } _ => self.executor.model().render_conversation( @@ -467,11 +542,13 @@ impl Generator { )); } let reused = self.executor.align_prompt(&tokens)?; + self.metrics.kv_prefix_reused(reused); progress(self.executor.position(), self.executor.context(), None); let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645)); let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct; let mut generated = ChatTurn { user: false, + skip_previous_eos: false, reasoning: reasoning.then(String::new), reasoning_complete: !reasoning, content: String::new(), @@ -479,24 +556,40 @@ impl Generator { let mut emitted_reasoning = 0; let mut emitted_content = 0; let prompt_tokens = tokens.len(); - for (index, token) in tokens.into_iter().enumerate().skip(reused) { - if cancelled.load(Ordering::Relaxed) { - return Ok(( - GenerationOutput { - message: generated, - prompt_tokens: prompt_tokens as u32, - cached_tokens: reused as u32, - completion_tokens: 0, - finish_reason: "stop", - }, - false, - )); - } - self.executor.eval(token)?; - if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens { + let suffix = &tokens[reused..]; + let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 { + let context = self.executor.context(); + self.executor.prefill(suffix, |used| { + progress(used, context, None); + !cancelled.load(Ordering::Relaxed) + })? + } else { + let mut completed = 0; + for &token in suffix { + if cancelled.load(Ordering::Relaxed) { + break; + } + self.executor.eval(token)?; + completed += 1; progress(self.executor.position(), self.executor.context(), None); } + completed + }; + if completed != suffix.len() { + return Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: 0, + finish_reason: "stop", + previous_checkpoint_bytes: None, + checkpoint_bytes: 0, + }, + false, + )); } + progress(self.executor.position(), self.executor.context(), Some(0.0)); let generation_started = Instant::now(); let mut generated_tokens = 0_u32; for _ in 0..settings @@ -519,6 +612,8 @@ impl Generator { cached_tokens: reused as u32, completion_tokens: generated_tokens, finish_reason: "stop", + previous_checkpoint_bytes: None, + checkpoint_bytes: 0, }, true, )); @@ -550,6 +645,8 @@ impl Generator { cached_tokens: reused as u32, completion_tokens: generated_tokens, finish_reason: "stop", + previous_checkpoint_bytes: None, + checkpoint_bytes: 0, }, true, )); @@ -575,6 +672,8 @@ impl Generator { cached_tokens: reused as u32, completion_tokens: generated_tokens + 1, finish_reason: "stop", + previous_checkpoint_bytes: None, + checkpoint_bytes: 0, }, false, )); @@ -615,6 +714,8 @@ impl Generator { cached_tokens: reused as u32, completion_tokens: generated_tokens + 1, finish_reason: "stop", + previous_checkpoint_bytes: None, + checkpoint_bytes: 0, }, false, )); @@ -645,6 +746,8 @@ impl Generator { cached_tokens: reused as u32, completion_tokens: generated_tokens, finish_reason: "length", + previous_checkpoint_bytes: None, + checkpoint_bytes: 0, }, true, )) @@ -891,6 +994,7 @@ mod sampling_tests { fn checkpoint_tag_covers_the_canonical_chat_state() { let messages = [ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: "Hello".into(), @@ -924,7 +1028,7 @@ mod sampling_tests { ReasoningMode::Direct, ); assert_eq!(tokens.len(), 10); - let mut executor = metal::Executor::open(model, 32_768, false).unwrap(); + let mut executor = metal::Executor::open(model, 32_768, false, 0).unwrap(); assert_eq!(executor.context(), 32_768); for &token in &tokens { executor.eval(token).unwrap(); @@ -947,7 +1051,9 @@ mod sampling_tests { let next = argmax.0 as i32; let checkpoint = std::env::temp_dir().join(format!("ds4-rust-kv-{}.bin", std::process::id())); - executor.save_checkpoint(&checkpoint, [7; 32]).unwrap(); + executor + .save_checkpoint(&checkpoint, [7; 32], &mut |_| {}) + .unwrap(); executor.eval(next).unwrap(); let continued_logit = executor.logits()[0]; let continued_argmax = executor @@ -960,8 +1066,8 @@ mod sampling_tests { drop(executor); let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap(); - let mut restored = metal::Executor::open(model, 32_768, false).unwrap(); - assert!(restored.load_checkpoint(&checkpoint).unwrap()); + let mut restored = metal::Executor::open(model, 32_768, false, 0).unwrap(); + assert!(restored.load_checkpoint(&checkpoint, &mut |_| {}).unwrap()); assert_eq!(restored.position(), tokens.len() as u32); assert_eq!(restored.tokens(), tokens); assert_eq!(restored.checkpoint_tag(), [7; 32]); @@ -1872,6 +1978,7 @@ mod tests { model.render_prompt("", "Hello", ReasoningMode::Direct), [0, 128_803, 19_923, 128_804, 128_822] ); + assert_eq!(model.tokenizer.tokenize_rendered("｜DSML｜").len(), 1); assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct)); assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High)); let think_start = *model @@ -1884,18 +1991,21 @@ mod tests { &[ ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: "Hello".into(), }, ChatTurn { user: false, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: "Hello".into(), }, ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: "Hello".into(), @@ -1908,27 +2018,34 @@ mod tests { ] ); assert_eq!( - model.render_continuation("Hello", ReasoningMode::Direct), + model.render_continuation("Hello", ReasoningMode::Direct, false), [1, 128_803, 19_923, 128_804, 128_822] ); + assert_eq!( + model.render_continuation("Hello", ReasoningMode::Direct, true), + [128_803, 19_923, 128_804, 128_822] + ); assert_eq!( model.render_conversation( "", &[ ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: "Hello".into(), }, ChatTurn { user: false, + skip_previous_eos: false, reasoning: Some("Hello".into()), reasoning_complete: true, content: "Hello".into(), }, ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: "Hello".into(), diff --git a/src/engine/gguf.rs b/src/engine/gguf.rs index 322f057..68341dc 100644 --- a/src/engine/gguf.rs +++ b/src/engine/gguf.rs @@ -46,6 +46,8 @@ pub(super) struct Tensor { pub(super) struct Gguf { path: PathBuf, map: Mmap, + data_offset: u64, + max_tensor_bytes: u64, pub(super) metadata: HashMap, pub(super) tensors: HashMap, } @@ -141,6 +143,7 @@ impl Gguf { } let data_start = align(cursor.position() as u64, alignment)?; + let mut max_tensor_bytes = 0; for (name, tensor) in &mut tensors { tensor.offset = data_start .checked_add(tensor.offset) @@ -152,11 +155,14 @@ impl Gguf { if end > map.len() as u64 { return Err(format!("tensor {name} points outside the GGUF file")); } + max_tensor_bytes = max_tensor_bytes.max(tensor.bytes); } Ok(Self { path: path.to_owned(), map, + data_offset: data_start, + max_tensor_bytes, metadata, tensors, }) @@ -174,6 +180,14 @@ impl Gguf { self.map.as_ptr() } + pub(super) fn data_offset(&self) -> u64 { + self.data_offset + } + + pub(super) fn max_tensor_bytes(&self) -> u64 { + self.max_tensor_bytes + } + pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> { self.tensors .get(name) diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 8576e1c..809a769 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -11,6 +11,7 @@ use std::time::UNIX_EPOCH; const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4RKV01"; const CHECKPOINT_VERSION: u32 = 1; const CHECKPOINT_IO_CHUNK: usize = 8 * 1024 * 1024; +const DEFAULT_PREFILL_CHUNK: u32 = 4096; const SOURCES: [(&str, &str); 19] = [ ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"), @@ -74,9 +75,16 @@ struct GpuTensor { unsafe extern "C" { fn ds4_gpu_init() -> i32; fn ds4_gpu_cleanup(); - fn ds4_gpu_set_model_map(model_map: *const c_void, model_size: u64) -> i32; + fn ds4_gpu_set_model_map_range( + model_map: *const c_void, + model_size: u64, + map_offset: u64, + map_size: u64, + max_tensor_bytes: u64, + ) -> i32; fn ds4_gpu_set_quality(quality: bool); fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor; + fn ds4_gpu_tensor_view(base: *const GpuTensor, offset: u64, bytes: u64) -> *mut GpuTensor; fn ds4_gpu_tensor_free(tensor: *mut GpuTensor); fn ds4_gpu_tensor_fill_f32(tensor: *mut GpuTensor, value: f32, count: u64) -> i32; fn ds4_gpu_tensor_read( @@ -108,6 +116,18 @@ unsafe extern "C" { fn ds4_gpu_begin_commands() -> i32; fn ds4_gpu_end_commands() -> i32; + fn ds4_gpu_embed_tokens_hc_tensor( + out: *mut GpuTensor, + tokens: *const GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + vocab: u32, + rows: u32, + embd: u32, + hc: u32, + ) -> i32; + fn ds4_gpu_embed_token_hc_tensor( out: *mut GpuTensor, map: *const c_void, @@ -133,6 +153,18 @@ unsafe extern "C" { n: u32, eps: f32, ) -> i32; + fn ds4_gpu_hc_rms_scale_project_f16_tensor( + out: *mut GpuTensor, + scale: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u32, + output: u32, + x: *const GpuTensor, + rows: u32, + eps: f32, + ) -> i32; fn ds4_gpu_matmul_f16_tensor( out: *mut GpuTensor, map: *const c_void, @@ -226,8 +258,15 @@ unsafe extern "C" { rows: u32, eps: f32, ) -> i32; - fn ds4_gpu_head_rms_norm_rope_tail_tensor( - x: *mut GpuTensor, + fn ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( + out: *mut GpuTensor, + half: *mut GpuTensor, + map: *const c_void, + size: u64, + weight: u64, + input: u64, + output: u64, + x: *const GpuTensor, rows: u32, heads: u32, head_dim: u32, @@ -250,6 +289,23 @@ unsafe extern "C" { head_dim: u32, eps: f32, ) -> i32; + fn ds4_gpu_head_rms_norm_rope_tail_tensor( + x: *mut GpuTensor, + rows: u32, + heads: u32, + head_dim: u32, + rot: u32, + pos: u32, + original_context: u32, + inverse: bool, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + eps: f32, + ) -> i32; fn ds4_gpu_rope_tail_tensor( x: *mut GpuTensor, rows: u32, @@ -284,6 +340,44 @@ unsafe extern "C" { heads: u32, head_dim: u32, ) -> i32; + fn ds4_gpu_attention_decode_raw_batch_heads_tensor( + heads: *mut GpuTensor, + map: *const c_void, + size: u64, + sinks: u64, + q: *const GpuTensor, + raw_kv: *const GpuTensor, + tokens: u32, + pos: u32, + n_raw: u32, + raw_cap: u32, + raw_start: u32, + window: u32, + heads_count: u32, + head_dim: u32, + ) -> i32; + fn ds4_gpu_attention_decode_mixed_batch_heads_tensor( + heads: *mut GpuTensor, + map: *const c_void, + size: u64, + sinks: u64, + q: *const GpuTensor, + raw_kv: *const GpuTensor, + compressed: *const GpuTensor, + compressed_f16: u32, + compressed_mask: *const GpuTensor, + use_mask: u32, + tokens: u32, + pos: u32, + n_raw: u32, + raw_cap: u32, + raw_start: u32, + n_comp: u32, + window: u32, + ratio: u32, + heads_count: u32, + head_dim: u32, + ) -> i32; fn ds4_gpu_attention_indexed_mixed_batch_heads_tensor( heads: *mut GpuTensor, map: *const c_void, @@ -316,6 +410,19 @@ unsafe extern "C" { head_dim: u32, scale: f32, ) -> i32; + fn ds4_gpu_indexer_scores_decode_batch_tensor( + scores: *mut GpuTensor, + q: *const GpuTensor, + weights: *const GpuTensor, + index_comp: *const GpuTensor, + n_comp: u32, + tokens: u32, + pos: u32, + heads: u32, + head_dim: u32, + ratio: u32, + scale: f32, + ) -> i32; fn ds4_gpu_indexer_topk_tensor( selected: *mut GpuTensor, scores: *const GpuTensor, @@ -351,6 +458,18 @@ unsafe extern "C" { rms_eps: f32, state_already_stored: bool, ) -> i32; + fn ds4_gpu_compressor_prefill_state_ratio4_tensor( + state_kv: *mut GpuTensor, + state_score: *mut GpuTensor, + kv_tail: *const GpuTensor, + score_tail: *const GpuTensor, + map: *const c_void, + size: u64, + ape: u64, + ape_type: u32, + head_dim: u32, + pos: u32, + ) -> i32; fn ds4_gpu_dsv4_fp8_kv_quantize_tensor( x: *mut GpuTensor, rows: u32, @@ -365,6 +484,221 @@ unsafe extern "C" { head_dim: u32, rot: u32, ) -> i32; + fn ds4_gpu_store_raw_kv_batch_tensor( + raw_cache: *mut GpuTensor, + kv: *const GpuTensor, + raw_cap: u32, + pos: u32, + rows: u32, + head_dim: u32, + ) -> i32; + fn ds4_gpu_compressor_prefill_tensor( + cache: *mut GpuTensor, + state_kv: *mut GpuTensor, + state_score: *mut GpuTensor, + kv: *const GpuTensor, + score: *const GpuTensor, + map: *const c_void, + size: u64, + ape: u64, + ape_type: u32, + norm: u64, + norm_type: u32, + head_dim: u32, + ratio: u32, + pos: u32, + rows: u32, + rot: u32, + original_context: u32, + quantize_fp8: bool, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + rms_eps: f32, + ) -> i32; + fn ds4_gpu_compressor_prefill_ratio4_replay_tensor( + cache: *mut GpuTensor, + state_kv: *mut GpuTensor, + state_score: *mut GpuTensor, + kv: *const GpuTensor, + score: *const GpuTensor, + map: *const c_void, + size: u64, + ape: u64, + ape_type: u32, + norm: u64, + norm_type: u32, + head_dim: u32, + pos: u32, + rows: u32, + rot: u32, + original_context: u32, + quantize_fp8: bool, + freq_base: f32, + freq_scale: f32, + ext_factor: f32, + attn_factor: f32, + beta_fast: f32, + beta_slow: f32, + rms_eps: f32, + ) -> i32; + fn ds4_gpu_attention_prefill_raw_heads_tensor( + heads: *mut GpuTensor, + map: *const c_void, + size: u64, + sinks: u64, + q: *const GpuTensor, + raw: *const GpuTensor, + rows: u32, + window: u32, + heads_count: u32, + head_dim: u32, + ) -> i32; + fn ds4_gpu_attention_prefill_static_mixed_heads_tensor( + heads: *mut GpuTensor, + map: *const c_void, + size: u64, + sinks: u64, + q: *const GpuTensor, + raw: *const GpuTensor, + compressed: *const GpuTensor, + compressed_f16: u32, + rows: u32, + compressed_rows: u32, + window: u32, + ratio: u32, + heads_count: u32, + head_dim: u32, + ) -> i32; + fn ds4_gpu_indexer_scores_prefill_tensor( + scores: *mut GpuTensor, + q: *const GpuTensor, + weights: *const GpuTensor, + compressed: *const GpuTensor, + compressed_rows: u32, + rows: u32, + heads: u32, + head_dim: u32, + ratio: u32, + scale: f32, + ) -> i32; + fn ds4_gpu_attention_output_q8_batch_f16_tensor( + out_half: *mut GpuTensor, + low: *mut GpuTensor, + map: *const c_void, + size: u64, + weight_a: u64, + weight_b: u64, + group_dim: u64, + rank: u64, + groups: u32, + output: u64, + heads: *const GpuTensor, + rows: u32, + ) -> i32; + fn ds4_gpu_attention_output_q8_batch_tensor( + out: *mut GpuTensor, + low: *mut GpuTensor, + group_scratch: *mut GpuTensor, + low_scratch: *mut GpuTensor, + map: *const c_void, + size: u64, + weight_a: u64, + weight_b: u64, + group_dim: u64, + rank: u64, + groups: u32, + output: u64, + heads: *const GpuTensor, + rows: u32, + ) -> i32; + fn ds4_gpu_router_select_batch_tensor( + selected: *mut GpuTensor, + weights: *mut GpuTensor, + probs: *mut GpuTensor, + map: *const c_void, + size: u64, + bias: u64, + hash: u64, + hash_rows: u32, + expert_groups: u32, + groups_used: u32, + has_bias: bool, + hash_mode: bool, + logits: *const GpuTensor, + tokens: *const GpuTensor, + experts: u32, + experts_used: u32, + scale: f32, + rows: u32, + ) -> i32; + fn ds4_gpu_routed_moe_batch_tensor( + out: *mut GpuTensor, + gate: *mut GpuTensor, + up: *mut GpuTensor, + mid: *mut GpuTensor, + experts_out: *mut GpuTensor, + map: *const c_void, + size: u64, + gate_weight: u64, + up_weight: u64, + down_weight: u64, + gate_type: u32, + down_type: u32, + gate_expert_bytes: u64, + gate_row_bytes: u64, + down_expert_bytes: u64, + down_row_bytes: u64, + input: u32, + middle: u32, + output: u32, + selected: *const GpuTensor, + weights: *const GpuTensor, + total_experts: u32, + used_experts: u32, + clamp: f32, + x: *const GpuTensor, + layer: u32, + rows: u32, + mid_f16: *mut bool, + force_resident: bool, + ) -> i32; + fn ds4_gpu_swiglu_tensor( + out: *mut GpuTensor, + gate: *const GpuTensor, + up: *const GpuTensor, + count: u32, + clamp: f32, + scale: f32, + ) -> i32; + fn ds4_gpu_hc_expand_split_half_tensor( + out: *mut GpuTensor, + block_half: *const GpuTensor, + residual: *const GpuTensor, + split: *const GpuTensor, + embd: u32, + hc: u32, + ) -> i32; + fn ds4_gpu_hc_expand_split_tensor( + out: *mut GpuTensor, + block: *const GpuTensor, + residual: *const GpuTensor, + split: *const GpuTensor, + embd: u32, + hc: u32, + ) -> i32; + fn ds4_gpu_hc_expand_add_split_tensor( + out: *mut GpuTensor, + block: *const GpuTensor, + add: *const GpuTensor, + residual: *const GpuTensor, + split: *const GpuTensor, + embd: u32, + hc: u32, + ) -> i32; fn ds4_gpu_attention_output_low_q8_tensor( low: *mut GpuTensor, map: *const c_void, @@ -502,8 +836,17 @@ struct Context; impl Context { fn open(model: &Model, quality: bool) -> Result { check(unsafe { ds4_gpu_init() }, "Metal initialization")?; + let data_offset = model.main.data_offset(); if let Err(error) = check( - unsafe { ds4_gpu_set_model_map(model.main.map_ptr().cast(), model.main.len()) }, + unsafe { + ds4_gpu_set_model_map_range( + model.main.map_ptr().cast(), + model.main.len(), + data_offset, + model.main.len() - data_offset, + model.main.max_tensor_bytes(), + ) + }, "model mapping", ) { unsafe { ds4_gpu_cleanup() }; @@ -561,6 +904,12 @@ impl Buffer { .ok_or_else(|| format!("Metal could not allocate {bytes} bytes")) } + fn view(&self, offset: u64, bytes: u64) -> Result { + NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) }) + .map(Self) + .ok_or_else(|| "Metal could not create a tensor view".to_owned()) + } + fn read_f32(&self, values: &mut [f32]) -> Result<(), String> { check( unsafe { @@ -603,6 +952,20 @@ impl Buffer { ) } + fn write_i32(&self, values: &[i32]) -> Result<(), String> { + check( + unsafe { + ds4_gpu_tensor_write( + self.raw(), + 0, + values.as_ptr().cast(), + std::mem::size_of_val(values) as u64, + ) + }, + "uploading tokens", + ) + } + fn fill(&self, value: f32, count: u64) -> Result<(), String> { call( unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) }, @@ -854,6 +1217,106 @@ struct Scratch { logits: Buffer, } +struct BatchScratch { + tokens: Buffer, + current_hc: Buffer, + next_hc: Buffer, + after_attention_hc: Buffer, + flat_hc: Buffer, + hc_mix: Buffer, + hc_split: Buffer, + current: Buffer, + norm: Buffer, + q_rank: Buffer, + q_rank_norm: Buffer, + q: Buffer, + q_half: Buffer, + kv_raw: Buffer, + kv: Buffer, + compressed_kv: Buffer, + compressed_score: Buffer, + compressed_stage: Buffer, + indexer_q: Buffer, + indexer_weights: Buffer, + indexer_scores: Buffer, + indexer_selected: Buffer, + heads: Buffer, + attention_low: Buffer, + attention_group_tmp: Buffer, + attention_low_tmp: Buffer, + attention_out: Buffer, + router_logits: Buffer, + router_probs: Buffer, + router_selected: Buffer, + router_weights: Buffer, + routed_gate: Buffer, + routed_up: Buffer, + routed_mid: Buffer, + routed_experts: Buffer, + routed_out: Buffer, + shared_gate: Buffer, + shared_up: Buffer, + shared_mid: Buffer, + shared_out: Buffer, +} + +impl BatchScratch { + fn allocate(model: &Model, pos: u32, rows: u32) -> Result { + let shape = model.shape; + let rows = u64::from(rows); + let hc_dim = shape.hc * shape.embd; + let mix_hc = 2 * shape.hc + shape.hc * shape.hc; + let q_dim = shape.heads * shape.head_dim; + let group_dim = shape.head_dim * (shape.heads / shape.out_groups); + let low_dim = shape.out_groups * shape.lora_o; + let routed = shape.experts_used * shape.ff_expert; + Ok(Self { + tokens: Buffer::bytes(rows * 4)?, + current_hc: Buffer::floats(rows * hc_dim)?, + next_hc: Buffer::floats(rows * hc_dim)?, + after_attention_hc: Buffer::floats(rows * hc_dim)?, + flat_hc: Buffer::floats(rows * hc_dim)?, + hc_mix: Buffer::floats(rows * mix_hc)?, + hc_split: Buffer::floats(rows * mix_hc)?, + current: Buffer::floats(rows * shape.embd)?, + norm: Buffer::floats(rows * shape.embd)?, + q_rank: Buffer::floats(rows * shape.lora_q)?, + q_rank_norm: Buffer::floats(rows * shape.lora_q)?, + q: Buffer::floats(rows * q_dim)?, + q_half: Buffer::bytes(rows * q_dim.max(shape.embd) * 2)?, + kv_raw: Buffer::floats(rows * shape.head_dim)?, + kv: Buffer::floats(rows * shape.head_dim)?, + compressed_kv: Buffer::floats(rows * 2 * shape.head_dim)?, + compressed_score: Buffer::floats(rows * 2 * shape.head_dim)?, + compressed_stage: Buffer::floats(rows * shape.head_dim)?, + indexer_q: Buffer::floats(rows * shape.indexer_heads * shape.indexer_head_dim)?, + indexer_weights: Buffer::floats(rows * shape.indexer_heads)?, + indexer_scores: Buffer::floats( + rows * (u64::from(pos) + rows).div_ceil(4).saturating_add(2), + )?, + indexer_selected: Buffer::bytes(rows * shape.indexer_top_k * 4)?, + heads: Buffer::floats(rows * q_dim)?, + attention_low: Buffer::floats(rows * low_dim)?, + attention_group_tmp: Buffer::floats(rows * group_dim)?, + attention_low_tmp: Buffer::floats(rows * shape.lora_o)?, + attention_out: Buffer::floats(rows * shape.embd)?, + router_logits: Buffer::floats(rows * shape.experts)?, + router_probs: Buffer::floats(rows * shape.experts)?, + router_selected: Buffer::bytes(rows * shape.experts_used * 4)?, + router_weights: Buffer::floats(rows * shape.experts_used)?, + routed_gate: Buffer::floats(rows * routed)?, + routed_up: Buffer::floats(rows * routed)?, + routed_mid: Buffer::floats(rows * routed)?, + routed_experts: Buffer::floats(rows * shape.experts_used * shape.embd)?, + routed_out: Buffer::floats(rows * shape.embd)?, + shared_gate: Buffer::floats(rows * shape.ff_expert)?, + shared_up: Buffer::floats(rows * shape.ff_expert)?, + shared_mid: Buffer::floats(rows * shape.ff_expert)?, + shared_out: Buffer::floats(rows * shape.embd)?, + }) + } +} + impl Scratch { fn allocate(model: &Model, context: u32) -> Result { let shape = model.shape; @@ -913,6 +1376,7 @@ pub(super) struct Session { scratch: Scratch, layers: Vec, context: u32, + prefill_cap: u32, raw_cap: u32, position: u32, } @@ -979,11 +1443,19 @@ impl LayerState { } impl Session { - fn new(model: &Model, context: u32) -> Result { + fn new(model: &Model, context: u32, prefill_chunk: u32) -> Result { if context == 0 { return Err("context must contain at least one token".into()); } - let raw_cap = (model.shape.sliding_window as u32).min(context); + let prefill_cap = prefill_chunk.max(1).min(context); + let raw_cap = prefill_cap + .saturating_add(model.shape.sliding_window as u32) + .min(context) + .div_ceil(256) + .saturating_mul(256) + .min(8192) + .min(context) + .max(1); let layers = (0..model.shape.layers) .map(|index| LayerState::allocate(model, index, context, raw_cap)) .collect::>()?; @@ -991,6 +1463,7 @@ impl Session { scratch: Scratch::allocate(model, context)?, layers, context, + prefill_cap, raw_cap, position: 0, }) @@ -1010,10 +1483,23 @@ pub(super) struct Executor { } impl Executor { - pub(super) fn open(model: Model, context: u32, quality: bool) -> Result { + pub(super) fn open( + model: Model, + context: u32, + quality: bool, + prefill_chunk: u32, + ) -> Result { let weights = Weights::bind(&model)?; let context_handle = Context::open(&model, quality)?; - let session = Session::new(&model, context)?; + let session = Session::new( + &model, + context, + if prefill_chunk == 0 { + DEFAULT_PREFILL_CHUNK + } else { + prefill_chunk + }, + )?; let model_modified = fs::metadata(model.main.path()) .and_then(|metadata| metadata.modified()) .ok() @@ -1052,6 +1538,138 @@ impl Executor { Ok(()) } + pub(super) fn prefill( + &mut self, + tokens: &[i32], + mut progress: impl FnMut(u32) -> bool, + ) -> Result { + if tokens.len() < 4 && self.session.position != 0 { + for (index, &token) in tokens.iter().enumerate() { + self.eval(token)?; + if !progress(self.session.position) { + return Ok(index + 1); + } + } + return Ok(tokens.len()); + } + + let end = self + .session + .position + .checked_add(u32::try_from(tokens.len()).map_err(|_| "prefill is too large")?) + .ok_or("prefill position overflow")?; + if end > self.session.context { + return Err(format!( + "the Rust Metal executor currently supports {} tokens per session", + self.session.context + )); + } + let mut consumed = 0_usize; + while consumed < tokens.len() { + if !progress(self.session.position) { + return Ok(consumed); + } + let pos = self.session.position; + let mut cap = self.session.prefill_cap; + if pos != 0 { + cap = cap.min(self.session.raw_cap); + let offset = pos % self.session.prefill_cap; + if offset != 0 { + cap = cap.min(self.session.prefill_cap - offset); + } + } + let rows = (tokens.len() - consumed).min(cap as usize); + self.eval_batch(&tokens[consumed..consumed + rows])?; + consumed += rows; + if !progress(self.session.position) { + return Ok(consumed); + } + } + Ok(tokens.len()) + } + + fn eval_batch(&mut self, tokens: &[i32]) -> Result<(), String> { + let rows = u32::try_from(tokens.len()).map_err(|_| "prefill batch is too large")?; + if rows == 0 || rows > self.session.prefill_cap { + return Err("prefill batch exceeds the configured prefill workspace".into()); + } + if tokens + .iter() + .any(|token| *token < 0 || *token as u64 >= self.model.shape.vocab) + { + return Err("prefill contains a token outside the vocabulary".into()); + } + let mut batch = BatchScratch::allocate(&self.model, self.session.position, rows)?; + batch.tokens.write_i32(tokens)?; + let map = self.model.main.map_ptr().cast(); + let size = self.model.main.len(); + let shape = self.model.shape; + let pos = self.session.position; + + let commands = Commands::begin()?; + call( + unsafe { + ds4_gpu_embed_tokens_hc_tensor( + batch.current_hc.raw(), + batch.tokens.raw(), + map, + size, + self.weights.token_embedding.offset, + shape.vocab as u32, + rows, + shape.embd as u32, + shape.hc as u32, + ) + }, + "batch token embedding", + )?; + commands.finish()?; + + for (index, (weights, state)) in self + .weights + .layers + .iter() + .zip(&mut self.session.layers) + .enumerate() + { + let commands = Commands::begin()?; + encode_batch_layer( + &batch, + state, + weights, + shape, + map, + size, + index as u32, + pos, + rows, + self.session.raw_cap, + )?; + commands.finish()?; + std::mem::swap(&mut batch.current_hc, &mut batch.next_hc); + } + + let commands = Commands::begin()?; + call( + unsafe { + ds4_gpu_tensor_copy( + self.session.scratch.current_hc.raw(), + 0, + batch.current_hc.raw(), + u64::from(rows - 1) * shape.hc * shape.embd * 4, + shape.hc * shape.embd * 4, + ) + }, + "selecting the final prefill row", + )?; + encode_output(&self.session.scratch, &self.weights, shape, map, size)?; + commands.finish()?; + self.session.scratch.logits.read_f32(&mut self.logits)?; + self.session.position += rows; + self.tokens.extend_from_slice(tokens); + Ok(()) + } + pub(super) fn logits(&self) -> &[f32] { &self.logits } @@ -1069,7 +1687,7 @@ impl Executor { } pub(super) fn reset(&mut self) -> Result<(), String> { - self.session = Session::new(&self.model, self.session.context)?; + self.session = Session::new(&self.model, self.session.context, self.session.prefill_cap)?; self.tokens.clear(); self.checkpoint_tag = [0; 32]; Ok(()) @@ -1090,28 +1708,68 @@ impl Executor { self.checkpoint_tag } - pub(super) fn save_checkpoint(&self, path: &Path, tag: [u8; 32]) -> Result<(), String> { + pub(super) fn save_checkpoint( + &mut self, + path: &Path, + tag: [u8; 32], + progress: &mut impl FnMut(u64), + ) -> Result<(), String> { if let Some(parent) = path.parent() { fs::create_dir_all(parent).map_err(|error| error.to_string())?; } let temporary = path.with_extension("tmp"); let mut file = File::create(&temporary).map_err(|error| error.to_string())?; - self.write_checkpoint(&mut file, tag)?; + let mut reported = 0_u64; + let mut pending = 0_u64; + self.write_checkpoint(&mut file, tag, &mut |bytes| { + reported += bytes; + pending += bytes; + if pending >= CHECKPOINT_IO_CHUNK as u64 { + progress(pending); + pending = 0; + } + })?; + progress(pending); + let total = file.metadata().map_err(|error| error.to_string())?.len(); + progress(total.saturating_sub(reported)); file.sync_all().map_err(|error| error.to_string())?; - fs::rename(&temporary, path).map_err(|error| error.to_string()) + fs::rename(&temporary, path).map_err(|error| error.to_string())?; + self.checkpoint_tag = tag; + Ok(()) } - pub(super) fn load_checkpoint(&mut self, path: &Path) -> Result { + pub(super) fn load_checkpoint( + &mut self, + path: &Path, + progress: &mut impl FnMut(u64), + ) -> Result { let mut file = match File::open(path) { Ok(file) => file, Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), Err(error) => return Err(error.to_string()), }; - self.read_checkpoint(&mut file)?; + let mut reported = 0_u64; + let mut pending = 0_u64; + self.read_checkpoint(&mut file, &mut |bytes| { + reported += bytes; + pending += bytes; + if pending >= CHECKPOINT_IO_CHUNK as u64 { + progress(pending); + pending = 0; + } + })?; + progress(pending); + let total = file.metadata().map_err(|error| error.to_string())?.len(); + progress(total.saturating_sub(reported)); Ok(true) } - fn write_checkpoint(&self, file: &mut File, tag: [u8; 32]) -> Result<(), String> { + fn write_checkpoint( + &self, + file: &mut File, + tag: [u8; 32], + progress: &mut impl FnMut(u64), + ) -> Result<(), String> { let shape = self.model.shape; file.write_all(CHECKPOINT_MAGIC) .map_err(|error| error.to_string())?; @@ -1168,6 +1826,7 @@ impl Executor { u64::from(physical) * shape.head_dim * 4, shape.head_dim * 4, &mut chunk, + progress, )?; } if let Some(state) = &layer.compression { @@ -1177,10 +1836,11 @@ impl Executor { 0, u64::from(state.rows) * shape.head_dim * 2, &mut chunk, + progress, )?; let bytes = compressor_state_bytes(state.ratio, shape.head_dim); - write_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; - write_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?; + write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?; } if let Some(state) = &layer.indexer { write_buffer( @@ -1189,16 +1849,21 @@ impl Executor { 0, u64::from(state.rows) * shape.indexer_head_dim * 4, &mut chunk, + progress, )?; let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim); - write_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; - write_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + write_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?; + write_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?; } } Ok(()) } - fn read_checkpoint(&mut self, file: &mut File) -> Result<(), String> { + fn read_checkpoint( + &mut self, + file: &mut File, + progress: &mut impl FnMut(u64), + ) -> Result<(), String> { let mut magic = [0; 8]; file.read_exact(&mut magic) .map_err(|error| error.to_string())?; @@ -1290,6 +1955,7 @@ impl Executor { u64::from(physical) * shape.head_dim * 4, shape.head_dim * 4, &mut chunk, + progress, )?; } if let Some(state) = &mut layer.compression { @@ -1300,10 +1966,11 @@ impl Executor { 0, u64::from(state.rows) * shape.head_dim * 2, &mut chunk, + progress, )?; let bytes = compressor_state_bytes(state.ratio, shape.head_dim); - read_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; - read_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?; + read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?; } if let Some(state) = &mut layer.indexer { state.rows = indexer_rows[index]; @@ -1313,10 +1980,11 @@ impl Executor { 0, u64::from(state.rows) * shape.indexer_head_dim * 4, &mut chunk, + progress, )?; let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim); - read_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; - read_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + read_buffer(file, &state.state_kv, 0, bytes, &mut chunk, progress)?; + read_buffer(file, &state.state_score, 0, bytes, &mut chunk, progress)?; } } let mut trailing = [0]; @@ -1374,6 +2042,1283 @@ impl Executor { } } +#[allow(clippy::too_many_arguments)] +fn refresh_ratio4_compressor_state( + s: &BatchScratch, + state: &CompressionState, + weights: CompressorWeights, + shape: super::Shape, + map: *const c_void, + size: u64, + pos: u32, + rows: u32, + head_dim: u32, +) -> Result<(), String> { + if rows < 4 { + return Ok(()); + } + let width = 2 * u64::from(head_dim); + let tail = s + .norm + .view(u64::from(rows - 4) * shape.embd * 4, 4 * shape.embd * 4)?; + f16_rows( + &s.compressed_kv, + weights.kv, + shape.embd, + width, + &tail, + 4, + map, + size, + )?; + f16_rows( + &s.compressed_score, + weights.gate, + shape.embd, + width, + &tail, + 4, + map, + size, + )?; + call( + unsafe { + ds4_gpu_compressor_prefill_state_ratio4_tensor( + state.state_kv.raw(), + state.state_score.raw(), + s.compressed_kv.raw(), + s.compressed_score.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + head_dim, + pos + rows - 4, + ) + }, + "refreshing ratio-4 compressor state", + ) +} + +#[allow(clippy::too_many_arguments)] +fn compress_attention_batch( + s: &BatchScratch, + state: &mut CompressionState, + weights: CompressorWeights, + shape: super::Shape, + map: *const c_void, + size: u64, + pos: u32, + rows: u32, + original: u32, + freq_base: f32, + freq_scale: f32, + ext: f32, + attn_factor: f32, +) -> Result { + let ratio = state.ratio; + let chunk = rows / ratio; + if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) { + let before = if pos == 0 { 0 } else { state.rows }; + let target = s + .compressed_stage + .view(0, u64::from(chunk) * shape.head_dim * 4)?; + let result = if pos == 0 || ratio != 4 { + unsafe { + ds4_gpu_compressor_prefill_tensor( + target.raw(), + state.state_kv.raw(), + state.state_score.raw(), + s.compressed_kv.raw(), + s.compressed_score.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.head_dim as u32, + ratio, + pos, + rows, + shape.rot as u32, + original, + true, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + } + } else { + unsafe { + ds4_gpu_compressor_prefill_ratio4_replay_tensor( + target.raw(), + state.state_kv.raw(), + state.state_score.raw(), + s.compressed_kv.raw(), + s.compressed_score.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.head_dim as u32, + pos, + rows, + shape.rot as u32, + original, + true, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + } + }; + call(result, "batch attention compression")?; + if chunk != 0 { + call( + unsafe { + ds4_gpu_tensor_copy_f32_to_f16( + state.cache.raw(), + u64::from(before) * shape.head_dim * 2, + target.raw(), + 0, + u64::from(chunk) * shape.head_dim, + ) + }, + "committing compressed attention KV", + )?; + } + state.rows = before + chunk; + if ratio == 4 { + refresh_ratio4_compressor_state( + s, + state, + weights, + shape, + map, + size, + pos, + rows, + shape.head_dim as u32, + )?; + } + } else { + let width = if ratio == 4 { 2 } else { 1 } * shape.head_dim; + let target = s.compressed_stage.view(0, shape.head_dim * 4)?; + for row in 0..rows { + let absolute = pos + row; + let emit = (absolute + 1).is_multiple_of(ratio); + let kv = s + .compressed_kv + .view(u64::from(row) * width * 4, width * 4)?; + let score = s + .compressed_score + .view(u64::from(row) * width * 4, width * 4)?; + call( + unsafe { + ds4_gpu_compressor_update_tensor( + kv.raw(), + score.raw(), + state.state_kv.raw(), + state.state_score.raw(), + target.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.head_dim as u32, + ratio, + absolute, + 0, + shape.rot as u32, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + false, + ) + }, + "updating batched attention compression", + )?; + if emit { + call( + unsafe { + ds4_gpu_dsv4_fp8_kv_quantize_tensor( + target.raw(), + 1, + shape.head_dim as u32, + shape.rot as u32, + ) + }, + "rounding compressed attention KV", + )?; + call( + unsafe { + ds4_gpu_tensor_copy_f32_to_f16( + state.cache.raw(), + u64::from(state.rows) * shape.head_dim * 2, + target.raw(), + 0, + shape.head_dim, + ) + }, + "committing compressed attention KV", + )?; + state.rows += 1; + } + } + } + Ok(state.rows) +} + +#[allow(clippy::too_many_arguments)] +fn compress_index_batch( + s: &BatchScratch, + state: &mut CompressionState, + weights: CompressorWeights, + shape: super::Shape, + map: *const c_void, + size: u64, + pos: u32, + rows: u32, + original: u32, + freq_base: f32, + freq_scale: f32, + ext: f32, + attn_factor: f32, +) -> Result<(), String> { + let ratio = state.ratio; + if pos == 0 || (pos.is_multiple_of(ratio) && rows.is_multiple_of(ratio)) { + let before = if pos == 0 { 0 } else { state.rows }; + let chunk = rows / ratio; + let target = state.cache.view( + u64::from(before) * shape.indexer_head_dim * 4, + u64::from(chunk) * shape.indexer_head_dim * 4, + )?; + let result = if pos == 0 { + unsafe { + ds4_gpu_compressor_prefill_tensor( + target.raw(), + state.state_kv.raw(), + state.state_score.raw(), + s.compressed_kv.raw(), + s.compressed_score.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.indexer_head_dim as u32, + ratio, + pos, + rows, + shape.rot as u32, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + } + } else { + unsafe { + ds4_gpu_compressor_prefill_ratio4_replay_tensor( + target.raw(), + state.state_kv.raw(), + state.state_score.raw(), + s.compressed_kv.raw(), + s.compressed_score.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.indexer_head_dim as u32, + pos, + rows, + shape.rot as u32, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + } + }; + call(result, "batch indexer compression")?; + if chunk != 0 { + call( + unsafe { + ds4_gpu_dsv4_indexer_qat_tensor( + target.raw(), + chunk, + shape.indexer_head_dim as u32, + ) + }, + "batch compressed index rounding", + )?; + } + state.rows = before + chunk; + refresh_ratio4_compressor_state( + s, + state, + weights, + shape, + map, + size, + pos, + rows, + shape.indexer_head_dim as u32, + )?; + } else { + let width = 2 * shape.indexer_head_dim; + for row in 0..rows { + let absolute = pos + row; + let emit = (absolute + 1).is_multiple_of(ratio); + let kv = s + .compressed_kv + .view(u64::from(row) * width * 4, width * 4)?; + let score = s + .compressed_score + .view(u64::from(row) * width * 4, width * 4)?; + let target = state.cache.view( + u64::from(state.rows) * shape.indexer_head_dim * 4, + shape.indexer_head_dim * 4, + )?; + call( + unsafe { + ds4_gpu_compressor_update_tensor( + kv.raw(), + score.raw(), + state.state_kv.raw(), + state.state_score.raw(), + target.raw(), + map, + size, + weights.ape.offset, + weights.ape.kind, + weights.norm.offset, + weights.norm.kind, + shape.indexer_head_dim as u32, + ratio, + absolute, + 0, + shape.rot as u32, + original, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + false, + ) + }, + "updating batched indexer compression", + )?; + if emit { + call( + unsafe { + ds4_gpu_dsv4_indexer_qat_tensor( + target.raw(), + 1, + shape.indexer_head_dim as u32, + ) + }, + "rounding compressed index", + )?; + state.rows += 1; + } + } + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn encode_batch_layer( + s: &BatchScratch, + state: &mut LayerState, + w: &Layer, + shape: super::Shape, + map: *const c_void, + size: u64, + layer: u32, + pos: u32, + rows: u32, + raw_cap: u32, +) -> Result<(), String> { + let hc_dim = shape.hc * shape.embd; + let mix_hc = 2 * shape.hc + shape.hc * shape.hc; + let q_dim = shape.heads * shape.head_dim; + let compressed = layer >= 2; + let ratio = compression_ratio(layer); + let freq_base = if compressed { + shape.compress_rope_base + } else { + shape.rope_base + }; + let freq_scale = if compressed { + 1.0 / shape.rope_scale + } else { + 1.0 + }; + let ext = if compressed && shape.rope_scale > 1.0 { + 1.0 + } else { + 0.0 + }; + let attn_factor = if ext != 0.0 { + 1.0 / (1.0 + 0.1 * (1.0 / freq_scale).ln()) + } else { + 1.0 + }; + let original = if compressed { + shape.original_context as u32 + } else { + 0 + }; + + call( + unsafe { + ds4_gpu_hc_rms_scale_project_f16_tensor( + s.hc_mix.raw(), + s.flat_hc.raw(), + map, + size, + w.hc_attn_fn.offset, + hc_dim as u32, + mix_hc as u32, + s.current_hc.raw(), + rows, + shape.rms_epsilon, + ) + }, + "batch attention HC projection", + )?; + call( + unsafe { + ds4_gpu_hc_split_weighted_sum_norm_tensor( + s.current.raw(), + s.norm.raw(), + s.hc_split.raw(), + s.hc_mix.raw(), + s.current_hc.raw(), + map, + size, + w.hc_attn_scale.offset, + w.hc_attn_base.offset, + w.attn_norm.offset, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.hc_epsilon, + shape.rms_epsilon, + ) + }, + "batch attention HC mix", + )?; + q8_rows( + &s.q_rank, + w.attn_q_a, + shape.embd, + shape.lora_q, + &s.norm, + rows, + map, + size, + )?; + q8_rows( + &s.kv_raw, + w.attn_kv, + shape.embd, + shape.head_dim, + &s.norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_dsv4_qkv_rms_norm_rows_tensor( + s.q_rank_norm.raw(), + s.q_rank.raw(), + map, + size, + w.attn_q_a_norm.offset, + shape.lora_q as u32, + s.kv.raw(), + s.kv_raw.raw(), + w.attn_kv_norm.offset, + shape.head_dim as u32, + rows, + shape.rms_epsilon, + ) + }, + "batch Q/KV norm", + )?; + let fused_q = unsafe { + ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor( + s.q.raw(), + s.q_half.raw(), + map, + size, + w.attn_q_b.offset, + shape.lora_q, + q_dim, + s.q_rank_norm.raw(), + rows, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + shape.rms_epsilon, + ) + } != 0; + if !fused_q { + q8_rows( + &s.q, + w.attn_q_b, + shape.lora_q, + q_dim, + &s.q_rank_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_head_rms_norm_tensor( + s.q.raw(), + rows, + shape.heads as u32, + shape.head_dim as u32, + shape.rms_epsilon, + ) + }, + "batch Q norm", + )?; + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.q.raw(), + rows, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "batch Q RoPE", + )?; + } + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.kv.raw(), + rows, + shape.head_kv as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "batch KV RoPE", + )?; + call( + unsafe { + ds4_gpu_dsv4_fp8_kv_quantize_tensor( + s.kv.raw(), + rows, + shape.head_dim as u32, + shape.rot as u32, + ) + }, + "batch KV rounding", + )?; + call( + unsafe { + ds4_gpu_store_raw_kv_batch_tensor( + state.raw_cache.raw(), + s.kv.raw(), + raw_cap, + pos, + rows, + shape.head_dim as u32, + ) + }, + "batch raw KV store", + )?; + + let mut compressed_rows = 0; + if let (Some(weights), Some(compression)) = + (w.attn_compressor.as_ref(), state.compression.as_mut()) + { + let width = if ratio == 4 { 2 } else { 1 } * shape.head_dim; + f16_rows( + &s.compressed_kv, + weights.kv, + shape.embd, + width, + &s.norm, + rows, + map, + size, + )?; + f16_rows( + &s.compressed_score, + weights.gate, + shape.embd, + width, + &s.norm, + rows, + map, + size, + )?; + compressed_rows = compress_attention_batch( + s, + compression, + *weights, + shape, + map, + size, + pos, + rows, + original, + freq_base, + freq_scale, + ext, + attn_factor, + )?; + } + + if ratio == 4 { + let weights = w + .indexer + .ok_or("ratio-4 layer is missing indexer weights")?; + let indexer = state + .indexer + .as_mut() + .ok_or("ratio-4 layer is missing indexer state")?; + let width = 2 * shape.indexer_head_dim; + f16_rows( + &s.compressed_kv, + weights.compressor.kv, + shape.embd, + width, + &s.norm, + rows, + map, + size, + )?; + f16_rows( + &s.compressed_score, + weights.compressor.gate, + shape.embd, + width, + &s.norm, + rows, + map, + size, + )?; + compress_index_batch( + s, + indexer, + weights.compressor, + shape, + map, + size, + pos, + rows, + original, + freq_base, + freq_scale, + ext, + attn_factor, + )?; + matmul_rows( + &s.indexer_q, + weights.q, + shape.lora_q, + shape.indexer_heads * shape.indexer_head_dim, + &s.q_rank_norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.indexer_q.raw(), + rows, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + shape.rot as u32, + pos, + original, + false, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "batch indexer Q RoPE", + )?; + call( + unsafe { + ds4_gpu_dsv4_indexer_qat_tensor( + s.indexer_q.raw(), + rows * shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + ) + }, + "batch indexer Q rounding", + )?; + f16_rows( + &s.indexer_weights, + weights.proj, + shape.embd, + shape.indexer_heads, + &s.norm, + rows, + map, + size, + )?; + } + + if pos == 0 && compressed_rows == 0 { + call( + unsafe { + ds4_gpu_attention_prefill_raw_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + s.kv.raw(), + rows, + shape.sliding_window as u32, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "batch raw prefill attention", + )?; + } else if pos == 0 && ratio == 4 && compressed_rows > shape.indexer_top_k as u32 { + let indexer = state.indexer.as_ref().unwrap(); + call( + unsafe { + ds4_gpu_indexer_scores_prefill_tensor( + s.indexer_scores.raw(), + s.indexer_q.raw(), + s.indexer_weights.raw(), + indexer.cache.raw(), + compressed_rows, + rows, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + ratio, + ((shape.indexer_heads * shape.indexer_head_dim) as f32) + .sqrt() + .recip(), + ) + }, + "batch indexer scoring", + )?; + call( + unsafe { + ds4_gpu_indexer_topk_tensor( + s.indexer_selected.raw(), + s.indexer_scores.raw(), + compressed_rows, + rows, + shape.indexer_top_k as u32, + ) + }, + "batch indexer top-k", + )?; + let compression = state.compression.as_ref().unwrap(); + call( + unsafe { + ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + state.raw_cache.raw(), + compression.cache.raw(), + 1, + s.indexer_selected.raw(), + rows, + pos, + rows, + raw_cap, + 0, + compressed_rows, + shape.indexer_top_k as u32, + shape.sliding_window as u32, + ratio, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "batch indexed prefill attention", + )?; + } else if pos == 0 { + let compression = state.compression.as_ref().unwrap(); + call( + unsafe { + ds4_gpu_attention_prefill_static_mixed_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + s.kv.raw(), + compression.cache.raw(), + 1, + rows, + compressed_rows, + shape.sliding_window as u32, + ratio, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "batch mixed prefill attention", + )?; + } else { + let (n_raw, raw_start) = raw_batch_span(pos, rows, raw_cap, shape.sliding_window as u32); + if ratio == 0 { + call( + unsafe { + ds4_gpu_attention_decode_raw_batch_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + state.raw_cache.raw(), + rows, + pos, + n_raw, + raw_cap, + raw_start, + shape.sliding_window as u32, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "batch resumed raw attention", + )?; + } else if ratio == 4 && compressed_rows > shape.indexer_top_k as u32 { + let indexer = state.indexer.as_ref().unwrap(); + call( + unsafe { + ds4_gpu_indexer_scores_decode_batch_tensor( + s.indexer_scores.raw(), + s.indexer_q.raw(), + s.indexer_weights.raw(), + indexer.cache.raw(), + compressed_rows, + rows, + pos, + shape.indexer_heads as u32, + shape.indexer_head_dim as u32, + ratio, + ((shape.indexer_heads * shape.indexer_head_dim) as f32) + .sqrt() + .recip(), + ) + }, + "batch resumed indexer scoring", + )?; + call( + unsafe { + ds4_gpu_indexer_topk_tensor( + s.indexer_selected.raw(), + s.indexer_scores.raw(), + compressed_rows, + rows, + shape.indexer_top_k as u32, + ) + }, + "batch resumed indexer top-k", + )?; + let compression = state.compression.as_ref().unwrap(); + call( + unsafe { + ds4_gpu_attention_indexed_mixed_batch_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + state.raw_cache.raw(), + compression.cache.raw(), + 1, + s.indexer_selected.raw(), + rows, + pos, + n_raw, + raw_cap, + raw_start, + compressed_rows, + shape.indexer_top_k as u32, + shape.sliding_window as u32, + ratio, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "batch resumed indexed attention", + )?; + } else { + let compression = state.compression.as_ref().unwrap(); + call( + unsafe { + ds4_gpu_attention_decode_mixed_batch_heads_tensor( + s.heads.raw(), + map, + size, + w.attn_sinks.offset, + s.q.raw(), + state.raw_cache.raw(), + compression.cache.raw(), + 1, + std::ptr::null(), + 0, + rows, + pos, + n_raw, + raw_cap, + raw_start, + compressed_rows, + shape.sliding_window as u32, + ratio, + shape.heads as u32, + shape.head_dim as u32, + ) + }, + "batch resumed mixed attention", + )?; + } + } + call( + unsafe { + ds4_gpu_rope_tail_tensor( + s.heads.raw(), + rows, + shape.heads as u32, + shape.head_dim as u32, + shape.rot as u32, + pos, + original, + true, + freq_base, + freq_scale, + ext, + attn_factor, + shape.rope_beta_fast, + shape.rope_beta_slow, + ) + }, + "batch inverse attention RoPE", + )?; + let group_dim = shape.head_dim * (shape.heads / shape.out_groups); + let half_output = unsafe { + ds4_gpu_attention_output_q8_batch_f16_tensor( + s.q_half.raw(), + s.attention_low.raw(), + map, + size, + w.attn_output_a.offset, + w.attn_output_b.offset, + group_dim, + shape.lora_o, + shape.out_groups as u32, + shape.embd, + s.heads.raw(), + rows, + ) + } != 0; + if half_output { + call( + unsafe { + ds4_gpu_hc_expand_split_half_tensor( + s.after_attention_hc.raw(), + s.q_half.raw(), + s.current_hc.raw(), + s.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "batch attention HC expansion", + )?; + } else { + call( + unsafe { + ds4_gpu_attention_output_q8_batch_tensor( + s.attention_out.raw(), + s.attention_low.raw(), + s.attention_group_tmp.raw(), + s.attention_low_tmp.raw(), + map, + size, + w.attn_output_a.offset, + w.attn_output_b.offset, + group_dim, + shape.lora_o, + shape.out_groups as u32, + shape.embd, + s.heads.raw(), + rows, + ) + }, + "batch attention output", + )?; + call( + unsafe { + ds4_gpu_hc_expand_split_tensor( + s.after_attention_hc.raw(), + s.attention_out.raw(), + s.current_hc.raw(), + s.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "batch attention HC expansion", + )?; + } + + call( + unsafe { + ds4_gpu_hc_rms_scale_project_f16_tensor( + s.hc_mix.raw(), + s.flat_hc.raw(), + map, + size, + w.hc_ffn_fn.offset, + hc_dim as u32, + mix_hc as u32, + s.after_attention_hc.raw(), + rows, + shape.rms_epsilon, + ) + }, + "batch FFN HC projection", + )?; + call( + unsafe { + ds4_gpu_hc_split_weighted_sum_norm_tensor( + s.current.raw(), + s.norm.raw(), + s.hc_split.raw(), + s.hc_mix.raw(), + s.after_attention_hc.raw(), + map, + size, + w.hc_ffn_scale.offset, + w.hc_ffn_base.offset, + w.ffn_norm.offset, + shape.embd as u32, + shape.hc as u32, + shape.hc_sinkhorn as u32, + shape.hc_epsilon, + shape.rms_epsilon, + ) + }, + "batch FFN HC mix", + )?; + f16_rows( + &s.router_logits, + w.router, + shape.embd, + shape.experts, + &s.norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_router_select_batch_tensor( + s.router_selected.raw(), + s.router_weights.raw(), + s.router_probs.raw(), + map, + size, + w.router_bias.map_or(0, |value| value.offset), + w.router_hash.map_or(0, |value| value.offset), + w.router_hash.map_or(0, |value| value.dims[1] as u32), + 0, + 0, + w.router_bias.is_some(), + w.router_hash.is_some(), + s.router_logits.raw(), + s.tokens.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.expert_weight_scale, + rows, + ) + }, + "batch expert routing", + )?; + let gate_row = w.expert_gate.bytes / (w.expert_gate.dims[1] * w.expert_gate.dims[2]); + let down_row = w.expert_down.bytes / (w.expert_down.dims[1] * w.expert_down.dims[2]); + let mut mid_f16 = false; + call( + unsafe { + ds4_gpu_routed_moe_batch_tensor( + s.routed_out.raw(), + s.routed_gate.raw(), + s.routed_up.raw(), + s.routed_mid.raw(), + s.routed_experts.raw(), + map, + size, + w.expert_gate.offset, + w.expert_up.offset, + w.expert_down.offset, + w.expert_gate.kind, + w.expert_down.kind, + w.expert_gate.dims[1] * gate_row, + gate_row, + w.expert_down.dims[1] * down_row, + down_row, + w.expert_gate.dims[0] as u32, + w.expert_down.dims[0] as u32, + w.expert_down.dims[1] as u32, + s.router_selected.raw(), + s.router_weights.raw(), + shape.experts as u32, + shape.experts_used as u32, + shape.swiglu_clamp, + s.norm.raw(), + layer, + rows, + &mut mid_f16, + false, + ) + }, + "batch routed experts", + )?; + q8_rows( + &s.shared_gate, + w.shared_gate, + shape.embd, + shape.ff_expert, + &s.norm, + rows, + map, + size, + )?; + q8_rows( + &s.shared_up, + w.shared_up, + shape.embd, + shape.ff_expert, + &s.norm, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_swiglu_tensor( + s.shared_mid.raw(), + s.shared_gate.raw(), + s.shared_up.raw(), + rows * shape.ff_expert as u32, + shape.swiglu_clamp, + 1.0, + ) + }, + "batch shared expert activation", + )?; + q8_rows( + &s.shared_out, + w.shared_down, + shape.ff_expert, + shape.embd, + &s.shared_mid, + rows, + map, + size, + )?; + call( + unsafe { + ds4_gpu_hc_expand_add_split_tensor( + s.next_hc.raw(), + s.routed_out.raw(), + s.shared_out.raw(), + s.after_attention_hc.raw(), + s.hc_split.raw(), + shape.embd as u32, + shape.hc as u32, + ) + }, + "batch FFN HC expansion", + ) +} + // Mirrors one fixed DS4 tape invocation; grouping these scalar dimensions would // only hide the Metal kernel contract behind another one-use type. #[allow(clippy::too_many_arguments)] @@ -1595,12 +3540,7 @@ fn encode_layer( }, "KV cache write", )?; - let n_raw = (pos + 1).min(raw_cap); - let raw_start = if pos + 1 > raw_cap { - (raw_row + 1) % raw_cap - } else { - 0 - }; + let (n_raw, raw_start) = raw_decode_span(pos, raw_cap, shape.sliding_window as u32); if let (Some(weights), Some(compression)) = (w.attn_compressor.as_ref(), state.compression.as_mut()) { @@ -2210,6 +4150,29 @@ fn compression_ratio(layer: u32) -> u32 { } } +fn raw_decode_span(pos: u32, raw_cap: u32, window: u32) -> (u32, u32) { + let n_raw = (pos + 1) + .min(raw_cap) + .min(if window == 0 { raw_cap } else { window }); + (n_raw, (pos + 1 - n_raw) % raw_cap) +} + +fn raw_batch_span(pos: u32, rows: u32, raw_cap: u32, window: u32) -> (u32, u32) { + if rows == 0 || raw_cap == 0 { + return (0, 0); + } + let last = pos + rows - 1; + let needed = rows + .saturating_add(if rows == 1 { + window.saturating_sub(1) + } else { + window + }) + .min(last + 1) + .min(raw_cap); + (needed, (last + 1 - needed) % raw_cap) +} + fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 { let coefficient = if ratio == 4 { 2 } else { 1 }; coefficient * head_dim * coefficient * u64::from(ratio) * 4 @@ -2221,12 +4184,14 @@ fn write_buffer( mut offset: u64, mut bytes: u64, chunk: &mut [u8], + progress: &mut impl FnMut(u64), ) -> Result<(), String> { while bytes != 0 { let length = bytes.min(chunk.len() as u64) as usize; buffer.read(offset, &mut chunk[..length])?; file.write_all(&chunk[..length]) .map_err(|error| error.to_string())?; + progress(length as u64); offset += length as u64; bytes -= length as u64; } @@ -2239,12 +4204,14 @@ fn read_buffer( mut offset: u64, mut bytes: u64, chunk: &mut [u8], + progress: &mut impl FnMut(u64), ) -> Result<(), String> { while bytes != 0 { let length = bytes.min(chunk.len() as u64) as usize; file.read_exact(&mut chunk[..length]) .map_err(|error| error.to_string())?; buffer.write(offset, &chunk[..length])?; + progress(length as u64); offset += length as u64; bytes -= length as u64; } @@ -2397,6 +4364,52 @@ fn f16( ) } +#[allow(clippy::too_many_arguments)] +fn f16_rows( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + rows: u32, + map: *const c_void, + size: u64, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_matmul_f16_tensor( + out.raw(), + map, + size, + weight.offset, + input, + output, + x.raw(), + u64::from(rows), + ) + }, + "batch F16 projection", + ) +} + +#[allow(clippy::too_many_arguments)] +fn matmul_rows( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + rows: u32, + map: *const c_void, + size: u64, +) -> Result<(), String> { + match weight.kind { + F16 => f16_rows(out, weight, input, output, x, rows, map, size), + Q8_0 => q8_rows(out, weight, input, output, x, rows, map, size), + kind => Err(format!("unsupported Metal batch projection type {kind}")), + } +} + fn matmul( out: &Buffer, weight: Weight, @@ -2439,6 +4452,34 @@ fn q8( ) } +#[allow(clippy::too_many_arguments)] +fn q8_rows( + out: &Buffer, + weight: Weight, + input: u64, + output: u64, + x: &Buffer, + rows: u32, + map: *const c_void, + size: u64, +) -> Result<(), String> { + call( + unsafe { + ds4_gpu_matmul_q8_0_tensor( + out.raw(), + map, + size, + weight.offset, + input, + output, + x.raw(), + u64::from(rows), + ) + }, + "batch Q8 projection", + ) +} + fn call(result: i32, operation: &str) -> Result<(), String> { if result == 0 { Err(format!("Metal failed while {operation}")) @@ -2450,3 +4491,24 @@ fn call(result: i32, operation: &str) -> Result<(), String> { fn check(result: i32, operation: &str) -> Result<(), String> { call(result, operation) } + +#[cfg(test)] +mod tests { + use super::{raw_batch_span, raw_decode_span}; + + #[test] + fn decode_raw_span_tracks_the_logical_sliding_window() { + assert_eq!(raw_decode_span(0, 4_352, 128), (1, 0)); + assert_eq!(raw_decode_span(127, 4_352, 128), (128, 0)); + assert_eq!(raw_decode_span(128, 4_352, 128), (128, 1)); + assert_eq!(raw_decode_span(362, 4_352, 128), (128, 235)); + assert_eq!(raw_decode_span(4_500, 4_352, 128), (128, 4_373 % 4_352)); + } + + #[test] + fn batch_raw_span_includes_the_chunk_and_previous_window() { + assert_eq!(raw_batch_span(0, 4_096, 4_352, 128), (4_096, 0)); + assert_eq!(raw_batch_span(4_096, 1_925, 4_352, 128), (2_053, 3_968)); + assert_eq!(raw_batch_span(4_500, 1, 4_352, 128), (128, 21)); + } +} diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index 27ae1fa..2fc7119 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -21,6 +21,7 @@ pub(super) struct Tokenizer { sop: i32, think_start: i32, think_end: i32, + rendered_specials: Vec<(&'static [u8], i32)>, } impl Tokenizer { @@ -85,6 +86,32 @@ impl Tokenizer { { return Err("tokenizer does not provide the required DS4 chat markers".into()); } + let rendered_specials = [ + ("<｜begin▁of▁sentence｜>".as_bytes(), bos), + ("<｜end▁of▁sentence｜>".as_bytes(), eos), + (b"[gMASK]".as_slice(), bos), + (b"".as_slice(), sop), + (b"<|system|>".as_slice(), system), + ("<｜User｜>".as_bytes(), user), + ("<｜Assistant｜>".as_bytes(), assistant), + (b"<|user|>".as_slice(), user), + (b"<|assistant|>".as_slice(), assistant), + (b"<|observation|>".as_slice(), observation), + (b"".as_slice(), think_start), + (b"".as_slice(), think_end), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + (b"".as_slice(), lookup(b"")), + ("｜DSML｜".as_bytes(), lookup("｜DSML｜".as_bytes())), + ] + .into_iter() + .filter(|(_, token)| *token >= 0) + .collect(); Ok(Self { family, @@ -100,6 +127,7 @@ impl Tokenizer { sop, think_start, think_end, + rendered_specials, }) } @@ -126,6 +154,37 @@ impl Tokenizer { output } + pub(super) fn tokenize_rendered(&self, text: &str) -> Vec { + let bytes = text.as_bytes(); + let mut output = Vec::new(); + let mut span = 0; + let mut position = 0; + while position < bytes.len() { + let special = self + .rendered_specials + .iter() + .find(|(marker, _)| bytes[position..].starts_with(marker)); + if let Some((marker, token)) = special { + self.tokenize_plain(&text[span..position], &mut output); + output.push(*token); + position += marker.len(); + span = position; + } else { + position += 1; + } + } + self.tokenize_plain(&text[span..], &mut output); + output + } + + fn tokenize_plain(&self, text: &str, output: &mut Vec) { + if self.family == ModelFamily::Glm { + self.tokenize_glm(text, output); + } else { + self.tokenize_joyai(text, output); + } + } + pub(super) fn encode_chat( &self, system_prompt: &str, @@ -136,6 +195,7 @@ impl Tokenizer { system_prompt, &[ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content: prompt.to_owned(), @@ -172,7 +232,7 @@ impl Tokenizer { if self.family == ModelFamily::Glm { output.push(self.system); } - output.extend(self.tokenize(system_prompt)); + output.extend(self.tokenize_rendered(system_prompt)); } for message in messages { output.push(if message.user { @@ -183,7 +243,7 @@ impl Tokenizer { if !message.user { if let Some(reasoning) = &message.reasoning { output.push(self.think_start); - output.extend(self.tokenize(reasoning)); + output.extend(self.tokenize_rendered(reasoning)); if message.reasoning_complete { output.push(self.think_end); } @@ -193,7 +253,7 @@ impl Tokenizer { output.push(self.think_end); } } - output.extend(self.tokenize(&message.content)); + output.extend(self.tokenize_rendered(&message.content)); if !message.user && self.family == ModelFamily::DeepSeek { output.push(self.eos); } @@ -209,13 +269,18 @@ impl Tokenizer { output } - pub(super) fn encode_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec { + pub(super) fn encode_continuation( + &self, + prompt: &str, + reasoning: ReasoningMode, + skip_previous_eos: bool, + ) -> Vec { let mut output = Vec::new(); - if self.family == ModelFamily::DeepSeek { + if self.family == ModelFamily::DeepSeek && !skip_previous_eos { output.push(self.eos); } output.push(self.user); - output.extend(self.tokenize(prompt)); + output.extend(self.tokenize_rendered(prompt)); output.push(self.assistant); if reasoning != ReasoningMode::Direct { output.push(self.think_start); diff --git a/src/main.rs b/src/main.rs index d9c2a02..6b67606 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ mod app; mod database; mod engine; +mod metrics; mod model; #[cfg(target_os = "macos")] mod native_edit; diff --git a/src/metrics.rs b/src/metrics.rs new file mode 100644 index 0000000..eabc658 --- /dev/null +++ b/src/metrics.rs @@ -0,0 +1,727 @@ +use crate::model::ModelChoice; +use std::fs; +use std::path::Path; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU16, AtomicU32, AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[repr(u8)] +pub(crate) enum RuntimePhase { + #[default] + Unloaded, + Loading, + Prefilling, + Generating, + Ready, + Failed, +} + +impl RuntimePhase { + pub(crate) fn label(self) -> &'static str { + match self { + Self::Unloaded => "Unloaded", + Self::Loading => "Loading", + Self::Prefilling => "Prefilling", + Self::Generating => "Generating", + Self::Ready => "Ready", + Self::Failed => "Failed", + } + } +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[repr(u8)] +pub(crate) enum WorkSource { + #[default] + None, + LocalChat, + Http, +} + +impl WorkSource { + pub(crate) fn label(self) -> &'static str { + match self { + Self::None => "None", + Self::LocalChat => "Local chat", + Self::Http => "HTTP endpoint", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum KvLookup { + MemoryHit, + DiskHit, + Miss, + Invalid, +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct MetricsSnapshot { + pub(crate) uptime_seconds: u64, + pub(crate) phase: RuntimePhase, + pub(crate) source: WorkSource, + pub(crate) model: &'static str, + pub(crate) queue_depth: u32, + pub(crate) runtime_requests: u64, + pub(crate) local_requests: u64, + pub(crate) endpoint_generations: u64, + pub(crate) completed_requests: u64, + pub(crate) failed_requests: u64, + pub(crate) model_loads: u64, + pub(crate) model_unloads: u64, + pub(crate) model_load_ms: u64, + pub(crate) model_bytes: u64, + pub(crate) tensor_count: u64, + pub(crate) vocabulary_size: u64, + pub(crate) context_used: u32, + pub(crate) context_limit: u32, + pub(crate) decode_tokens_per_second: f32, + pub(crate) prefill_tokens_per_second: f32, + pub(crate) last_runtime_ms: u64, + pub(crate) average_runtime_ms: u64, + pub(crate) last_prompt_tokens: u64, + pub(crate) last_cached_tokens: u64, + pub(crate) last_completion_tokens: u64, + pub(crate) prompt_tokens: u64, + pub(crate) cached_tokens: u64, + pub(crate) completion_tokens: u64, + pub(crate) kv_lookups: u64, + pub(crate) kv_hits: u64, + pub(crate) kv_memory_hits: u64, + pub(crate) kv_disk_hits: u64, + pub(crate) kv_misses: u64, + pub(crate) kv_invalid: u64, + pub(crate) kv_prefix_hits: u64, + pub(crate) checkpoint_writes: u64, + pub(crate) kv_read_active: bool, + pub(crate) kv_write_active: bool, + pub(crate) kv_read_operations: u64, + pub(crate) kv_write_operations: u64, + pub(crate) kv_read_errors: u64, + pub(crate) kv_write_errors: u64, + pub(crate) kv_read_bytes: u64, + pub(crate) kv_write_bytes: u64, + pub(crate) last_kv_read_ms: u64, + pub(crate) last_kv_write_ms: u64, + pub(crate) kv_files: u64, + pub(crate) kv_bytes: u64, + pub(crate) local_kv_files: u64, + pub(crate) local_kv_bytes: u64, + pub(crate) http_kv_files: u64, + pub(crate) http_kv_bytes: u64, + pub(crate) server_listening: bool, + pub(crate) server_port: u16, + pub(crate) http_active: u32, + pub(crate) http_requests: u64, + pub(crate) http_completed: u64, + pub(crate) http_errors: u64, + pub(crate) http_chat_requests: u64, + pub(crate) http_model_requests: u64, + pub(crate) http_streaming_requests: u64, + pub(crate) http_bytes_received: u64, + pub(crate) last_http_ms: u64, + pub(crate) average_http_ms: u64, +} + +pub(crate) struct Metrics { + started: Instant, + phase: AtomicU8, + source: AtomicU8, + model: AtomicU8, + queue_depth: AtomicU32, + runtime_requests: AtomicU64, + local_requests: AtomicU64, + endpoint_generations: AtomicU64, + completed_requests: AtomicU64, + failed_requests: AtomicU64, + model_loads: AtomicU64, + model_unloads: AtomicU64, + model_load_ms: AtomicU64, + model_bytes: AtomicU64, + tensor_count: AtomicU64, + vocabulary_size: AtomicU64, + context_used: AtomicU32, + context_limit: AtomicU32, + decode_tps: AtomicU32, + prefill_tps: AtomicU32, + last_runtime_ms: AtomicU64, + total_runtime_ms: AtomicU64, + last_prompt_tokens: AtomicU64, + last_cached_tokens: AtomicU64, + last_completion_tokens: AtomicU64, + prompt_tokens: AtomicU64, + cached_tokens: AtomicU64, + completion_tokens: AtomicU64, + kv_lookups: AtomicU64, + kv_hits: AtomicU64, + kv_memory_hits: AtomicU64, + kv_disk_hits: AtomicU64, + kv_misses: AtomicU64, + kv_invalid: AtomicU64, + kv_prefix_hits: AtomicU64, + checkpoint_writes: AtomicU64, + kv_read_active: AtomicBool, + kv_write_active: AtomicBool, + kv_read_operations: AtomicU64, + kv_write_operations: AtomicU64, + kv_read_errors: AtomicU64, + kv_write_errors: AtomicU64, + kv_read_bytes: AtomicU64, + kv_write_bytes: AtomicU64, + last_kv_read_ms: AtomicU64, + last_kv_write_ms: AtomicU64, + kv_files: AtomicU64, + kv_bytes: AtomicU64, + local_kv_files: AtomicU64, + local_kv_bytes: AtomicU64, + http_kv_files: AtomicU64, + http_kv_bytes: AtomicU64, + server_listening: AtomicBool, + server_port: AtomicU16, + http_active: AtomicU32, + http_requests: AtomicU64, + http_completed: AtomicU64, + http_errors: AtomicU64, + http_chat_requests: AtomicU64, + http_model_requests: AtomicU64, + http_streaming_requests: AtomicU64, + http_bytes_received: AtomicU64, + last_http_ms: AtomicU64, + total_http_ms: AtomicU64, +} + +impl Metrics { + pub(crate) fn new(cache_root: &Path) -> Self { + let cache = cache_usage(cache_root); + Self { + started: Instant::now(), + phase: AtomicU8::new(RuntimePhase::Unloaded as u8), + source: AtomicU8::new(WorkSource::None as u8), + model: AtomicU8::new(0), + queue_depth: AtomicU32::new(0), + runtime_requests: AtomicU64::new(0), + local_requests: AtomicU64::new(0), + endpoint_generations: AtomicU64::new(0), + completed_requests: AtomicU64::new(0), + failed_requests: AtomicU64::new(0), + model_loads: AtomicU64::new(0), + model_unloads: AtomicU64::new(0), + model_load_ms: AtomicU64::new(0), + model_bytes: AtomicU64::new(0), + tensor_count: AtomicU64::new(0), + vocabulary_size: AtomicU64::new(0), + context_used: AtomicU32::new(0), + context_limit: AtomicU32::new(0), + decode_tps: AtomicU32::new(0), + prefill_tps: AtomicU32::new(0), + last_runtime_ms: AtomicU64::new(0), + total_runtime_ms: AtomicU64::new(0), + last_prompt_tokens: AtomicU64::new(0), + last_cached_tokens: AtomicU64::new(0), + last_completion_tokens: AtomicU64::new(0), + prompt_tokens: AtomicU64::new(0), + cached_tokens: AtomicU64::new(0), + completion_tokens: AtomicU64::new(0), + kv_lookups: AtomicU64::new(0), + kv_hits: AtomicU64::new(0), + kv_memory_hits: AtomicU64::new(0), + kv_disk_hits: AtomicU64::new(0), + kv_misses: AtomicU64::new(0), + kv_invalid: AtomicU64::new(0), + kv_prefix_hits: AtomicU64::new(0), + checkpoint_writes: AtomicU64::new(0), + kv_read_active: AtomicBool::new(false), + kv_write_active: AtomicBool::new(false), + kv_read_operations: AtomicU64::new(0), + kv_write_operations: AtomicU64::new(0), + kv_read_errors: AtomicU64::new(0), + kv_write_errors: AtomicU64::new(0), + kv_read_bytes: AtomicU64::new(0), + kv_write_bytes: AtomicU64::new(0), + last_kv_read_ms: AtomicU64::new(0), + last_kv_write_ms: AtomicU64::new(0), + kv_files: AtomicU64::new(cache.files), + kv_bytes: AtomicU64::new(cache.bytes), + local_kv_files: AtomicU64::new(cache.local_files), + local_kv_bytes: AtomicU64::new(cache.local_bytes), + http_kv_files: AtomicU64::new(cache.http_files), + http_kv_bytes: AtomicU64::new(cache.http_bytes), + server_listening: AtomicBool::new(false), + server_port: AtomicU16::new(0), + http_active: AtomicU32::new(0), + http_requests: AtomicU64::new(0), + http_completed: AtomicU64::new(0), + http_errors: AtomicU64::new(0), + http_chat_requests: AtomicU64::new(0), + http_model_requests: AtomicU64::new(0), + http_streaming_requests: AtomicU64::new(0), + http_bytes_received: AtomicU64::new(0), + last_http_ms: AtomicU64::new(0), + total_http_ms: AtomicU64::new(0), + } + } + + pub(crate) fn request_queued(&self, source: WorkSource) { + self.runtime_requests.fetch_add(1, Ordering::Relaxed); + match source { + WorkSource::LocalChat => self.local_requests.fetch_add(1, Ordering::Relaxed), + WorkSource::Http => self.endpoint_generations.fetch_add(1, Ordering::Relaxed), + WorkSource::None => 0, + }; + self.queue_depth.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn request_rejected(&self) { + self.queue_depth.fetch_sub(1, Ordering::Relaxed); + self.failed_requests.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn request_started(&self, source: WorkSource) { + self.queue_depth.fetch_sub(1, Ordering::Relaxed); + self.source.store(source as u8, Ordering::Relaxed); + self.decode_tps.store(0, Ordering::Relaxed); + self.prefill_tps.store(0, Ordering::Relaxed); + } + + pub(crate) fn loading(&self) { + self.phase + .store(RuntimePhase::Loading as u8, Ordering::Relaxed); + } + + pub(crate) fn loaded( + &self, + model: ModelChoice, + elapsed: Duration, + mapped_bytes: u64, + tensors: usize, + vocabulary_size: usize, + ) { + self.model.store(model_code(model), Ordering::Relaxed); + self.model_loads.fetch_add(1, Ordering::Relaxed); + self.model_load_ms + .store(milliseconds(elapsed), Ordering::Relaxed); + self.model_bytes.store(mapped_bytes, Ordering::Relaxed); + self.tensor_count.store(tensors as u64, Ordering::Relaxed); + self.vocabulary_size + .store(vocabulary_size as u64, Ordering::Relaxed); + } + + pub(crate) fn prefill_progress(&self, used: u32, limit: u32, speed: f32) { + self.phase + .store(RuntimePhase::Prefilling as u8, Ordering::Relaxed); + self.context_used.store(used, Ordering::Relaxed); + self.context_limit.store(limit, Ordering::Relaxed); + store_f32(&self.prefill_tps, speed); + } + + pub(crate) fn generation_progress(&self, used: u32, limit: u32, speed: f32) { + self.phase + .store(RuntimePhase::Generating as u8, Ordering::Relaxed); + self.context_used.store(used, Ordering::Relaxed); + self.context_limit.store(limit, Ordering::Relaxed); + store_f32(&self.decode_tps, speed); + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn request_finished( + &self, + source: WorkSource, + elapsed: Duration, + prompt_tokens: u32, + cached_tokens: u32, + completion_tokens: u32, + previous_checkpoint_bytes: Option, + checkpoint_bytes: u64, + ) { + let elapsed = milliseconds(elapsed); + self.completed_requests.fetch_add(1, Ordering::Relaxed); + self.last_runtime_ms.store(elapsed, Ordering::Relaxed); + self.total_runtime_ms.fetch_add(elapsed, Ordering::Relaxed); + self.last_prompt_tokens + .store(u64::from(prompt_tokens), Ordering::Relaxed); + self.last_cached_tokens + .store(u64::from(cached_tokens), Ordering::Relaxed); + self.last_completion_tokens + .store(u64::from(completion_tokens), Ordering::Relaxed); + self.prompt_tokens + .fetch_add(u64::from(prompt_tokens), Ordering::Relaxed); + self.cached_tokens + .fetch_add(u64::from(cached_tokens), Ordering::Relaxed); + self.completion_tokens + .fetch_add(u64::from(completion_tokens), Ordering::Relaxed); + self.record_checkpoint(source, previous_checkpoint_bytes, checkpoint_bytes); + self.phase + .store(RuntimePhase::Ready as u8, Ordering::Relaxed); + self.source.store(WorkSource::None as u8, Ordering::Relaxed); + } + + pub(crate) fn request_failed(&self, elapsed: Duration) { + let elapsed = milliseconds(elapsed); + self.failed_requests.fetch_add(1, Ordering::Relaxed); + self.last_runtime_ms.store(elapsed, Ordering::Relaxed); + self.total_runtime_ms.fetch_add(elapsed, Ordering::Relaxed); + self.phase + .store(RuntimePhase::Failed as u8, Ordering::Relaxed); + self.source.store(WorkSource::None as u8, Ordering::Relaxed); + } + + pub(crate) fn unloaded(&self) { + self.phase + .store(RuntimePhase::Unloaded as u8, Ordering::Relaxed); + self.model.store(0, Ordering::Relaxed); + self.model_bytes.store(0, Ordering::Relaxed); + self.tensor_count.store(0, Ordering::Relaxed); + self.vocabulary_size.store(0, Ordering::Relaxed); + self.context_used.store(0, Ordering::Relaxed); + self.decode_tps.store(0, Ordering::Relaxed); + self.prefill_tps.store(0, Ordering::Relaxed); + self.model_unloads.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn server_listening(&self, port: u16) { + self.server_port.store(port, Ordering::Relaxed); + self.server_listening.store(true, Ordering::Relaxed); + } + + pub(crate) fn kv_lookup(&self, result: KvLookup) { + self.kv_lookups.fetch_add(1, Ordering::Relaxed); + match result { + KvLookup::MemoryHit => { + self.kv_hits.fetch_add(1, Ordering::Relaxed); + self.kv_memory_hits.fetch_add(1, Ordering::Relaxed); + } + KvLookup::DiskHit => { + self.kv_hits.fetch_add(1, Ordering::Relaxed); + self.kv_disk_hits.fetch_add(1, Ordering::Relaxed); + } + KvLookup::Miss => { + self.kv_misses.fetch_add(1, Ordering::Relaxed); + } + KvLookup::Invalid => { + self.kv_misses.fetch_add(1, Ordering::Relaxed); + self.kv_invalid.fetch_add(1, Ordering::Relaxed); + } + } + } + + pub(crate) fn kv_prefix_reused(&self, tokens: usize) { + if tokens > 0 { + self.kv_prefix_hits.fetch_add(1, Ordering::Relaxed); + } + } + + pub(crate) fn kv_read_started(&self) { + self.kv_read_active.store(true, Ordering::Relaxed); + self.kv_read_operations.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn kv_read_bytes(&self, bytes: u64) { + self.kv_read_bytes.fetch_add(bytes, Ordering::Relaxed); + } + + pub(crate) fn kv_read_finished(&self, elapsed: Duration, failed: bool) { + self.kv_read_active.store(false, Ordering::Relaxed); + self.last_kv_read_ms + .store(milliseconds(elapsed), Ordering::Relaxed); + if failed { + self.kv_read_errors.fetch_add(1, Ordering::Relaxed); + } + } + + pub(crate) fn kv_write_started(&self) { + self.kv_write_active.store(true, Ordering::Relaxed); + self.kv_write_operations.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn kv_write_bytes(&self, bytes: u64) { + self.kv_write_bytes.fetch_add(bytes, Ordering::Relaxed); + } + + pub(crate) fn kv_write_finished(&self, elapsed: Duration, failed: bool) { + self.kv_write_active.store(false, Ordering::Relaxed); + self.last_kv_write_ms + .store(milliseconds(elapsed), Ordering::Relaxed); + if failed { + self.kv_write_errors.fetch_add(1, Ordering::Relaxed); + } + } + + pub(crate) fn server_stopped(&self, port: u16) { + if self.server_port.load(Ordering::Relaxed) == port { + self.server_listening.store(false, Ordering::Relaxed); + self.http_active.store(0, Ordering::Relaxed); + } + } + + pub(crate) fn http_started(&self, path: &str, bytes: usize) { + self.http_active.fetch_add(1, Ordering::Relaxed); + self.http_requests.fetch_add(1, Ordering::Relaxed); + self.http_bytes_received + .fetch_add(bytes as u64, Ordering::Relaxed); + match path { + "/v1/chat/completions" => { + self.http_chat_requests.fetch_add(1, Ordering::Relaxed); + } + path if path == "/v1/models" || path.starts_with("/v1/models/") => { + self.http_model_requests.fetch_add(1, Ordering::Relaxed); + } + _ => {} + } + } + + pub(crate) fn http_streaming(&self) { + self.http_streaming_requests.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn http_finished(&self, elapsed: Duration, failed: bool) { + self.http_active.fetch_sub(1, Ordering::Relaxed); + self.http_completed.fetch_add(1, Ordering::Relaxed); + if failed { + self.http_errors.fetch_add(1, Ordering::Relaxed); + } + let elapsed = milliseconds(elapsed); + self.last_http_ms.store(elapsed, Ordering::Relaxed); + self.total_http_ms.fetch_add(elapsed, Ordering::Relaxed); + } + + pub(crate) fn snapshot(&self) -> MetricsSnapshot { + let completed = self.completed_requests.load(Ordering::Relaxed); + let failed = self.failed_requests.load(Ordering::Relaxed); + let http_completed = self.http_completed.load(Ordering::Relaxed); + MetricsSnapshot { + uptime_seconds: self.started.elapsed().as_secs(), + phase: phase(self.phase.load(Ordering::Relaxed)), + source: source(self.source.load(Ordering::Relaxed)), + model: model_name(self.model.load(Ordering::Relaxed)), + queue_depth: self.queue_depth.load(Ordering::Relaxed), + runtime_requests: self.runtime_requests.load(Ordering::Relaxed), + local_requests: self.local_requests.load(Ordering::Relaxed), + endpoint_generations: self.endpoint_generations.load(Ordering::Relaxed), + completed_requests: completed, + failed_requests: failed, + model_loads: self.model_loads.load(Ordering::Relaxed), + model_unloads: self.model_unloads.load(Ordering::Relaxed), + model_load_ms: self.model_load_ms.load(Ordering::Relaxed), + model_bytes: self.model_bytes.load(Ordering::Relaxed), + tensor_count: self.tensor_count.load(Ordering::Relaxed), + vocabulary_size: self.vocabulary_size.load(Ordering::Relaxed), + context_used: self.context_used.load(Ordering::Relaxed), + context_limit: self.context_limit.load(Ordering::Relaxed), + decode_tokens_per_second: load_f32(&self.decode_tps), + prefill_tokens_per_second: load_f32(&self.prefill_tps), + last_runtime_ms: self.last_runtime_ms.load(Ordering::Relaxed), + average_runtime_ms: average( + self.total_runtime_ms.load(Ordering::Relaxed), + completed + failed, + ), + last_prompt_tokens: self.last_prompt_tokens.load(Ordering::Relaxed), + last_cached_tokens: self.last_cached_tokens.load(Ordering::Relaxed), + last_completion_tokens: self.last_completion_tokens.load(Ordering::Relaxed), + prompt_tokens: self.prompt_tokens.load(Ordering::Relaxed), + cached_tokens: self.cached_tokens.load(Ordering::Relaxed), + completion_tokens: self.completion_tokens.load(Ordering::Relaxed), + kv_lookups: self.kv_lookups.load(Ordering::Relaxed), + kv_hits: self.kv_hits.load(Ordering::Relaxed), + kv_memory_hits: self.kv_memory_hits.load(Ordering::Relaxed), + kv_disk_hits: self.kv_disk_hits.load(Ordering::Relaxed), + kv_misses: self.kv_misses.load(Ordering::Relaxed), + kv_invalid: self.kv_invalid.load(Ordering::Relaxed), + kv_prefix_hits: self.kv_prefix_hits.load(Ordering::Relaxed), + checkpoint_writes: self.checkpoint_writes.load(Ordering::Relaxed), + kv_read_active: self.kv_read_active.load(Ordering::Relaxed), + kv_write_active: self.kv_write_active.load(Ordering::Relaxed), + kv_read_operations: self.kv_read_operations.load(Ordering::Relaxed), + kv_write_operations: self.kv_write_operations.load(Ordering::Relaxed), + kv_read_errors: self.kv_read_errors.load(Ordering::Relaxed), + kv_write_errors: self.kv_write_errors.load(Ordering::Relaxed), + kv_read_bytes: self.kv_read_bytes.load(Ordering::Relaxed), + kv_write_bytes: self.kv_write_bytes.load(Ordering::Relaxed), + last_kv_read_ms: self.last_kv_read_ms.load(Ordering::Relaxed), + last_kv_write_ms: self.last_kv_write_ms.load(Ordering::Relaxed), + kv_files: self.kv_files.load(Ordering::Relaxed), + kv_bytes: self.kv_bytes.load(Ordering::Relaxed), + local_kv_files: self.local_kv_files.load(Ordering::Relaxed), + local_kv_bytes: self.local_kv_bytes.load(Ordering::Relaxed), + http_kv_files: self.http_kv_files.load(Ordering::Relaxed), + http_kv_bytes: self.http_kv_bytes.load(Ordering::Relaxed), + server_listening: self.server_listening.load(Ordering::Relaxed), + server_port: self.server_port.load(Ordering::Relaxed), + http_active: self.http_active.load(Ordering::Relaxed), + http_requests: self.http_requests.load(Ordering::Relaxed), + http_completed, + http_errors: self.http_errors.load(Ordering::Relaxed), + http_chat_requests: self.http_chat_requests.load(Ordering::Relaxed), + http_model_requests: self.http_model_requests.load(Ordering::Relaxed), + http_streaming_requests: self.http_streaming_requests.load(Ordering::Relaxed), + http_bytes_received: self.http_bytes_received.load(Ordering::Relaxed), + last_http_ms: self.last_http_ms.load(Ordering::Relaxed), + average_http_ms: average(self.total_http_ms.load(Ordering::Relaxed), http_completed), + } + } + + fn record_checkpoint(&self, source: WorkSource, previous_bytes: Option, bytes: u64) { + if bytes == 0 { + return; + } + self.checkpoint_writes.fetch_add(1, Ordering::Relaxed); + let previous = previous_bytes.unwrap_or(0); + update_bytes(&self.kv_bytes, previous, bytes); + let (files, total) = match source { + WorkSource::LocalChat => (&self.local_kv_files, &self.local_kv_bytes), + WorkSource::Http => (&self.http_kv_files, &self.http_kv_bytes), + WorkSource::None => return, + }; + update_bytes(total, previous, bytes); + if previous_bytes.is_none() { + self.kv_files.fetch_add(1, Ordering::Relaxed); + files.fetch_add(1, Ordering::Relaxed); + } + } +} + +fn store_f32(target: &AtomicU32, value: f32) { + target.store(value.max(0.0).to_bits(), Ordering::Relaxed); +} + +fn load_f32(source: &AtomicU32) -> f32 { + f32::from_bits(source.load(Ordering::Relaxed)) +} + +fn milliseconds(duration: Duration) -> u64 { + duration.as_millis().min(u128::from(u64::MAX)) as u64 +} + +fn average(total: u64, count: u64) -> u64 { + total.checked_div(count).unwrap_or(0) +} + +fn update_bytes(target: &AtomicU64, previous: u64, current: u64) { + let total = target.load(Ordering::Relaxed); + target.store(total.saturating_sub(previous) + current, Ordering::Relaxed); +} + +fn phase(value: u8) -> RuntimePhase { + match value { + 1 => RuntimePhase::Loading, + 2 => RuntimePhase::Prefilling, + 3 => RuntimePhase::Generating, + 4 => RuntimePhase::Ready, + 5 => RuntimePhase::Failed, + _ => RuntimePhase::Unloaded, + } +} + +fn source(value: u8) -> WorkSource { + match value { + 1 => WorkSource::LocalChat, + 2 => WorkSource::Http, + _ => WorkSource::None, + } +} + +fn model_code(model: ModelChoice) -> u8 { + match model { + ModelChoice::DeepSeekV4Flash => 1, + ModelChoice::DeepSeekV4Pro => 2, + ModelChoice::Glm52 => 3, + } +} + +fn model_name(value: u8) -> &'static str { + match value { + 1 => "DeepSeek V4 Flash", + 2 => "DeepSeek V4 Pro", + 3 => "GLM 5.2", + _ => "No model loaded", + } +} + +#[derive(Default)] +struct CacheUsage { + files: u64, + bytes: u64, + local_files: u64, + local_bytes: u64, + http_files: u64, + http_bytes: u64, +} + +fn cache_usage(root: &Path) -> CacheUsage { + let mut usage = CacheUsage::default(); + for (path, http) in [(root.to_path_buf(), false), (root.join("http"), true)] { + let Ok(entries) = fs::read_dir(path) else { + continue; + }; + for entry in entries.flatten() { + let Ok(metadata) = entry.metadata() else { + continue; + }; + if !metadata.is_file() { + continue; + } + usage.files += 1; + usage.bytes += metadata.len(); + if http { + usage.http_files += 1; + usage.http_bytes += metadata.len(); + } else { + usage.local_files += 1; + usage.local_bytes += metadata.len(); + } + } + } + usage +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn snapshots_track_runtime_and_server_counters() { + let metrics = Metrics::new(Path::new("/path/that/does/not/exist")); + metrics.request_queued(WorkSource::LocalChat); + metrics.request_started(WorkSource::LocalChat); + metrics.prefill_progress(100, 1_000, 250.0); + metrics.generation_progress(120, 1_000, 20.0); + metrics.kv_lookup(KvLookup::DiskHit); + metrics.kv_lookup(KvLookup::MemoryHit); + metrics.kv_lookup(KvLookup::Miss); + metrics.kv_lookup(KvLookup::Invalid); + metrics.kv_prefix_reused(80); + metrics.kv_prefix_reused(0); + metrics.kv_read_started(); + metrics.kv_read_bytes(2_048); + metrics.kv_read_finished(Duration::from_millis(10), false); + metrics.kv_write_started(); + metrics.kv_write_bytes(4_096); + metrics.kv_write_finished(Duration::from_millis(20), false); + metrics.request_finished( + WorkSource::LocalChat, + Duration::from_millis(250), + 100, + 80, + 20, + None, + 4_096, + ); + metrics.http_started("/v1/models", 0); + metrics.http_finished(Duration::from_millis(5), false); + + let snapshot = metrics.snapshot(); + assert_eq!(snapshot.phase, RuntimePhase::Ready); + assert_eq!(snapshot.cached_tokens, 80); + assert_eq!(snapshot.kv_lookups, 4); + assert_eq!(snapshot.kv_hits, 2); + assert_eq!(snapshot.kv_memory_hits, 1); + assert_eq!(snapshot.kv_disk_hits, 1); + assert_eq!(snapshot.kv_misses, 2); + assert_eq!(snapshot.kv_invalid, 1); + assert_eq!(snapshot.kv_prefix_hits, 1); + assert_eq!(snapshot.kv_read_bytes, 2_048); + assert_eq!(snapshot.kv_write_bytes, 4_096); + assert_eq!(snapshot.local_kv_bytes, 4_096); + assert_eq!(snapshot.http_model_requests, 1); + } +} diff --git a/src/runtime.rs b/src/runtime.rs index 9a64fb6..7a13dd3 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,4 +1,5 @@ use crate::engine::{ChatTurn, GenerationOutput, Generator}; +use crate::metrics::{Metrics, WorkSource}; use crate::settings::{EngineSettings, TurnSettings}; use std::path::PathBuf; use std::sync::Arc; @@ -10,6 +11,7 @@ use std::time::{Duration, Instant}; #[derive(Clone)] pub(crate) struct GenerationService { commands: Sender, + metrics: Arc, } pub(crate) struct ActiveGeneration { @@ -47,13 +49,14 @@ struct Command { } impl GenerationService { - pub(crate) fn spawn() -> Result { + pub(crate) fn spawn(metrics: Arc) -> Result { let (commands, receiver) = mpsc::channel::(); + let worker_metrics = Arc::clone(&metrics); thread::Builder::new() .name("model-runtime".into()) - .spawn(move || run(receiver)) + .spawn(move || run(receiver, worker_metrics)) .map_err(|error| format!("Could not start the model runtime: {error}"))?; - Ok(Self { commands }) + Ok(Self { commands, metrics }) } pub(crate) fn generate( @@ -64,9 +67,15 @@ impl GenerationService { checkpoint: CheckpointTarget, idle_timeout: Duration, ) -> Result { + let source = match &checkpoint { + CheckpointTarget::Local(_) => WorkSource::LocalChat, + CheckpointTarget::Transient(_) => WorkSource::Http, + }; let cancel = Arc::new(AtomicBool::new(false)); let (events, receiver) = mpsc::channel(); - self.commands + self.metrics.request_queued(source); + if self + .commands .send(Command { engine, turn, @@ -76,7 +85,11 @@ impl GenerationService { cancel: Arc::clone(&cancel), events, }) - .map_err(|_| "The model runtime stopped unexpectedly.".to_owned())?; + .is_err() + { + self.metrics.request_rejected(); + return Err("The model runtime stopped unexpectedly.".to_owned()); + } Ok(ActiveGeneration { events: receiver, cancel, @@ -84,34 +97,70 @@ impl GenerationService { } } -fn run(commands: Receiver) { +fn run(commands: Receiver, metrics: Arc) { let mut loaded = None::<(EngineSettings, Generator)>; let mut last_used = Instant::now(); let mut idle_timeout = Duration::from_secs(15 * 60); loop { match commands.recv_timeout(Duration::from_secs(1)) { Ok(command) => { + let request_started = Instant::now(); + let source = match &command.checkpoint { + CheckpointTarget::Local(_) => WorkSource::LocalChat, + CheckpointTarget::Transient(_) => WorkSource::Http, + }; + metrics.request_started(source); idle_timeout = command.idle_timeout; if loaded .as_ref() .is_none_or(|(settings, _)| settings != &command.engine) { + if loaded.take().is_some() { + metrics.unloaded(); + } let _ = command.events.send(GenerationEvent::Loading); - loaded = match Generator::open(&command.engine) { - Ok(generator) => Some((command.engine.clone(), generator)), + metrics.loading(); + let load_started = Instant::now(); + loaded = match Generator::open(&command.engine, Arc::clone(&metrics)) { + Ok(generator) => { + let summary = generator.summary(); + metrics.loaded( + summary.model, + load_started.elapsed(), + summary.mapped_bytes, + summary.tensor_count, + summary.vocabulary_size, + ); + Some((command.engine.clone(), generator)) + } Err(error) => { + metrics.request_failed(request_started.elapsed()); let _ = command.events.send(GenerationEvent::Finished(Err(error))); None } }; } if let Some((_, generator)) = &mut loaded { + let mut prefill_started = None::<(Instant, u32)>; let mut emit = |reasoning, content| { let _ = command .events .send(GenerationEvent::Chunk { reasoning, content }); }; let mut progress = |used, limit, tokens_per_second| { + if let Some(speed) = tokens_per_second { + metrics.generation_progress(used, limit, speed); + } else { + let (started, initial) = + prefill_started.get_or_insert_with(|| (Instant::now(), used)); + let elapsed = started.elapsed().as_secs_f32(); + let speed = if elapsed > 0.0 { + used.saturating_sub(*initial) as f32 / elapsed + } else { + 0.0 + }; + metrics.prefill_progress(used, limit, speed); + } let _ = command.events.send(GenerationEvent::Context { used, limit, @@ -136,6 +185,18 @@ fn run(commands: Receiver) { &mut progress, ), }; + match &result { + Ok(output) => metrics.request_finished( + source, + request_started.elapsed(), + output.prompt_tokens, + output.cached_tokens, + output.completion_tokens, + output.previous_checkpoint_bytes, + output.checkpoint_bytes, + ), + Err(_) => metrics.request_failed(request_started.elapsed()), + } let _ = command.events.send(GenerationEvent::Finished(result)); last_used = Instant::now(); } @@ -143,6 +204,7 @@ fn run(commands: Receiver) { Err(mpsc::RecvTimeoutError::Timeout) => { if loaded.is_some() && last_used.elapsed() >= idle_timeout { loaded = None; + metrics.unloaded(); } } Err(mpsc::RecvTimeoutError::Disconnected) => break, diff --git a/src/server.rs b/src/server.rs index 01c65a0..b75a820 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,21 +1,25 @@ use crate::database::AppPreferences; use crate::engine::ChatTurn; +use crate::metrics::Metrics; use crate::model::{self, ModelChoice}; use crate::runtime::{CheckpointTarget, GenerationEvent, GenerationService}; use crate::settings::{ReasoningMode, effective_settings}; use serde::Deserialize; +use serde_json::value::RawValue; use serde_json::{Map, Value, json}; use std::collections::HashMap; +use std::fs::File; use std::io::{Read, Write}; use std::net::{TcpListener, TcpStream}; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use std::thread::{self, JoinHandle}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; const MAX_HEADER_BYTES: usize = 64 * 1024; const MAX_BODY_BYTES: usize = 64 * 1024 * 1024; +const PREFILL_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5); const TOOLS_PROMPT: &str = "## Tools\n\n\ You have access to a set of tools to help answer the user question. You can invoke tools by writing a \"<｜DSML｜tool_calls>\" block like the following:\n\n\ <｜DSML｜tool_calls>\n\ @@ -35,6 +39,8 @@ Otherwise, output directly after with tool calls or final response.\n\n pub(crate) struct ServerHandle { stop: Arc, thread: Option>, + metrics: Arc, + port: u16, } struct State { @@ -44,6 +50,7 @@ struct State { cache_path: PathBuf, sequence: AtomicU64, tool_memory: Mutex>, + metrics: Arc, } #[derive(Deserialize)] @@ -54,6 +61,8 @@ struct ChatRequest { messages: Vec, #[serde(default)] tools: Vec, + #[serde(skip)] + tool_schemas: Vec, #[serde(default)] tool_choice: Option, #[serde(default)] @@ -94,6 +103,8 @@ struct ApiMessage { reasoning_content: Value, #[serde(default)] tool_calls: Vec, + #[serde(default)] + tool_call_id: String, } #[derive(Clone, Deserialize)] @@ -124,6 +135,7 @@ enum OneOrMany { } struct ParsedRequest { + protocol: Protocol, model_id: String, messages: Vec, turn: crate::settings::TurnSettings, @@ -135,18 +147,338 @@ struct ParsedRequest { } struct ResponseOptions { + protocol: Protocol, + reasoning_summary: bool, model_id: String, stream: bool, include_usage: bool, has_tools: bool, } +#[derive(Clone, Copy, PartialEq, Eq)] +enum Protocol { + Chat, + Completion, + Anthropic, + Responses, +} + +enum ToolProjectionState { + Seeking, + Invokes, + Parameters, + Value, + Done, + Failed, +} + +enum ToolProjectionEvent { + Text(String), + Start { + index: usize, + id: String, + name: String, + }, + Arguments { + index: usize, + fragment: String, + }, + End { + index: usize, + }, +} + +struct ToolProjector { + raw: String, + position: usize, + text_emitted: usize, + state: ToolProjectionState, + index: usize, + ids: Vec, + first_parameter: bool, + string_parameter: bool, + syntax: Option, +} + +impl ToolProjector { + fn new() -> Self { + Self { + raw: String::new(), + position: 0, + text_emitted: 0, + state: ToolProjectionState::Seeking, + index: 0, + ids: Vec::new(), + first_parameter: true, + string_parameter: false, + syntax: None, + } + } + + fn push(&mut self, chunk: &str, final_chunk: bool, prefix: &str) -> Vec { + self.raw.push_str(chunk); + let mut events = Vec::new(); + loop { + match self.state { + ToolProjectionState::Seeking => { + if let Some((start, syntax)) = TOOL_SYNTAXES + .iter() + .filter_map(|syntax| { + self.raw + .find(syntax.tool_start) + .map(|start| (start, *syntax)) + }) + .min_by_key(|(start, _)| *start) + { + if start > self.text_emitted { + let text = self.raw[self.text_emitted..start].trim_end(); + if !text.is_empty() { + events.push(ToolProjectionEvent::Text(text.to_owned())); + } + } + self.position = start + syntax.tool_start.len(); + self.text_emitted = start; + self.syntax = Some(syntax); + self.state = ToolProjectionState::Invokes; + } else { + let limit = if final_chunk { + self.raw.len() + } else { + TOOL_SYNTAXES + .iter() + .map(|syntax| { + safe_before_partial_marker(&self.raw, syntax.tool_start) + }) + .min() + .unwrap_or(self.raw.len()) + }; + if limit > self.text_emitted { + let text = &self.raw[self.text_emitted..limit]; + if !text.trim().is_empty() { + events.push(ToolProjectionEvent::Text(text.to_owned())); + self.text_emitted = limit; + } + } + break; + } + } + ToolProjectionState::Invokes => { + let syntax = self.syntax.unwrap(); + self.skip_whitespace(); + if self.full_at(syntax.tool_end) { + self.position += syntax.tool_end.len(); + self.state = ToolProjectionState::Done; + break; + } + if self.partial_at(syntax.tool_end) || self.partial_at(syntax.invoke_start) { + break; + } + if !self.full_at(syntax.invoke_start) { + self.state = ToolProjectionState::Failed; + break; + } + let Some(tag_end) = self.raw[self.position..].find('>') else { + break; + }; + let tag_end = self.position + tag_end + 1; + let Some(name) = dsml_attribute(&self.raw[self.position..tag_end], "name") + else { + self.state = ToolProjectionState::Failed; + break; + }; + let id = random_tool_id(prefix); + self.ids.push(id.clone()); + events.push(ToolProjectionEvent::Start { + index: self.index, + id, + name, + }); + events.push(ToolProjectionEvent::Arguments { + index: self.index, + fragment: "{".into(), + }); + self.position = tag_end; + self.first_parameter = true; + self.state = ToolProjectionState::Parameters; + } + ToolProjectionState::Parameters => { + let syntax = self.syntax.unwrap(); + self.skip_whitespace(); + if self.full_at(syntax.invoke_end) { + events.push(ToolProjectionEvent::Arguments { + index: self.index, + fragment: "}".into(), + }); + events.push(ToolProjectionEvent::End { index: self.index }); + self.position += syntax.invoke_end.len(); + self.index += 1; + self.state = ToolProjectionState::Invokes; + continue; + } + if self.partial_at(syntax.invoke_end) || self.partial_at(syntax.parameter_start) + { + break; + } + if !self.full_at(syntax.parameter_start) { + self.state = ToolProjectionState::Failed; + break; + } + let Some(tag_end) = self.raw[self.position..].find('>') else { + break; + }; + let tag_end = self.position + tag_end + 1; + let tag = &self.raw[self.position..tag_end]; + let Some(name) = dsml_attribute(tag, "name") else { + self.state = ToolProjectionState::Failed; + break; + }; + self.string_parameter = + dsml_attribute(tag, "string").as_deref() != Some("false"); + let mut fragment = if self.first_parameter { + String::new() + } else { + ",".into() + }; + self.first_parameter = false; + fragment + .push_str(&serde_json::to_string(&name).unwrap_or_else(|_| "\"\"".into())); + fragment.push(':'); + if self.string_parameter { + fragment.push('"'); + } + events.push(ToolProjectionEvent::Arguments { + index: self.index, + fragment, + }); + self.position = tag_end; + self.state = ToolProjectionState::Value; + } + ToolProjectionState::Value => { + let syntax = self.syntax.unwrap(); + if let Some(relative_end) = self.raw[self.position..].find(syntax.parameter_end) + { + let end = self.position + relative_end; + self.emit_value(end, &mut events); + if self.string_parameter { + events.push(ToolProjectionEvent::Arguments { + index: self.index, + fragment: "\"".into(), + }); + } + self.position = end + syntax.parameter_end.len(); + self.state = ToolProjectionState::Parameters; + continue; + } + let limit = safe_parameter_value_limit( + &self.raw, + self.position, + syntax.parameter_end, + self.string_parameter, + ); + self.emit_value(limit, &mut events); + break; + } + ToolProjectionState::Done | ToolProjectionState::Failed => break, + } + } + events + } + + fn emit_value(&mut self, end: usize, events: &mut Vec) { + if end <= self.position { + return; + } + let raw = &self.raw[self.position..end]; + let fragment = if self.string_parameter { + let value = unescape_dsml(raw); + let encoded = serde_json::to_string(&value).unwrap_or_else(|_| "\"\"".into()); + encoded[1..encoded.len() - 1].to_owned() + } else { + raw.to_owned() + }; + events.push(ToolProjectionEvent::Arguments { + index: self.index, + fragment, + }); + self.position = end; + } + + fn skip_whitespace(&mut self) { + while self.raw[self.position..] + .chars() + .next() + .is_some_and(char::is_whitespace) + { + self.position += self.raw[self.position..].chars().next().unwrap().len_utf8(); + } + } + + fn full_at(&self, marker: &str) -> bool { + self.raw.as_bytes()[self.position..].starts_with(marker.as_bytes()) + } + + fn partial_at(&self, marker: &str) -> bool { + let tail = &self.raw.as_bytes()[self.position..]; + tail.len() < marker.len() && marker.as_bytes().starts_with(tail) + } +} + +fn dsml_attribute(tag: &str, name: &str) -> Option { + let start = tag.find(&format!("{name}=\""))? + name.len() + 2; + let end = start + tag[start..].find('"')?; + Some(unescape_dsml(&tag[start..end])) +} + +fn safe_before_partial_marker(text: &str, marker: &str) -> usize { + let mut limit = text.len().saturating_sub(marker.len().saturating_sub(1)); + while !text.is_char_boundary(limit) { + limit -= 1; + } + limit +} + +fn safe_parameter_value_limit(text: &str, start: usize, end_marker: &str, string: bool) -> usize { + let bytes = text.as_bytes(); + let marker = end_marker.as_bytes(); + let mut limit = bytes.len(); + for length in (1..marker.len().min(bytes.len().saturating_sub(start) + 1)).rev() { + if bytes[start..].ends_with(&marker[..length]) { + limit -= length; + break; + } + } + if string { + for entity in ["&", "<", ">", """, "'"] { + let entity = entity.as_bytes(); + for length in 1..entity.len() { + if bytes[start..limit].ends_with(&entity[..length]) { + limit -= length; + return limit; + } + } + } + } + limit +} + struct HttpRequest { method: String, path: String, body: Vec, } +#[derive(Deserialize)] +struct RawToolsRequest<'a> { + #[serde(default, borrow)] + tools: Vec<&'a RawValue>, +} + +#[derive(Deserialize)] +struct RawOpenAiTool<'a> { + #[serde(default, borrow)] + function: Option<&'a RawValue>, +} + impl ServerHandle { pub(crate) fn spawn( generation: GenerationService, @@ -154,6 +486,7 @@ impl ServerHandle { models_path: PathBuf, cache_path: PathBuf, port: u16, + metrics: Arc, ) -> Result { let address = format!("127.0.0.1:{port}"); let listener = TcpListener::bind(("127.0.0.1", port)) @@ -170,14 +503,18 @@ impl ServerHandle { cache_path, sequence: AtomicU64::new(0), tool_memory: Mutex::new(HashMap::new()), + metrics: Arc::clone(&metrics), }); let thread = thread::Builder::new() .name("local-http".into()) .spawn(move || serve(listener, state, worker_stop)) .map_err(|error| format!("Could not start the local HTTP service: {error}"))?; + metrics.server_listening(port); Ok(Self { stop, thread: Some(thread), + metrics, + port, }) } } @@ -188,6 +525,7 @@ impl Drop for ServerHandle { if let Some(thread) = self.thread.take() { let _ = thread.join(); } + self.metrics.server_stopped(self.port); } } @@ -213,24 +551,59 @@ fn serve(listener: TcpListener, state: Arc, stop: Arc) { fn handle(mut stream: TcpStream, state: &State) { let _ = stream.set_read_timeout(Some(Duration::from_secs(30))); + let started = Instant::now(); let request = match read_request(&mut stream) { Ok(request) => request, Err(error) => { + state.metrics.http_started("", 0); let _ = send_error(&mut stream, 400, &error); + state.metrics.http_finished(started.elapsed(), true); return; } }; - match (request.method.as_str(), request.path.as_str()) { - ("OPTIONS", _) => { - let _ = send_response(&mut stream, 204, None, &[]); - } + state + .metrics + .http_started(&request.path, request.body.len()); + let result = match (request.method.as_str(), request.path.as_str()) { + ("OPTIONS", _) => send_response(&mut stream, 204, None, &[]), ("GET", "/v1/models") => { let body = models_json(state); - let _ = send_json(&mut stream, 200, &body); + send_json(&mut stream, 200, &body) } ("POST", "/v1/chat/completions") => { - if let Err(error) = chat_completion(&mut stream, state, &request.body) { - let _ = send_error(&mut stream, error.0, &error.1); + match chat_completion(&mut stream, state, &request.body) { + Ok(()) => Ok(()), + Err(error) => { + let _ = send_error(&mut stream, error.0, &error.1); + Err(error.1) + } + } + } + ("POST", "/v1/completions") => { + match compatible_completion(&mut stream, state, &request.body, Protocol::Completion) { + Ok(()) => Ok(()), + Err(error) => { + let _ = send_error(&mut stream, error.0, &error.1); + Err(error.1) + } + } + } + ("POST", "/v1/messages") => { + match compatible_completion(&mut stream, state, &request.body, Protocol::Anthropic) { + Ok(()) => Ok(()), + Err(error) => { + let _ = send_error(&mut stream, error.0, &error.1); + Err(error.1) + } + } + } + ("POST", "/v1/responses") => { + match compatible_completion(&mut stream, state, &request.body, Protocol::Responses) { + Ok(()) => Ok(()), + Err(error) => { + let _ = send_error(&mut stream, error.0, &error.1); + Err(error.1) + } } } ("GET", path) if path.starts_with("/v1/models/") => { @@ -243,15 +616,20 @@ fn handle(mut stream: TcpStream, state: &State) { .preferences .read() .map_or(32_768, |preferences| preferences.context_tokens); - let _ = send_json(&mut stream, 200, &model_json(model, context)); + send_json(&mut stream, 200, &model_json(model, context)) } else { let _ = send_error(&mut stream, 404, "unknown model"); + Err("unknown model".into()) } } _ => { let _ = send_error(&mut stream, 404, "unknown endpoint"); + Err("unknown endpoint".into()) } - } + }; + state + .metrics + .http_finished(started.elapsed(), result.is_err()); } fn chat_completion( @@ -259,10 +637,16 @@ fn chat_completion( state: &State, body: &[u8], ) -> Result<(), (u16, String)> { - let request: ChatRequest = + let mut request: ChatRequest = serde_json::from_slice(body).map_err(|_| (400, "invalid JSON request".to_owned()))?; - let parsed = parse_chat_request(state, request)?; + request.tool_schemas = raw_tool_schemas(body, Protocol::Chat)?; + let parsed = parse_chat_request(state, request, Protocol::Chat)?; + if parsed.stream { + state.metrics.http_streaming(); + } let response = ResponseOptions { + protocol: parsed.protocol, + reasoning_summary: false, model_id: parsed.model_id, stream: parsed.stream, include_usage: parsed.include_usage, @@ -287,7 +671,539 @@ fn chat_completion( } } -fn parse_chat_request(state: &State, request: ChatRequest) -> Result { +fn compatible_completion( + stream: &mut TcpStream, + state: &State, + body: &[u8], + protocol: Protocol, +) -> Result<(), (u16, String)> { + let value: Value = + serde_json::from_slice(body).map_err(|_| (400, "invalid JSON request".to_owned()))?; + let reasoning_summary = protocol == Protocol::Responses + && matches!( + value.pointer("/reasoning/summary").and_then(Value::as_str), + Some("auto" | "concise" | "detailed") + ); + let mut request = match protocol { + Protocol::Completion => completion_request(value)?, + Protocol::Anthropic => anthropic_request(value)?, + Protocol::Responses => responses_request(value)?, + Protocol::Chat => unreachable!(), + }; + request.tool_schemas = raw_tool_schemas(body, protocol)?; + let parsed = parse_chat_request(state, request, protocol)?; + if parsed.stream { + state.metrics.http_streaming(); + } + let response = ResponseOptions { + protocol, + reasoning_summary, + model_id: parsed.model_id, + stream: parsed.stream, + include_usage: parsed.include_usage, + has_tools: parsed.has_tools, + }; + let active = state + .generation + .generate( + parsed.engine, + parsed.turn, + parsed.messages, + CheckpointTarget::Transient(state.cache_path.clone()), + parsed.idle_timeout, + ) + .map_err(|error| (500, error))?; + let sequence = state.sequence.fetch_add(1, Ordering::Relaxed) + 1; + let id = match protocol { + Protocol::Anthropic => format!("chatcmpl-{sequence}"), + Protocol::Responses => random_id("resp_"), + Protocol::Completion => format!("cmpl-{sequence}"), + Protocol::Chat => unreachable!(), + }; + if response.stream { + stream_response(stream, state, response, active, &id) + } else { + final_response(stream, state, response, active, &id) + } +} + +fn completion_request(mut value: Value) -> Result { + let object = value + .as_object_mut() + .ok_or_else(|| (400, "invalid JSON request".to_owned()))?; + let prompt = object + .remove("prompt") + .ok_or_else(|| (400, "missing prompt".to_owned()))?; + let prompt = match prompt { + Value::String(text) => text, + Value::Array(values) => values + .into_iter() + .next() + .and_then(|value| value.as_str().map(str::to_owned)) + .unwrap_or_default(), + _ => String::new(), + }; + object.insert( + "messages".into(), + json!([ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": prompt} + ]), + ); + decode_chat_request(value) +} + +fn anthropic_request(value: Value) -> Result { + let object = value + .as_object() + .ok_or_else(|| (400, "invalid JSON request".to_owned()))?; + let source = object + .get("messages") + .and_then(Value::as_array) + .ok_or_else(|| (400, "missing messages".to_owned()))?; + let mut messages = Vec::new(); + let system = object.get("system").map(content_text).unwrap_or_default(); + if !system.is_empty() { + messages.push(json!({"role": "system", "content": system})); + } + for message in source { + let role = message + .get("role") + .and_then(Value::as_str) + .unwrap_or("user"); + let Some(blocks) = message.get("content").and_then(Value::as_array) else { + messages.push(json!({"role": role, "content": message.get("content").cloned().unwrap_or(Value::Null)})); + continue; + }; + let mut text = String::new(); + let mut reasoning = String::new(); + let mut calls = Vec::new(); + for block in blocks { + match block.get("type").and_then(Value::as_str).unwrap_or("text") { + "text" => text.push_str(block.get("text").and_then(Value::as_str).unwrap_or("")), + "thinking" | "redacted_thinking" => reasoning.push_str( + block + .get("thinking") + .or_else(|| block.get("data")) + .and_then(Value::as_str) + .unwrap_or(""), + ), + "tool_use" => calls.push(json!({ + "id": block.get("id").and_then(Value::as_str).unwrap_or(""), + "type": "function", + "function": { + "name": block.get("name").and_then(Value::as_str).unwrap_or(""), + "arguments": block.get("input").cloned().unwrap_or_else(|| json!({})).to_string() + } + })), + "tool_result" => { + if !text.is_empty() { + messages.push(json!({"role": role, "content": std::mem::take(&mut text)})); + } + messages.push(json!({ + "role": "tool", + "content": block.get("content").map(content_text).unwrap_or_default(), + "tool_call_id": block.get("tool_use_id").and_then(Value::as_str).unwrap_or("") + })); + } + _ => {} + } + } + if !text.is_empty() || !reasoning.is_empty() || !calls.is_empty() { + messages.push(json!({ + "role": role, + "content": text, + "reasoning_content": reasoning, + "tool_calls": calls + })); + } + } + let mut request = Map::new(); + copy_request_fields(object, &mut request); + request.insert("messages".into(), Value::Array(messages)); + if let Some(stops) = object.get("stop_sequences") { + request.insert("stop".into(), stops.clone()); + } + if let Some(effort) = object + .get("output_config") + .and_then(|value| value.get("effort")) + { + request.insert("reasoning_effort".into(), effort.clone()); + } + request.insert( + "tools".into(), + normalize_anthropic_tools(object.get("tools")), + ); + decode_chat_request(Value::Object(request)) +} + +fn responses_request(value: Value) -> Result { + let object = value + .as_object() + .ok_or_else(|| (400, "invalid JSON request".to_owned()))?; + for key in ["previous_response_id", "conversation"] { + if object.get(key).is_some_and(|value| !value.is_null()) { + return Err(( + 400, + format!("{key} is not supported; replay full input instead"), + )); + } + } + if let Some(choice) = object.get("tool_choice") { + match choice { + Value::String(choice) if choice == "none" || choice == "auto" => {} + Value::String(choice) => { + return Err((400, format!("tool_choice={choice} not supported"))); + } + Value::Object(_) => return Err((400, "forced tool_choice not supported".into())), + _ => {} + } + } + let input = object + .get("input") + .ok_or_else(|| (400, "missing input".to_owned()))?; + let mut messages = Vec::new(); + match input { + Value::String(text) => messages.push(json!({"role": "user", "content": text})), + Value::Array(items) => { + let mut pending_reasoning = String::new(); + for item in items { + let item = item + .as_object() + .ok_or_else(|| (400, "invalid JSON request".to_owned()))?; + if item + .get("status") + .and_then(Value::as_str) + .is_some_and(|status| status != "completed") + { + return Err((400, "invalid JSON request".into())); + } + let kind = item + .get("type") + .and_then(Value::as_str) + .unwrap_or("message"); + let role = item.get("role").and_then(Value::as_str).unwrap_or("user"); + let consumes_reasoning = (kind == "message" && role == "assistant") + || matches!( + kind, + "function_call" + | "custom_tool_call" + | "local_shell_call" + | "web_search_call" + | "tool_search_call" + | "image_generation_call" + ); + let bookkeeping = matches!(kind, "compaction" | "context_compaction"); + if !consumes_reasoning && !bookkeeping && !pending_reasoning.is_empty() { + messages.push(reasoning_message(std::mem::take(&mut pending_reasoning))); + } + match kind { + "message" => { + let mut message = json!({ + "role": role, + "content": responses_content_text(item.get("content").unwrap_or(&Value::Null))? + }); + if role == "assistant" && !pending_reasoning.is_empty() { + message["reasoning_content"] = + Value::String(std::mem::take(&mut pending_reasoning)); + } + messages.push(message); + } + "function_call" | "custom_tool_call" => { + let name = if kind == "function_call" { + format!( + "{}{}", + item.get("namespace").and_then(Value::as_str).unwrap_or(""), + item.get("name").and_then(Value::as_str).unwrap_or("") + ) + } else { + item.get("name") + .and_then(Value::as_str) + .unwrap_or("") + .to_owned() + }; + push_responses_call( + &mut messages, + &mut pending_reasoning, + item, + &name, + item.get("arguments").or_else(|| item.get("input")), + ); + } + "function_call_output" | "custom_tool_call_output" => messages.push(json!({ + "role": "tool", + "content": responses_output_text(item.get("output"))?, + "tool_call_id": item.get("call_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or("") + })), + "reasoning" => { + for value in [item.get("summary"), item.get("content")] + .into_iter() + .flatten() + { + let text = responses_content_text(value)?; + if !pending_reasoning.is_empty() && !text.is_empty() { + pending_reasoning.push('\n'); + } + pending_reasoning.push_str(&text); + } + } + "local_shell_call" + | "web_search_call" + | "tool_search_call" + | "image_generation_call" => { + let name = match kind { + "local_shell_call" => "local_shell", + "tool_search_call" => "tool_search", + other => other, + }; + push_responses_call( + &mut messages, + &mut pending_reasoning, + item, + name, + item.get("action") + .or_else(|| item.get("arguments")) + .or_else(|| item.get("input")), + ); + } + "local_shell_call_output" + | "web_search_call_output" + | "tool_search_output" + | "tool_search_call_output" + | "image_generation_call_output" => { + let content = if let Some(value) = + item.get("output").or_else(|| item.get("result")) + { + responses_output_text(Some(value))? + } else { + item.get("tools").map(Value::to_string).unwrap_or_default() + }; + messages.push(json!({ + "role": "tool", + "content": content, + "tool_call_id": item.get("call_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or("") + })); + } + "compaction" | "context_compaction" => {} + _ => return Err((400, "invalid JSON request".into())), + } + } + if !pending_reasoning.is_empty() { + messages.push(reasoning_message(pending_reasoning)); + } + } + _ => return Err((400, "invalid JSON request".into())), + } + if let Some(instructions) = object.get("instructions") { + let instructions = match instructions { + Value::Null => String::new(), + Value::String(text) => text.clone(), + _ => return Err((400, "invalid JSON request".into())), + }; + if !instructions.is_empty() { + messages.insert(0, json!({"role": "system", "content": instructions})); + } + } + let mut request = Map::new(); + copy_request_fields(object, &mut request); + request.insert("messages".into(), Value::Array(messages)); + if let Some(tokens) = object.get("max_output_tokens") { + request.insert("max_tokens".into(), tokens.clone()); + } + if let Some(effort) = object + .get("reasoning") + .and_then(|value| value.get("effort")) + { + request.insert("reasoning_effort".into(), effort.clone()); + } + request.insert( + "tools".into(), + normalize_responses_tools(object.get("tools")), + ); + decode_chat_request(Value::Object(request)) +} + +fn reasoning_message(reasoning: String) -> Value { + json!({"role": "assistant", "content": "", "reasoning_content": reasoning}) +} + +fn push_responses_call( + messages: &mut Vec, + pending_reasoning: &mut String, + item: &Map, + name: &str, + arguments: Option<&Value>, +) { + let call = json!({ + "id": item.get("call_id").or_else(|| item.get("id")).and_then(Value::as_str).unwrap_or(""), + "type": "function", + "function": {"name": name, "arguments": raw_json_text(arguments)} + }); + if let Some(last) = messages.last_mut() + && last.get("role").and_then(Value::as_str) == Some("assistant") + { + if !pending_reasoning.is_empty() + && last + .get("reasoning_content") + .and_then(Value::as_str) + .is_none_or(str::is_empty) + { + last["reasoning_content"] = Value::String(std::mem::take(pending_reasoning)); + } + if !last.get("tool_calls").is_some_and(Value::is_array) { + last["tool_calls"] = json!([]); + } + last["tool_calls"].as_array_mut().unwrap().push(call); + } else { + messages.push(json!({ + "role": "assistant", + "content": "", + "reasoning_content": std::mem::take(pending_reasoning), + "tool_calls": [call] + })); + } +} + +fn raw_json_text(value: Option<&Value>) -> String { + match value { + Some(Value::String(text)) => text.clone(), + Some(value) => value.to_string(), + None => "{}".into(), + } +} + +fn responses_output_text(value: Option<&Value>) -> Result { + match value { + Some(Value::Array(_)) => responses_content_text(value.unwrap()), + Some(Value::String(text)) => Ok(text.clone()), + Some(value) => Ok(value.to_string()), + None => Ok(String::new()), + } +} + +fn responses_content_text(value: &Value) -> Result { + match value { + Value::String(text) => Ok(text.clone()), + Value::Null => Ok(String::new()), + Value::Array(parts) => { + let mut text = String::new(); + for part in parts { + match part { + Value::String(part) => text.push_str(part), + Value::Object(part) + if matches!( + part.get("type").and_then(Value::as_str), + Some( + "input_text" + | "output_text" + | "text" + | "summary_text" + | "reasoning_text" + ) + ) && matches!( + part.get("text"), + Some(Value::String(_) | Value::Null) + ) => + { + if let Some(value) = part.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + _ => return Err((400, "invalid JSON request".into())), + } + } + Ok(text) + } + _ => Err((400, "invalid JSON request".into())), + } +} + +fn copy_request_fields(source: &Map, target: &mut Map) { + for key in [ + "model", + "max_tokens", + "max_completion_tokens", + "temperature", + "top_p", + "min_p", + "top_k", + "seed", + "stream", + "stream_options", + "thinking", + "think", + "reasoning_effort", + "tool_choice", + "stop", + ] { + if let Some(value) = source.get(key) { + target.insert(key.into(), value.clone()); + } + } +} + +fn normalize_anthropic_tools(tools: Option<&Value>) -> Value { + Value::Array( + tools + .and_then(Value::as_array) + .into_iter() + .flatten() + .map(|tool| { + json!({"type": "function", "function": { + "name": tool.get("name").and_then(Value::as_str).unwrap_or(""), + "description": tool.get("description").and_then(Value::as_str).unwrap_or(""), + "parameters": tool.get("input_schema").cloned().unwrap_or_else(|| json!({"type": "object"})) + }}) + }) + .collect(), + ) +} + +fn normalize_responses_tools(tools: Option<&Value>) -> Value { + Value::Array( + tools + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter(|tool| tool.get("type").and_then(Value::as_str) == Some("function")) + .map(|tool| { + json!({"type": "function", "function": { + "name": tool.get("name").and_then(Value::as_str).unwrap_or(""), + "description": tool.get("description").and_then(Value::as_str).unwrap_or(""), + "parameters": tool.get("parameters").cloned().unwrap_or_else(|| json!({"type": "object"})) + }}) + }) + .collect(), + ) +} + +fn decode_chat_request(value: Value) -> Result { + serde_json::from_value(value).map_err(|_| (400, "invalid JSON request".to_owned())) +} + +fn raw_tool_schemas(body: &[u8], protocol: Protocol) -> Result, (u16, String)> { + let request: RawToolsRequest<'_> = + serde_json::from_slice(body).map_err(|_| (400, "invalid JSON request".to_owned()))?; + request + .tools + .into_iter() + .map(|tool| { + if protocol == Protocol::Chat { + let wrapper: RawOpenAiTool<'_> = serde_json::from_str(tool.get()) + .map_err(|_| (400, "invalid JSON request".to_owned()))?; + Ok(wrapper.function.unwrap_or(tool).get().to_owned()) + } else { + Ok(tool.get().to_owned()) + } + }) + .collect() +} + +fn parse_chat_request( + state: &State, + request: ChatRequest, + protocol: Protocol, +) -> Result { if request.messages.is_empty() { return Err((400, "missing messages".into())); } @@ -322,10 +1238,16 @@ fn parse_chat_request(state: &State, request: ChatRequest) -> Result 0); generation.reasoning_mode = request_reasoning(&request, requested_id)?; - let tools_enabled = !request.tools.is_empty() + let tools_enabled = (!request.tools.is_empty() || !request.tool_schemas.is_empty()) && request.tool_choice.as_ref().and_then(Value::as_str) != Some("none"); - let (system, messages) = - render_messages(state, &request.messages, &request.tools, tools_enabled)?; + let (system, messages) = render_messages( + state, + &request.messages, + &request.tools, + &request.tool_schemas, + tools_enabled, + protocol, + )?; generation.system_prompt = system; let runtime = preferences.runtime().map_err(|error| (500, error))?; let mut effective = effective_settings(model, &generation, &runtime, &state.models_path) @@ -342,6 +1264,7 @@ fn parse_chat_request(state: &State, request: ChatRequest) -> Result Result<(String, Vec), (u16, String)> { + validate_tool_results(state, messages, protocol)?; let preserve_reasoning = tools_enabled || messages.iter().any(|message| { matches!(message.role.as_str(), "tool" | "function") || !message.tool_calls.is_empty() @@ -408,16 +1334,26 @@ fn render_messages( let mut system = String::new(); if tools_enabled { system.push_str(TOOLS_PROMPT); - for tool in tools { - let schema = tool.get("function").unwrap_or(tool); - if !system.ends_with("\n\n") { + if tool_schemas.is_empty() { + for tool in tools { + let schema = tool.get("function").unwrap_or(tool); + if !system.ends_with("\n\n") { + system.push('\n'); + } + system.push_str( + &serde_json::to_string(schema) + .map_err(|error| (400, format!("invalid tool schema: {error}")))?, + ); + system.push('\n'); + } + } else { + for schema in tool_schemas { + if !system.ends_with("\n\n") { + system.push('\n'); + } + system.push_str(schema); system.push('\n'); } - system.push_str( - &serde_json::to_string(schema) - .map_err(|error| (400, format!("invalid tool schema: {error}")))?, - ); - system.push('\n'); } system.push_str( "\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. Use the exact parameter names from the schemas.", @@ -436,6 +1372,7 @@ fn render_messages( } "user" => turns.push(ChatTurn { user: true, + skip_previous_eos: false, reasoning: None, reasoning_complete: true, content, @@ -453,6 +1390,7 @@ fn render_messages( } else { turns.push(ChatTurn { user: true, + skip_previous_eos: protocol == Protocol::Responses, reasoning: None, reasoning_complete: true, content: wrapped, @@ -467,6 +1405,7 @@ fn render_messages( let reasoning = content_text(&message.reasoning_content); turns.push(ChatTurn { user: false, + skip_previous_eos: false, reasoning: (preserve_reasoning && !reasoning.is_empty()).then_some(reasoning), reasoning_complete: true, content, @@ -478,6 +1417,44 @@ fn render_messages( Ok((system, turns)) } +fn validate_tool_results( + state: &State, + messages: &[ApiMessage], + protocol: Protocol, +) -> Result<(), (u16, String)> { + if !matches!(protocol, Protocol::Anthropic | Protocol::Responses) { + return Ok(()); + } + let memory = state.tool_memory.lock().ok(); + for (index, message) in messages.iter().enumerate() { + if !matches!(message.role.as_str(), "tool" | "function") || message.tool_call_id.is_empty() + { + continue; + } + let id = &message.tool_call_id; + let live = memory + .as_ref() + .is_some_and(|memory| memory.contains_key(id)); + let replayed = messages[..index].iter().any(|message| { + message.role == "assistant" && message.tool_calls.iter().any(|call| call.id == *id) + }); + if live || replayed { + continue; + } + let message = match protocol { + Protocol::Anthropic => format!( + "Anthropic continuation state is not available for tool_use_id {id}; retry by replaying the full messages history" + ), + Protocol::Responses => format!( + "Responses continuation state is not available for call_id {id}; retry by replaying the full input history" + ), + _ => unreachable!(), + }; + return Err((400, message)); + } + Ok(()) +} + fn replayed_or_canonical_tools(state: &State, calls: &[ApiToolCall]) -> String { if let Ok(memory) = state.tool_memory.lock() && let Some(raw) = calls.iter().find_map(|call| { @@ -536,27 +1513,112 @@ fn final_response( id: &str, ) -> Result<(), (u16, String)> { let output = wait_for_output(active)?; - let (content, calls) = parse_generated_tools(state, &output.message.content); + let (content, calls) = parse_generated_tools(state, &output.message.content, request.protocol); let finish = if calls.is_empty() { output.finish_reason } else { "tool_calls" }; - let mut message = json!({"role": "assistant", "content": content}); - if let Some(reasoning) = output.message.reasoning.filter(|value| !value.is_empty()) { - message["reasoning_content"] = Value::String(reasoning); - } - if !calls.is_empty() { - message["tool_calls"] = tool_calls_json(&calls); - } - let body = json!({ - "id": id, - "object": "chat.completion", - "created": unix_time(), - "model": request.model_id, - "choices": [{"index": 0, "message": message, "finish_reason": finish}], - "usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens), - }); + let reasoning = output.message.reasoning.filter(|value| !value.is_empty()); + let usage = usage_json( + output.prompt_tokens, + output.cached_tokens, + output.completion_tokens, + ); + let body = match request.protocol { + Protocol::Chat => { + let mut message = json!({"role": "assistant", "content": content}); + if let Some(reasoning) = reasoning { + message["reasoning_content"] = Value::String(reasoning); + } + if !calls.is_empty() { + message["tool_calls"] = tool_calls_json(&calls); + } + json!({ + "id": id, "object": "chat.completion", "created": unix_time(), + "model": request.model_id, + "choices": [{"index": 0, "message": message, "finish_reason": finish}], + "usage": usage, + }) + } + Protocol::Completion => json!({ + "id": id, "object": "text_completion", "created": unix_time(), + "model": request.model_id, + "choices": [{"text": content, "index": 0, "finish_reason": finish}], + "usage": usage, + }), + Protocol::Anthropic => { + let mut blocks = Vec::new(); + if let Some(reasoning) = reasoning { + blocks.push(json!({"type": "thinking", "thinking": reasoning, "signature": id})); + } + if !content.is_empty() { + blocks.push(json!({"type": "text", "text": content})); + } + for call in &calls { + blocks.push(json!({ + "type": "tool_use", "id": call.id, "name": call.function.name, + "input": serde_json::from_str::(&call.function.arguments).unwrap_or_else(|_| json!({})) + })); + } + if blocks.is_empty() || (blocks.iter().all(|block| block["type"] == "thinking")) { + blocks.push(json!({"type": "text", "text": ""})); + } + let cached = output.cached_tokens.min(output.prompt_tokens); + let written = output.prompt_tokens - cached; + json!({ + "id": id, "type": "message", "role": "assistant", "model": request.model_id, + "content": blocks, + "stop_reason": if finish == "tool_calls" { "tool_use" } else if finish == "length" { "max_tokens" } else { "end_turn" }, + "stop_sequence": Value::Null, + "usage": { + "input_tokens": output.prompt_tokens - cached - written, + "output_tokens": output.completion_tokens, + "cache_read_input_tokens": cached, + "cache_creation_input_tokens": written + } + }) + } + Protocol::Responses => { + let status = if finish == "length" { + "incomplete" + } else if finish == "error" { + "failed" + } else { + "completed" + }; + let mut items = Vec::new(); + if let Some(reasoning) = reasoning { + items.push(json!({ + "id": random_id("rs_"), "type": "reasoning", "status": status, + "summary": [{"type": "summary_text", "text": reasoning}] + })); + } + if !content.is_empty() { + items.push(json!({ + "id": random_id("msg_"), "type": "message", "status": status, + "role": "assistant", "content": [{"type": "output_text", "text": content, "annotations": []}] + })); + } + for call in &calls { + items.push(json!({ + "id": random_id("fc_"), "type": "function_call", "status": status, + "name": call.function.name, "call_id": call.id, "arguments": call.function.arguments + })); + } + json!({ + "id": id, "object": "response", "created_at": unix_time(), "status": status, + "model": request.model_id, "output": items, + "usage": { + "input_tokens": output.prompt_tokens, + "input_tokens_details": {"cached_tokens": output.cached_tokens.min(output.prompt_tokens), "cache_write_tokens": output.prompt_tokens - output.cached_tokens.min(output.prompt_tokens)}, + "output_tokens": output.completion_tokens, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": output.prompt_tokens + output.completion_tokens + } + }) + } + }; send_json(stream, 200, &body).map_err(|error| (500, error)) } @@ -566,26 +1628,645 @@ fn stream_response( request: ResponseOptions, active: crate::runtime::ActiveGeneration, id: &str, +) -> Result<(), (u16, String)> { + if request.protocol != Protocol::Chat { + return structured_stream_response(stream, state, request, active, id); + } + stream_response_with_keepalive( + stream, + state, + request, + active, + id, + PREFILL_KEEPALIVE_INTERVAL, + ) +} + +fn structured_stream_response( + stream: &mut TcpStream, + state: &State, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, +) -> Result<(), (u16, String)> { + match request.protocol { + Protocol::Completion => completion_stream_response(stream, request, active, id), + Protocol::Anthropic => anthropic_stream_response(stream, state, request, active, id), + Protocol::Responses => responses_stream_response(stream, state, request, active, id), + Protocol::Chat => unreachable!(), + } +} + +fn completion_stream_response( + stream: &mut TcpStream, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, ) -> Result<(), (u16, String)> { send_sse_headers(stream).map_err(|error| (500, error))?; - let role = chunk_json(id, &request.model_id, json!({"role": "assistant"}), None); - send_sse(stream, &role).map_err(|error| (500, error))?; - let mut buffered_content = String::new(); - let mut buffered_reasoning = String::new(); - let mut output = None; while let Ok(event) = active.events.recv() { match event { - GenerationEvent::Loading | GenerationEvent::Context { .. } => {} + GenerationEvent::Chunk { + reasoning: false, + content, + } if !content.is_empty() => send_sse( + stream, + &json!({ + "id": id, "object": "text_completion", "created": unix_time(), + "model": request.model_id, + "choices": [{"text": content, "index": 0, "finish_reason": Value::Null}] + }), + ) + .map_err(|error| (500, error))?, + GenerationEvent::Finished(Ok(output)) => { + send_sse( + stream, + &json!({ + "id": id, "object": "text_completion", "created": unix_time(), + "model": request.model_id, + "choices": [{"text": "", "index": 0, "finish_reason": output.finish_reason}] + }), + ) + .map_err(|error| (500, error))?; + if request.include_usage { + send_sse( + stream, + &json!({ + "id": id, "object": "text_completion", "created": unix_time(), + "model": request.model_id, "choices": [], + "usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens) + }), + ) + .map_err(|error| (500, error))?; + } + return stream + .write_all(b"data: [DONE]\n\n") + .map_err(|error| (500, error.to_string())); + } + GenerationEvent::Finished(Err(error)) => { + let _ = send_sse_error(stream, &error); + return Ok(()); + } + _ => {} + } + } + Err((500, "The model runtime stopped unexpectedly.".into())) +} + +fn anthropic_stream_response( + stream: &mut TcpStream, + state: &State, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, +) -> Result<(), (u16, String)> { + send_sse_headers(stream).map_err(|error| (500, error))?; + let mut prompt_tokens = 0; + let mut started = false; + let mut block = None::<(usize, bool)>; + let mut next_index = 0; + let mut projector = ToolProjector::new(); + let mut tool_indices = Vec::new(); + while let Ok(event) = active.events.recv() { + match event { + GenerationEvent::Context { + used, + tokens_per_second, + .. + } => { + if tokens_per_second.is_none() { + prompt_tokens = used; + } else if !started { + anthropic_stream_start(stream, &request, id, prompt_tokens, 0)?; + started = true; + } + } GenerationEvent::Chunk { reasoning, content } => { - if reasoning { - buffered_reasoning.push_str(&content); + if !started { + anthropic_stream_start(stream, &request, id, prompt_tokens, 0)?; + started = true; + } + if !reasoning && request.has_tools { + let events = projector.push(&content, false, "toolu_"); + send_anthropic_projection_events( + stream, + events, + &mut block, + &mut next_index, + &mut tool_indices, + )?; + continue; + } + if block.is_some_and(|(_, current_reasoning)| current_reasoning != reasoning) { + let (index, _) = block.take().unwrap(); + send_named_sse( + stream, + "content_block_stop", + &json!({"type": "content_block_stop", "index": index}), + )?; + } + let index = if let Some((index, _)) = block { + index } else { - buffered_content.push_str(&content); + let index = next_index; + next_index += 1; + send_named_sse( + stream, + "content_block_start", + &json!({ + "type": "content_block_start", "index": index, + "content_block": if reasoning { json!({"type": "thinking", "thinking": "", "signature": ""}) } else { json!({"type": "text", "text": ""}) } + }), + )?; + block = Some((index, reasoning)); + index + }; + let delta = if reasoning { + json!({"type": "thinking_delta", "thinking": content}) + } else { + json!({"type": "text_delta", "text": content}) + }; + send_named_sse( + stream, + "content_block_delta", + &json!({"type": "content_block_delta", "index": index, "delta": delta}), + )?; + } + GenerationEvent::Finished(Ok(output)) => { + if !started { + anthropic_stream_start( + stream, + &request, + id, + output.prompt_tokens, + output.cached_tokens, + )?; } - if request.has_tools && buffered_content.contains("") { - active.cancel.store(true, Ordering::Relaxed); + if request.has_tools { + let events = projector.push("", true, "toolu_"); + send_anthropic_projection_events( + stream, + events, + &mut block, + &mut next_index, + &mut tool_indices, + )?; } - if !request.has_tools && !content.is_empty() { + if let Some((index, _)) = block.take() { + send_named_sse( + stream, + "content_block_stop", + &json!({"type": "content_block_stop", "index": index}), + )?; + } + let (_, calls) = parse_generated_tools_with_ids( + state, + &output.message.content, + Protocol::Anthropic, + &projector.ids, + ); + if projector.ids.is_empty() { + for call in &calls { + send_named_sse( + stream, + "content_block_start", + &json!({"type": "content_block_start", "index": next_index, "content_block": {"type": "tool_use", "id": call.id, "name": call.function.name, "input": {}}}), + )?; + send_named_sse( + stream, + "content_block_delta", + &json!({"type": "content_block_delta", "index": next_index, "delta": {"type": "input_json_delta", "partial_json": call.function.arguments}}), + )?; + send_named_sse( + stream, + "content_block_stop", + &json!({"type": "content_block_stop", "index": next_index}), + )?; + next_index += 1; + } + } + let finish = if calls.is_empty() { + output.finish_reason + } else { + "tool_calls" + }; + send_named_sse( + stream, + "message_delta", + &json!({"type": "message_delta", "delta": {"stop_reason": if finish == "tool_calls" { "tool_use" } else if finish == "length" { "max_tokens" } else { "end_turn" }, "stop_sequence": Value::Null}, "usage": {"output_tokens": output.completion_tokens}}), + )?; + return send_named_sse(stream, "message_stop", &json!({"type": "message_stop"})); + } + GenerationEvent::Finished(Err(error)) => { + return send_named_sse( + stream, + "error", + &json!({"type": "error", "error": {"type": "api_error", "message": error}}), + ); + } + _ => {} + } + } + Err((500, "The model runtime stopped unexpectedly.".into())) +} + +fn send_anthropic_projection_events( + stream: &mut impl Write, + events: Vec, + block: &mut Option<(usize, bool)>, + next_index: &mut usize, + tool_indices: &mut Vec, +) -> Result<(), (u16, String)> { + for event in events { + match event { + ToolProjectionEvent::Text(text) if !text.is_empty() => { + if block.is_some_and(|(_, reasoning)| reasoning) { + let (index, _) = block.take().unwrap(); + send_named_sse( + stream, + "content_block_stop", + &json!({"type": "content_block_stop", "index": index}), + )?; + } + let index = if let Some((index, _)) = *block { + index + } else { + let index = *next_index; + *next_index += 1; + send_named_sse( + stream, + "content_block_start", + &json!({"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}}), + )?; + *block = Some((index, false)); + index + }; + send_named_sse( + stream, + "content_block_delta", + &json!({"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}}), + )?; + } + ToolProjectionEvent::Start { index, id, name } => { + if let Some((open_index, _)) = block.take() { + send_named_sse( + stream, + "content_block_stop", + &json!({"type": "content_block_stop", "index": open_index}), + )?; + } + let content_index = *next_index; + *next_index += 1; + if tool_indices.len() == index { + tool_indices.push(content_index); + } + send_named_sse( + stream, + "content_block_start", + &json!({"type": "content_block_start", "index": content_index, "content_block": {"type": "tool_use", "id": id, "name": name, "input": {}}}), + )?; + } + ToolProjectionEvent::Arguments { index, fragment } => { + if let Some(content_index) = tool_indices.get(index) { + send_named_sse( + stream, + "content_block_delta", + &json!({"type": "content_block_delta", "index": content_index, "delta": {"type": "input_json_delta", "partial_json": fragment}}), + )?; + } + } + ToolProjectionEvent::End { index } => { + if let Some(content_index) = tool_indices.get(index) { + send_named_sse( + stream, + "content_block_stop", + &json!({"type": "content_block_stop", "index": content_index}), + )?; + } + } + ToolProjectionEvent::Text(_) => {} + } + } + Ok(()) +} + +fn anthropic_stream_start( + stream: &mut impl Write, + request: &ResponseOptions, + id: &str, + prompt_tokens: u32, + cached_tokens: u32, +) -> Result<(), (u16, String)> { + let cached = cached_tokens.min(prompt_tokens); + let written = prompt_tokens - cached; + send_named_sse( + stream, + "message_start", + &json!({"type": "message_start", "message": {"id": id, "type": "message", "role": "assistant", "model": request.model_id, "content": [], "stop_reason": Value::Null, "stop_sequence": Value::Null, "usage": {"input_tokens": prompt_tokens - cached - written, "output_tokens": 0, "cache_read_input_tokens": cached, "cache_creation_input_tokens": written}}}), + ) +} + +fn responses_stream_response( + stream: &mut TcpStream, + state: &State, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, +) -> Result<(), (u16, String)> { + send_sse_headers(stream).map_err(|error| (500, error))?; + let created = unix_time(); + let message_id = random_id("msg_"); + let reasoning_id = random_id("rs_"); + let mut sequence = 0; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.created", "response": {"id": id, "object": "response", "created_at": created, "status": "in_progress", "model": request.model_id, "output": []}}), + )?; + let mut reasoning_open = false; + let mut message_open = false; + let mut reasoning = String::new(); + let mut content = String::new(); + while let Ok(event) = active.events.recv() { + match event { + GenerationEvent::Chunk { + reasoning: true, + content: chunk, + } => { + if !request.reasoning_summary { + continue; + } + if !reasoning_open { + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.added", "output_index": 0, "item": {"id": reasoning_id, "type": "reasoning", "status": "in_progress", "summary": []}}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.reasoning_summary_part.added", "item_id": reasoning_id, "output_index": 0, "summary_index": 0, "part": {"type": "summary_text", "text": ""}}), + )?; + reasoning_open = true; + } + reasoning.push_str(&chunk); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.reasoning_summary_text.delta", "item_id": reasoning_id, "output_index": 0, "summary_index": 0, "delta": chunk}), + )?; + } + GenerationEvent::Chunk { + reasoning: false, + content: chunk, + } => { + if request.has_tools { + content.push_str(&chunk); + continue; + } + if !message_open { + let output_index = usize::from(reasoning_open); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.added", "output_index": output_index, "item": {"id": message_id, "type": "message", "status": "in_progress", "role": "assistant", "content": []}}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.content_part.added", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}), + )?; + message_open = true; + } + content.push_str(&chunk); + let output_index = usize::from(reasoning_open); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_text.delta", "item_id": message_id, "output_index": output_index, "content_index": 0, "delta": chunk}), + )?; + } + GenerationEvent::Finished(Ok(output)) => { + let (parsed_content, calls) = + parse_generated_tools(state, &output.message.content, Protocol::Responses); + if request.has_tools { + content = parsed_content; + } + let finish = if calls.is_empty() { + output.finish_reason + } else { + "tool_calls" + }; + let status = if finish == "length" { + "incomplete" + } else if finish == "error" { + "failed" + } else { + "completed" + }; + let mut terminal_items = Vec::new(); + let mut output_index = 0; + if reasoning_open { + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.reasoning_summary_text.done", "item_id": reasoning_id, "output_index": output_index, "summary_index": 0, "text": reasoning}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.reasoning_summary_part.done", "item_id": reasoning_id, "output_index": output_index, "summary_index": 0, "part": {"type": "summary_text", "text": reasoning}}), + )?; + let item = json!({"id": reasoning_id, "type": "reasoning", "status": status, "summary": [{"type": "summary_text", "text": reasoning}]}); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.done", "output_index": output_index, "item": item}), + )?; + terminal_items.push(item); + output_index += 1; + } + if !content.is_empty() { + if !message_open { + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.added", "output_index": output_index, "item": {"id": message_id, "type": "message", "status": "in_progress", "role": "assistant", "content": []}}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.content_part.added", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_text.delta", "item_id": message_id, "output_index": output_index, "content_index": 0, "delta": content}), + )?; + } + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_text.done", "item_id": message_id, "output_index": output_index, "content_index": 0, "text": content}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.content_part.done", "item_id": message_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": content, "annotations": []}}), + )?; + let item = json!({"id": message_id, "type": "message", "status": status, "role": "assistant", "content": [{"type": "output_text", "text": content, "annotations": []}]}); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.done", "output_index": output_index, "item": item}), + )?; + terminal_items.push(item); + output_index += 1; + } + for call in &calls { + let item_id = random_id("fc_"); + let mut item = json!({"id": item_id, "type": "function_call", "status": status, "name": call.function.name, "call_id": call.id, "arguments": call.function.arguments}); + let mut added = item.clone(); + added["status"] = Value::String("in_progress".into()); + added["arguments"] = Value::String(String::new()); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.added", "output_index": output_index, "item": added}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.function_call_arguments.delta", "item_id": item_id, "output_index": output_index, "delta": call.function.arguments}), + )?; + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.function_call_arguments.done", "item_id": item_id, "output_index": output_index, "name": call.function.name, "arguments": call.function.arguments}), + )?; + item["id"] = Value::String(item_id); + send_responses_sse( + stream, + &mut sequence, + json!({"type": "response.output_item.done", "output_index": output_index, "item": item}), + )?; + terminal_items.push(item); + output_index += 1; + } + let event_type = if finish == "length" { + "response.incomplete" + } else if finish == "error" { + "response.failed" + } else { + "response.completed" + }; + let cached = output.cached_tokens.min(output.prompt_tokens); + return send_responses_sse( + stream, + &mut sequence, + json!({"type": event_type, "response": {"id": id, "object": "response", "created_at": created, "status": status, "model": request.model_id, "output": terminal_items, "usage": {"input_tokens": output.prompt_tokens, "input_tokens_details": {"cached_tokens": cached, "cache_write_tokens": output.prompt_tokens - cached}, "output_tokens": output.completion_tokens, "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": output.prompt_tokens + output.completion_tokens}}}), + ); + } + GenerationEvent::Finished(Err(error)) => { + let _ = send_sse_error(stream, &error); + return Ok(()); + } + _ => {} + } + } + Err((500, "The model runtime stopped unexpectedly.".into())) +} + +#[allow(clippy::too_many_arguments)] +fn send_named_sse( + stream: &mut impl Write, + event: &str, + value: &Value, +) -> Result<(), (u16, String)> { + let body = serde_json::to_vec(value).map_err(|error| (500, error.to_string()))?; + stream + .write_all(b"event: ") + .and_then(|()| stream.write_all(event.as_bytes())) + .and_then(|()| stream.write_all(b"\ndata: ")) + .and_then(|()| stream.write_all(&body)) + .and_then(|()| stream.write_all(b"\n\n")) + .map_err(|error| (500, error.to_string())) +} + +fn send_responses_sse( + stream: &mut impl Write, + sequence: &mut u32, + value: Value, +) -> Result<(), (u16, String)> { + let mut object = value + .as_object() + .cloned() + .ok_or_else(|| (500, "Responses event is not an object".to_owned()))?; + let event_type = object.shift_remove("type").unwrap_or(Value::Null); + let mut ordered = Map::new(); + ordered.insert("type".into(), event_type); + ordered.insert("sequence_number".into(), Value::from(*sequence)); + ordered.extend(object); + *sequence += 1; + send_sse(stream, &Value::Object(ordered)).map_err(|error| (500, error)) +} + +fn stream_response_with_keepalive( + stream: &mut TcpStream, + state: &State, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, + keepalive_interval: Duration, +) -> Result<(), (u16, String)> { + let mut projector = ToolProjector::new(); + let mut output = None; + let mut prefilling = true; + let mut headers_sent = false; + let mut role_sent = false; + let mut last_keepalive = Instant::now(); + loop { + let event = match receive_stream_event( + stream, + &active, + prefilling && headers_sent, + &mut last_keepalive, + keepalive_interval, + ) { + Ok(Some(event)) => event, + Ok(None) => break, + Err(_) => return Ok(()), + }; + match event { + GenerationEvent::Loading => {} + GenerationEvent::Context { + tokens_per_second, .. + } => { + prefilling = tokens_per_second.is_none(); + if prefilling && !headers_sent { + send_sse_headers(stream).map_err(|error| (500, error))?; + headers_sent = true; + last_keepalive = Instant::now(); + } else if !prefilling { + send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?; + } + } + GenerationEvent::Chunk { reasoning, content } => { + prefilling = false; + send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?; + if request.has_tools && !reasoning { + let events = projector.push(&content, false, "call_"); + if send_chat_projection_events(stream, &request, id, events).is_err() { + active.cancel.store(true, Ordering::Relaxed); + return Ok(()); + } + if TOOL_SYNTAXES + .iter() + .any(|syntax| projector.raw.contains(syntax.tool_end)) + { + active.cancel.store(true, Ordering::Relaxed); + } + } else if !content.is_empty() { let field = if reasoning { "reasoning_content" } else { @@ -600,57 +2281,51 @@ fn stream_response( } GenerationEvent::Finished(result) => { match result { - Ok(result) => output = Some(result), + Ok(result) => { + send_stream_start(stream, &request, id, &mut headers_sent, &mut role_sent)?; + output = Some(result); + } Err(error) => { - let event = json!({"error": {"message": error, "type": "server_error"}}); - let _ = send_sse(stream, &event); - let _ = stream.write_all(b"data: [DONE]\n\n"); - return Ok(()); + if headers_sent { + let _ = send_sse_error(stream, &error); + return Ok(()); + } + return Err((500, error)); } } break; } } } - let output = output.ok_or_else(|| (500, "The model runtime stopped unexpectedly.".into()))?; - let (content, calls) = parse_generated_tools(state, &output.message.content); + let output = match output { + Some(output) => output, + None if headers_sent => { + let _ = send_sse_error(stream, "The model runtime stopped unexpectedly."); + return Ok(()); + } + None => return Err((500, "The model runtime stopped unexpectedly.".into())), + }; if request.has_tools { - if let Some(reasoning) = output - .message - .reasoning - .as_deref() - .filter(|value| !value.is_empty()) - { - send_sse( - stream, - &chunk_json( - id, - &request.model_id, - json!({"reasoning_content": reasoning}), - None, - ), - ) - .map_err(|error| (500, error))?; - } - if !content.is_empty() { - send_sse( - stream, - &chunk_json(id, &request.model_id, json!({"content": content}), None), - ) - .map_err(|error| (500, error))?; - } - if !calls.is_empty() { - send_sse( - stream, - &chunk_json( - id, - &request.model_id, - json!({"tool_calls": tool_calls_json(&calls)}), - None, - ), - ) - .map_err(|error| (500, error))?; - } + let events = projector.push("", true, "call_"); + send_chat_projection_events(stream, &request, id, events)?; + } + let (_, calls) = parse_generated_tools_with_ids( + state, + &output.message.content, + Protocol::Chat, + &projector.ids, + ); + if request.has_tools && projector.ids.is_empty() && !calls.is_empty() { + send_sse( + stream, + &chunk_json( + id, + &request.model_id, + json!({"tool_calls": tool_calls_json(&calls)}), + None, + ), + ) + .map_err(|error| (500, error))?; } let finish = if calls.is_empty() { output.finish_reason @@ -678,6 +2353,81 @@ fn stream_response( .map_err(|error| (500, error.to_string())) } +fn send_chat_projection_events( + stream: &mut impl Write, + request: &ResponseOptions, + response_id: &str, + events: Vec, +) -> Result<(), (u16, String)> { + for event in events { + let delta = match event { + ToolProjectionEvent::Text(content) if !content.is_empty() => { + json!({"content": content}) + } + ToolProjectionEvent::Start { index, id, name } => json!({"tool_calls": [{ + "index": index, "id": id, "type": "function", + "function": {"name": name, "arguments": ""} + }]}), + ToolProjectionEvent::Arguments { index, fragment } => json!({ + "tool_calls": [{"index": index, "function": {"arguments": fragment}}] + }), + ToolProjectionEvent::End { .. } | ToolProjectionEvent::Text(_) => continue, + }; + send_sse( + stream, + &chunk_json(response_id, &request.model_id, delta, None), + ) + .map_err(|error| (500, error))?; + } + Ok(()) +} + +fn send_stream_start( + stream: &mut impl Write, + request: &ResponseOptions, + id: &str, + headers_sent: &mut bool, + role_sent: &mut bool, +) -> Result<(), (u16, String)> { + if !*headers_sent { + send_sse_headers(stream).map_err(|error| (500, error))?; + *headers_sent = true; + } + if !*role_sent { + let role = chunk_json(id, &request.model_id, json!({"role": "assistant"}), None); + send_sse(stream, &role).map_err(|error| (500, error))?; + *role_sent = true; + } + Ok(()) +} + +fn receive_stream_event( + stream: &mut impl Write, + active: &crate::runtime::ActiveGeneration, + prefilling: bool, + last_keepalive: &mut Instant, + keepalive_interval: Duration, +) -> Result, String> { + if !prefilling { + return Ok(active.events.recv().ok()); + } + loop { + if last_keepalive.elapsed() >= keepalive_interval { + if let Err(error) = stream.write_all(b": prefill\n\n") { + active.cancel.store(true, Ordering::Relaxed); + return Err(error.to_string()); + } + *last_keepalive = Instant::now(); + } + let remaining = keepalive_interval.saturating_sub(last_keepalive.elapsed()); + match active.events.recv_timeout(remaining) { + Ok(event) => return Ok(Some(event)), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {} + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => return Ok(None), + } + } +} + fn wait_for_output( active: crate::runtime::ActiveGeneration, ) -> Result { @@ -689,7 +2439,10 @@ fn wait_for_output( content: chunk, } => { content.push_str(&chunk); - if content.contains("") { + if TOOL_SYNTAXES + .iter() + .any(|syntax| content.contains(syntax.tool_end)) + { active.cancel.store(true, Ordering::Relaxed); } } @@ -702,46 +2455,83 @@ fn wait_for_output( Err((500, "The model runtime stopped unexpectedly.".into())) } -fn parse_generated_tools(state: &State, text: &str) -> (String, Vec) { - let Some(start) = text.find("<｜DSML｜tool_calls>") else { +fn parse_generated_tools( + state: &State, + text: &str, + protocol: Protocol, +) -> (String, Vec) { + parse_generated_tools_with_ids(state, text, protocol, &[]) +} + +fn parse_generated_tools_with_ids( + state: &State, + text: &str, + protocol: Protocol, + streamed_ids: &[String], +) -> (String, Vec) { + let Some((start, syntax)) = TOOL_SYNTAXES + .iter() + .filter_map(|syntax| text.find(syntax.tool_start).map(|start| (start, *syntax))) + .min_by_key(|(start, _)| *start) + else { return (text.to_owned(), Vec::new()); }; - let Some(relative_end) = text[start..].find("") else { + let Some(relative_end) = text[start..].find(syntax.tool_end) else { return (text.to_owned(), Vec::new()); }; - let end = start + relative_end + "".len(); - let raw = &text[start..end]; + let end = start + relative_end + syntax.tool_end.len(); + let content = text[..start].trim_end(); + let raw = &text[content.len()..end]; let mut calls = Vec::new(); - let mut cursor = 0; - while let Some(relative) = raw[cursor..].find("<｜DSML｜invoke name=\"") { - let name_start = cursor + relative + "<｜DSML｜invoke name=\"".len(); - let Some(name_end_relative) = raw[name_start..].find("\">") else { + let mut cursor = raw.find(syntax.tool_start).unwrap() + syntax.tool_start.len(); + loop { + skip_text_whitespace(raw, &mut cursor); + if raw[cursor..].starts_with(syntax.tool_end) { + break; + } + if !raw[cursor..].starts_with(syntax.invoke_start) { + return (text.to_owned(), Vec::new()); + } + let Some(tag_end) = raw[cursor..].find('>').map(|end| cursor + end + 1) else { return (text.to_owned(), Vec::new()); }; - let name_end = name_start + name_end_relative; - let body_start = name_end + 2; - let Some(body_end_relative) = raw[body_start..].find("") else { + let Some(name) = dsml_attribute(&raw[cursor..tag_end], "name") else { return (text.to_owned(), Vec::new()); }; - let body_end = body_start + body_end_relative; - let arguments = parse_dsml_arguments(&raw[body_start..body_end]); + cursor = tag_end; + let mut arguments = Map::new(); + loop { + skip_text_whitespace(raw, &mut cursor); + if raw[cursor..].starts_with(syntax.invoke_end) { + cursor += syntax.invoke_end.len(); + break; + } + let Some((name, value)) = parse_tool_parameter(raw, &mut cursor, syntax) else { + return (text.to_owned(), Vec::new()); + }; + arguments.insert(name, value); + } calls.push(ApiToolCall { id: String::new(), function: ApiFunction { - name: unescape_dsml(&raw[name_start..name_end]), + name, arguments: Value::Object(arguments).to_string(), }, }); - cursor = body_end + "".len(); } if calls.is_empty() { return (text.to_owned(), Vec::new()); } - let sequence = state - .sequence - .fetch_add(calls.len() as u64, Ordering::Relaxed); + let prefix = if protocol == Protocol::Anthropic { + "toolu_" + } else { + "call_" + }; for (index, call) in calls.iter_mut().enumerate() { - call.id = format!("call_{:032x}", sequence + index as u64 + 1); + call.id = streamed_ids + .get(index) + .cloned() + .unwrap_or_else(|| random_tool_id(prefix)); } if let Ok(mut memory) = state.tool_memory.lock() { // ponytail: one process-local replay table; add LRU eviction if 100k live tool ids is measured insufficient. @@ -752,38 +2542,100 @@ fn parse_generated_tools(state: &State, text: &str) -> (String, Vec memory.insert(call.id.clone(), raw.to_owned()); } } - (text[..start].trim_end().to_owned(), calls) + (content.to_owned(), calls) } -fn parse_dsml_arguments(body: &str) -> Map { - let mut arguments = Map::new(); - let mut cursor = 0; - while let Some(relative) = body[cursor..].find("<｜DSML｜parameter name=\"") { - let name_start = cursor + relative + "<｜DSML｜parameter name=\"".len(); - let Some(name_end_relative) = body[name_start..].find("\" string=\"") else { - break; - }; - let name_end = name_start + name_end_relative; - let flag_start = name_end + "\" string=\"".len(); - let Some(flag_end_relative) = body[flag_start..].find("\">") else { - break; - }; - let flag_end = flag_start + flag_end_relative; - let value_start = flag_end + 2; - let Some(value_end_relative) = body[value_start..].find("") else { - break; - }; - let value_end = value_start + value_end_relative; - let raw = unescape_dsml(&body[value_start..value_end]); - let value = if &body[flag_start..flag_end] == "true" { - Value::String(raw) - } else { - serde_json::from_str(&raw).unwrap_or(Value::String(raw)) - }; - arguments.insert(unescape_dsml(&body[name_start..name_end]), value); - cursor = value_end + "".len(); +#[derive(Clone, Copy)] +struct ToolSyntax { + tool_start: &'static str, + tool_end: &'static str, + invoke_start: &'static str, + invoke_end: &'static str, + parameter_start: &'static str, + parameter_end: &'static str, +} + +const TOOL_SYNTAXES: [ToolSyntax; 3] = [ + ToolSyntax { + tool_start: "<｜DSML｜tool_calls>", + tool_end: "", + invoke_start: "<｜DSML｜invoke", + invoke_end: "", + parameter_start: "<｜DSML｜parameter", + parameter_end: "", + }, + ToolSyntax { + tool_start: "", + tool_end: "", + invoke_start: "", + parameter_start: "", + }, + ToolSyntax { + tool_start: "", + tool_end: "", + invoke_start: "", + parameter_start: "", + }, +]; + +fn parse_tool_parameter( + text: &str, + cursor: &mut usize, + syntax: ToolSyntax, +) -> Option<(String, Value)> { + if !text[*cursor..].starts_with(syntax.parameter_start) { + return None; + } + let tag_end = text[*cursor..].find('>').map(|end| *cursor + end + 1)?; + let tag = &text[*cursor..tag_end]; + let name = dsml_attribute(tag, "name")?; + let is_string = dsml_attribute(tag, "string"); + *cursor = tag_end; + let mut nested_start = *cursor; + skip_text_whitespace(text, &mut nested_start); + if is_string.is_none() && text[nested_start..].starts_with(syntax.parameter_start) { + *cursor = nested_start; + let mut nested = Map::new(); + loop { + skip_text_whitespace(text, cursor); + if !text[*cursor..].starts_with(syntax.parameter_start) { + break; + } + let (name, value) = parse_tool_parameter(text, cursor, syntax)?; + nested.insert(name, value); + } + skip_text_whitespace(text, cursor); + if !text[*cursor..].starts_with(syntax.parameter_end) { + return None; + } + *cursor += syntax.parameter_end.len(); + return Some((name, Value::Object(nested))); + } + let value_end = text[*cursor..] + .find(syntax.parameter_end) + .map(|end| *cursor + end)?; + let raw = &text[*cursor..value_end]; + *cursor = value_end + syntax.parameter_end.len(); + let value = if is_string.as_deref().unwrap_or("true") == "true" { + Value::String(unescape_dsml(raw)) + } else { + serde_json::from_str(raw).unwrap_or(Value::Null) + }; + Some((name, value)) +} + +fn skip_text_whitespace(text: &str, cursor: &mut usize) { + while text[*cursor..] + .chars() + .next() + .is_some_and(char::is_whitespace) + { + *cursor += text[*cursor..].chars().next().unwrap().len_utf8(); } - arguments } fn tool_calls_json(calls: &[ApiToolCall]) -> Value { @@ -923,7 +2775,7 @@ fn chunk_json(id: &str, model: &str, delta: Value, finish: Option<&str>) -> Valu }) } -fn send_sse_headers(stream: &mut TcpStream) -> Result<(), String> { +fn send_sse_headers(stream: &mut impl Write) -> Result<(), String> { stream .write_all( b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n", @@ -931,7 +2783,7 @@ fn send_sse_headers(stream: &mut TcpStream) -> Result<(), String> { .map_err(|error| error.to_string()) } -fn send_sse(stream: &mut TcpStream, value: &Value) -> Result<(), String> { +fn send_sse(stream: &mut impl Write, value: &Value) -> Result<(), String> { let body = serde_json::to_vec(value).map_err(|error| error.to_string())?; stream .write_all(b"data: ") @@ -940,6 +2792,16 @@ fn send_sse(stream: &mut TcpStream, value: &Value) -> Result<(), String> { stream.write_all(b"\n\n").map_err(|error| error.to_string()) } +fn send_sse_error(stream: &mut impl Write, message: &str) -> Result<(), String> { + let body = serde_json::to_vec(&json!({"error": {"message": message, "type": "server_error"}})) + .map_err(|error| error.to_string())?; + stream + .write_all(b"event: error\ndata: ") + .map_err(|error| error.to_string())?; + stream.write_all(&body).map_err(|error| error.to_string())?; + stream.write_all(b"\n\n").map_err(|error| error.to_string()) +} + fn send_json(stream: &mut TcpStream, code: u16, value: &Value) -> Result<(), String> { let mut body = serde_json::to_vec(value).map_err(|error| error.to_string())?; body.push(b'\n'); @@ -1061,6 +2923,35 @@ fn unix_time() -> u64 { .as_secs() } +fn random_id(prefix: &str) -> String { + random_hex_id::<12>(prefix) +} + +fn random_tool_id(prefix: &str) -> String { + random_hex_id::<16>(prefix) +} + +fn random_hex_id(prefix: &str) -> String { + let mut bytes = [0_u8; N]; + if File::open("/dev/urandom") + .and_then(|mut file| file.read_exact(&mut bytes)) + .is_err() + { + let fallback = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + bytes.copy_from_slice(&fallback.to_le_bytes()[..N]); + } + let mut id = String::with_capacity(prefix.len() + N * 2); + id.push_str(prefix); + for byte in bytes { + use std::fmt::Write as _; + let _ = write!(id, "{byte:02x}"); + } + id +} + fn default_role() -> String { "user".into() } @@ -1073,6 +2964,76 @@ fn empty_arguments() -> String { mod tests { use super::*; + #[test] + fn tool_prompt_matches_ds4_server_text() { + assert_eq!(TOOLS_PROMPT.len(), 1_183); + } + + #[test] + fn dsml_tool_stream_projects_incremental_json_arguments() { + let mut projector = ToolProjector::new(); + let mut fragments = Vec::new(); + for (chunk, final_chunk) in [ + ("\n\n", false), + ( + "<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"echo\">\n<｜DSML｜parameter name=\"text\" string=\"true\">h&am", + false, + ), + ( + "p;i\n\n", + true, + ), + ] { + for event in projector.push(chunk, final_chunk, "call_") { + match event { + ToolProjectionEvent::Start { name, .. } => { + fragments.push(format!("start:{name}")); + } + ToolProjectionEvent::Arguments { fragment, .. } => fragments.push(fragment), + ToolProjectionEvent::End { .. } => fragments.push("end".into()), + ToolProjectionEvent::Text(_) => {} + } + } + } + assert_eq!( + fragments, + [ + "start:echo", + "{", + "\"text\":\"", + "h", + "&i", + "\"", + "}", + "end" + ] + ); + assert_eq!(projector.ids.len(), 1); + } + + #[test] + fn tool_stream_projects_ds4_recovery_syntaxes() { + for text in [ + "hi", + "hi", + ] { + let mut projector = ToolProjector::new(); + let events = projector.push(text, true, "call_"); + assert!(events.iter().any(|event| matches!( + event, + ToolProjectionEvent::Start { name, .. } if name == "echo" + ))); + let arguments = events + .into_iter() + .filter_map(|event| match event { + ToolProjectionEvent::Arguments { fragment, .. } => Some(fragment), + _ => None, + }) + .collect::(); + assert_eq!(arguments, r#"{"text":"hi"}"#); + } + } + #[test] fn canonical_tool_replay_preserves_argument_types() { let calls = vec![ApiToolCall { @@ -1090,20 +3051,153 @@ mod tests { #[test] fn generated_dsml_becomes_openai_tool_calls() { + let metrics = Arc::new(Metrics::new(PathBuf::new().as_path())); let state = State { - generation: GenerationService::spawn().unwrap(), + generation: GenerationService::spawn(Arc::clone(&metrics)).unwrap(), preferences: Arc::new(RwLock::new(AppPreferences::default())), models_path: PathBuf::new(), cache_path: PathBuf::new(), sequence: AtomicU64::new(0), tool_memory: Mutex::new(HashMap::new()), + metrics, }; let text = "done\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"shell\">\n<｜DSML｜parameter name=\"command\" string=\"true\">pwd\n\n"; - let (content, calls) = parse_generated_tools(&state, text); + let (content, calls) = parse_generated_tools(&state, text, Protocol::Chat); assert_eq!(content, "done"); assert_eq!(calls.len(), 1); assert_eq!(calls[0].function.name, "shell"); assert_eq!(calls[0].function.arguments, r#"{"command":"pwd"}"#); + assert_eq!(calls[0].id.len(), 37); + assert_eq!(state.sequence.load(Ordering::Relaxed), 0); + assert!( + state + .tool_memory + .lock() + .unwrap() + .get(&calls[0].id) + .unwrap() + .starts_with("\n\n<｜DSML｜tool_calls>") + ); + + let (_, calls) = parse_generated_tools(&state, text, Protocol::Anthropic); + assert!(calls[0].id.starts_with("toolu_")); + assert_eq!(calls[0].id.len(), 38); + } + + #[test] + fn generated_tool_parser_accepts_ds4_recovery_syntaxes() { + let metrics = Arc::new(Metrics::new(PathBuf::new().as_path())); + let state = State { + generation: GenerationService::spawn(Arc::clone(&metrics)).unwrap(), + preferences: Arc::new(RwLock::new(AppPreferences::default())), + models_path: PathBuf::new(), + cache_path: PathBuf::new(), + sequence: AtomicU64::new(0), + tool_memory: Mutex::new(HashMap::new()), + metrics, + }; + let short = "done\n\n/tmp"; + let (content, calls) = parse_generated_tools(&state, short, Protocol::Chat); + assert_eq!(content, "done"); + assert_eq!(calls[0].function.name, "run"); + assert_eq!( + calls[0].function.arguments, + r#"{"options":{"path":"/tmp"}}"# + ); + + let plain = "invalid"; + let (_, calls) = parse_generated_tools(&state, plain, Protocol::Chat); + assert_eq!(calls[0].function.arguments, r#"{"count":null}"#); + } + + #[test] + fn protocol_tool_results_require_live_or_replayed_call_ids() { + let metrics = Arc::new(Metrics::new(PathBuf::new().as_path())); + let state = State { + generation: GenerationService::spawn(Arc::clone(&metrics)).unwrap(), + preferences: Arc::new(RwLock::new(AppPreferences::default())), + models_path: PathBuf::new(), + cache_path: PathBuf::new(), + sequence: AtomicU64::new(0), + tool_memory: Mutex::new(HashMap::new()), + metrics, + }; + let tool_result: ApiMessage = serde_json::from_value(json!({ + "role": "tool", "content": "ok", "tool_call_id": "call_missing" + })) + .unwrap(); + let error = validate_tool_results(&state, &[tool_result], Protocol::Responses).unwrap_err(); + assert_eq!(error.0, 400); + assert!( + error + .1 + .contains("Responses continuation state is not available") + ); + + let replay: Vec = serde_json::from_value(json!([ + {"role": "assistant", "tool_calls": [{"id": "call_replay", "function": {"name": "echo", "arguments": "{}"}}]}, + {"role": "tool", "content": "ok", "tool_call_id": "call_replay"} + ])) + .unwrap(); + assert!(validate_tool_results(&state, &replay, Protocol::Responses).is_ok()); + + state + .tool_memory + .lock() + .unwrap() + .insert("toolu_live".into(), String::new()); + let live: ApiMessage = serde_json::from_value(json!({ + "role": "tool", "content": "ok", "tool_call_id": "toolu_live" + })) + .unwrap(); + assert!(validate_tool_results(&state, &[live], Protocol::Anthropic).is_ok()); + } + + #[test] + fn responses_parser_rejects_dropped_state_and_preserves_hosted_calls() { + let error = responses_request(json!({ + "input": "hello", "previous_response_id": "resp_old" + })) + .err() + .unwrap(); + assert_eq!( + error.1, + "previous_response_id is not supported; replay full input instead" + ); + assert_eq!( + responses_request(json!({"input": "hello", "tool_choice": "required"})) + .err() + .unwrap() + .1, + "tool_choice=required not supported" + ); + assert!(responses_request(json!({"input": [{"type": "unknown"}]})).is_err()); + assert!( + responses_request(json!({ + "input": [{"type": "message", "role": "user", "content": [{"type": "image", "url": "x"}]}] + })) + .is_err() + ); + + let request = responses_request(json!({ + "input": [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "why"}]}, + {"type": "local_shell_call", "id": "call_1", "action": {"command": "pwd"}, "status": "completed"}, + {"type": "local_shell_call_output", "call_id": "call_1", "output": "/tmp"} + ] + })) + .unwrap(); + assert_eq!(request.messages[0].reasoning_content, json!("why")); + assert_eq!( + request.messages[0].tool_calls[0].function.name, + "local_shell" + ); + assert_eq!( + request.messages[0].tool_calls[0].function.arguments, + r#"{"command":"pwd"}"# + ); + assert_eq!(request.messages[1].tool_call_id, "call_1"); + assert_eq!(request.messages[1].content, json!("/tmp")); } #[test] @@ -1113,4 +3207,63 @@ mod tests { "one two" ); } + + #[test] + fn streaming_prefill_keeps_the_client_connection_alive() { + let (events, receiver) = std::sync::mpsc::channel(); + let active = crate::runtime::ActiveGeneration { + events: receiver, + cancel: Arc::new(AtomicBool::new(false)), + }; + let producer = thread::spawn(move || { + thread::sleep(Duration::from_millis(25)); + events.send(GenerationEvent::Loading).unwrap(); + }); + let mut response = Vec::new(); + let event = receive_stream_event( + &mut response, + &active, + true, + &mut Instant::now(), + Duration::from_millis(10), + ) + .unwrap(); + producer.join().unwrap(); + assert!(matches!(event, Some(GenerationEvent::Loading))); + assert!(response.starts_with(b": prefill\n\n")); + + let request = ResponseOptions { + protocol: Protocol::Chat, + reasoning_summary: false, + model_id: "deepseek-v4-flash".into(), + stream: true, + include_usage: false, + has_tools: true, + }; + let mut response = Vec::new(); + send_sse_headers(&mut response).unwrap(); + response.write_all(b": prefill\n\n").unwrap(); + let mut headers_sent = true; + let mut role_sent = false; + send_stream_start( + &mut response, + &request, + "chatcmpl-test", + &mut headers_sent, + &mut role_sent, + ) + .unwrap(); + let response = String::from_utf8(response).unwrap(); + assert!(response.find(": prefill").unwrap() < response.find("assistant").unwrap()); + } + + #[test] + fn streaming_errors_match_ds4_server_framing() { + let mut response = Vec::new(); + send_sse_error(&mut response, "prefill failed").unwrap(); + assert_eq!( + String::from_utf8(response).unwrap(), + "event: error\ndata: {\"error\":{\"message\":\"prefill failed\",\"type\":\"server_error\"}}\n\n" + ); + } }