swadhin / justapi public

swadhinbiswas

[email protected]

2026-07-09 09:00:00 +0600 committed NaN years ago

d743f45

commit d743f45

23 files changed

+9975 -0

Files

crates

justapi-inference

src

real

lora.rs +334 mod.rs +17 model.rs +500 quant.rs +274 tokenizer.rs +144 autoscaler.rs +254 control_plane.rs +458 engine.rs +257 gateway.rs +302 kv_cache.rs +773 lib.rs +74 model.rs +187 openai.rs +1064 pd.rs +349 radix_cache.rs +776 router.rs +416 scheduler_engine.rs +693 scheduler.rs +1144 spec_decode_tree.rs +861 spec_decode.rs +535 supervisor.rs +355

tests

engine.rs +177 Cargo.toml +31

Showing 23 changed files+9975-0

▼crates/justapi-inference/Cargo.toml

+31-0

@@ -0,0 +1,31 @@
	1	+[package]
	2	+name = "justapi-inference"
	3	+version.workspace = true
	4	+edition.workspace = true
	5	+license.workspace = true
	6	+
	7	+[dependencies]
	8	+tokio = { workspace = true }
	9	+anyhow.workspace = true
	10	+thiserror.workspace = true
	11	+tracing.workspace = true
	12	+serde = { workspace = true }
	13	+serde_json = { workspace = true }
	14	+futures = "0.3"
	15	+bytes = "1.12.0"
	16	+async-stream = "0.3"
	17	+rand = "0.8"
	18	+candle-core = "0.11.0"
	19	+candle-nn = "0.11.0"
	20	+candle-transformers = { version = "0.11.0", optional = true }
	21	+tokenizers = { version = "0.22", optional = true }
	22	+
	23	+[features]
	24	+default = []
	25	+# GPU backends (feature-gated; require the matching toolkit to COMPILE).
	26	+# CPU is the default backend and always available.
	27	+cuda = ["candle-core/cuda", "candle-nn/cuda"]
	28	+# Real weight loading + tokenizer (GGUF/safetensors). Heavy: pulls in
	29	+# candle-transformers + tokenizers. NOT in default features so the lean
	30	+# build/test path stays fast.
	31	+real = ["dep:candle-transformers", "dep:tokenizers"]
0	32

▼crates/justapi-inference/src/autoscaler.rs

+254-0

@@ -0,0 +1,254 @@
	1	+//! LLM-specific autoscaler — scaling decisions from LLM health signals.
	2	+//!
	3	+//! Unlike a generic HTTP autoscaler (which watches QPS), this watches the
	4	+//! signals that actually matter for LLM serving:
	5	+//!
	6	+//! - TTFT (time to first token) — prefill latency / saturation.
	7	+//! - tokens/sec — decode throughput per replica.
	8	+//! - queue depth — pending requests waiting for a slot.
	9	+//! - KV-cache pressure — the binding memory constraint for long contexts.
	10	+//!
	11	+//! Pure decision logic (no timers, no metrics collection, no orchestration).
	12	+//! The [`Autoscaler`] is stateful only for cooldown/flap protection; pair it
	13	+//! with the metrics pipeline and the multi-replica supervisor.
	14	+
	15	+/// Live LLM-serving signals fed to the autoscaler each tick.
	16	+#[derive(Debug, Clone, Copy)]
	17	+pub struct LlmMetrics {
	18	+ /// Decode throughput per replica (tokens/sec).
	19	+ pub tokens_per_sec: f32,
	20	+ /// Time to first token, milliseconds.
	21	+ pub ttft_ms: f32,
	22	+ /// Number of requests waiting for a free slot.
	23	+ pub queue_depth: usize,
	24	+ /// KV-cache pressure across replicas (0–100%).
	25	+ pub kv_pressure_pct: f32,
	26	+}
	27	+
	28	+/// Tuning for the autoscaler.
	29	+#[derive(Debug, Clone, Copy)]
	30	+pub struct AutoscalerConfig {
	31	+ /// Never scale below this many replicas.
	32	+ pub min_replicas: usize,
	33	+ /// Never scale above this many replicas.
	34	+ pub max_replicas: usize,
	35	+ /// TTFT above which we consider the fleet saturated.
	36	+ pub target_ttft_ms: f32,
	37	+ /// KV-cache pressure (0–100) above which we consider the fleet saturated.
	38	+ pub target_kv_pressure_pct: f32,
	39	+ /// Queue depth at/above which we scale up.
	40	+ pub queue_depth_scale_up: usize,
	41	+ /// Queue depth at/below which we may scale down.
	42	+ pub queue_depth_scale_down: usize,
	43	+ /// Ticks to wait after a scaling action before another is allowed
	44	+ /// (prevents flapping on noisy metrics).
	45	+ pub cooldown_ticks: usize,
	46	+}
	47	+
	48	+impl Default for AutoscalerConfig {
	49	+ fn default() -> Self {
	50	+ Self {
	51	+ min_replicas: 1,
	52	+ max_replicas: 8,
	53	+ target_ttft_ms: 500.0,
	54	+ target_kv_pressure_pct: 80.0,
	55	+ queue_depth_scale_up: 16,
	56	+ queue_depth_scale_down: 2,
	57	+ cooldown_ticks: 10,
	58	+ }
	59	+ }
	60	+}
	61	+
	62	+/// What the autoscaler wants the supervisor to do.
	63	+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	64	+pub enum ScaleDecision {
	65	+ /// Add replicas up to `desired`.
	66	+ ScaleUp(usize),
	67	+ /// Remove replicas down to `desired`.
	68	+ ScaleDown(usize),
	69	+ /// Keep `desired` replicas (no change).
	70	+ Hold(usize),
	71	+}
	72	+
	73	+/// Stateful LLM autoscaler.
	74	+pub struct Autoscaler {
	75	+ config: AutoscalerConfig,
	76	+ current: usize,
	77	+ cooldown: usize,
	78	+}
	79	+
	80	+impl Autoscaler {
	81	+ /// Create an autoscaler starting at `config.min_replicas`.
	82	+ pub fn new(config: AutoscalerConfig) -> Self {
	83	+ let current = config.min_replicas.max(1).min(config.max_replicas.max(1));
	84	+ Self {
	85	+ config,
	86	+ current,
	87	+ cooldown: 0,
	88	+ }
	89	+ }
	90	+
	91	+ /// Current desired replica count.
	92	+ pub fn current(&self) -> usize {
	93	+ self.current
	94	+ }
	95	+
	96	+ /// Whether the fleet is currently in a cooldown window.
	97	+ pub fn in_cooldown(&self) -> bool {
	98	+ self.cooldown > 0
	99	+ }
	100	+
	101	+ /// Decide the next action from the latest metrics.
	102	+ ///
	103	+ /// - Scale up when KV pressure, TTFT, or queue depth breach their
	104	+ /// saturation thresholds.
	105	+ /// - Scale down when all signals are comfortably under half their
	106	+ /// targets and queue depth is low.
	107	+ /// - Hold otherwise, or while in cooldown.
	108	+ pub fn decide(&mut self, m: &LlmMetrics) -> ScaleDecision {
	109	+ if self.cooldown > 0 {
	110	+ self.cooldown -= 1;
	111	+ return ScaleDecision::Hold(self.current);
	112	+ }
	113	+
	114	+ let saturated = m.kv_pressure_pct >= self.config.target_kv_pressure_pct
	115	+ \|\| m.ttft_ms >= self.config.target_ttft_ms
	116	+ \|\| m.queue_depth >= self.config.queue_depth_scale_up;
	117	+
	118	+ let underused = m.kv_pressure_pct <= self.config.target_kv_pressure_pct * 0.5
	119	+ && m.ttft_ms <= self.config.target_ttft_ms * 0.5
	120	+ && m.queue_depth <= self.config.queue_depth_scale_down;
	121	+
	122	+ let desired = if saturated {
	123	+ (self.current + 1).min(self.config.max_replicas)
	124	+ } else if underused {
	125	+ self.current.saturating_sub(1).max(self.config.min_replicas)
	126	+ } else {
	127	+ self.current
	128	+ };
	129	+
	130	+ if desired == self.current {
	131	+ ScaleDecision::Hold(self.current)
	132	+ } else if desired > self.current {
	133	+ self.current = desired;
	134	+ self.cooldown = self.config.cooldown_ticks;
	135	+ ScaleDecision::ScaleUp(desired)
	136	+ } else {
	137	+ self.current = desired;
	138	+ self.cooldown = self.config.cooldown_ticks;
	139	+ ScaleDecision::ScaleDown(desired)
	140	+ }
	141	+ }
	142	+}
	143	+
	144	+// ---------------------------------------------------------------------------
	145	+// Tests
	146	+// ---------------------------------------------------------------------------
	147	+
	148	+#[cfg(test)]
	149	+mod tests {
	150	+ use super::*;
	151	+
	152	+ fn cfg() -> AutoscalerConfig {
	153	+ AutoscalerConfig {
	154	+ min_replicas: 1,
	155	+ max_replicas: 4,
	156	+ target_ttft_ms: 500.0,
	157	+ target_kv_pressure_pct: 80.0,
	158	+ queue_depth_scale_up: 16,
	159	+ queue_depth_scale_down: 2,
	160	+ cooldown_ticks: 2,
	161	+ }
	162	+ }
	163	+
	164	+ fn metrics(kv: f32, ttft: f32, q: usize) -> LlmMetrics {
	165	+ LlmMetrics {
	166	+ tokens_per_sec: 50.0,
	167	+ ttft_ms: ttft,
	168	+ queue_depth: q,
	169	+ kv_pressure_pct: kv,
	170	+ }
	171	+ }
	172	+
	173	+ #[test]
	174	+ fn scales_up_on_kv_pressure() {
	175	+ let mut a = Autoscaler::new(cfg());
	176	+ let d = a.decide(&metrics(95.0, 100.0, 0));
	177	+ assert_eq!(d, ScaleDecision::ScaleUp(2));
	178	+ assert!(a.in_cooldown());
	179	+ }
	180	+
	181	+ #[test]
	182	+ fn scales_up_on_ttft() {
	183	+ let mut a = Autoscaler::new(cfg());
	184	+ let d = a.decide(&metrics(10.0, 1200.0, 0));
	185	+ assert_eq!(d, ScaleDecision::ScaleUp(2));
	186	+ }
	187	+
	188	+ #[test]
	189	+ fn scales_up_on_queue_depth() {
	190	+ let mut a = Autoscaler::new(cfg());
	191	+ let d = a.decide(&metrics(10.0, 100.0, 40));
	192	+ assert_eq!(d, ScaleDecision::ScaleUp(2));
	193	+ }
	194	+
	195	+ #[test]
	196	+ fn holds_when_healthy() {
	197	+ let mut a = Autoscaler::new(cfg());
	198	+ let d = a.decide(&metrics(40.0, 200.0, 1));
	199	+ assert_eq!(d, ScaleDecision::Hold(1));
	200	+ }
	201	+
	202	+ #[test]
	203	+ fn scales_down_when_underused() {
	204	+ let mut a = Autoscaler::new(cfg());
	205	+ // First bump to 2, then drain metrics to trigger a scale down.
	206	+ let _ = a.decide(&metrics(95.0, 100.0, 0));
	207	+ // cooldown ticks
	208	+ let _ = a.decide(&metrics(10.0, 100.0, 0));
	209	+ let _ = a.decide(&metrics(10.0, 100.0, 0));
	210	+ let d = a.decide(&metrics(10.0, 100.0, 0));
	211	+ // still underused and out of cooldown → ScaleDown back to 1
	212	+ assert_eq!(d, ScaleDecision::ScaleDown(1));
	213	+ }
	214	+
	215	+ #[test]
	216	+ fn never_exceeds_max() {
	217	+ let mut a = Autoscaler::new(cfg());
	218	+ // Drive up to max with cooldown ticks between.
	219	+ for _ in 0..10 {
	220	+ let _ = a.decide(&metrics(99.0, 2000.0, 100));
	221	+ }
	222	+ assert_eq!(a.current(), 4);
	223	+ let d = a.decide(&metrics(99.0, 2000.0, 100));
	224	+ assert_eq!(d, ScaleDecision::Hold(4));
	225	+ }
	226	+
	227	+ #[test]
	228	+ fn never_below_min() {
	229	+ let mut a = Autoscaler::new(cfg());
	230	+ // Underused from the start → stays at min.
	231	+ for _ in 0..6 {
	232	+ let _ = a.decide(&metrics(5.0, 50.0, 0));
	233	+ }
	234	+ assert_eq!(a.current(), 1);
	235	+ let d = a.decide(&metrics(5.0, 50.0, 0));
	236	+ assert_eq!(d, ScaleDecision::Hold(1));
	237	+ }
	238	+
	239	+ #[test]
	240	+ fn cooldown_prevents_flapping() {
	241	+ let mut a = Autoscaler::new(cfg());
	242	+ assert_eq!(
	243	+ a.decide(&metrics(95.0, 100.0, 0)),
	244	+ ScaleDecision::ScaleUp(2)
	245	+ );
	246	+ // Immediately after scaling up we are in cooldown; metrics flipping to
	247	+ // healthy must NOT trigger an instant scale-down.
	248	+ assert_eq!(a.decide(&metrics(5.0, 50.0, 0)), ScaleDecision::Hold(2));
	249	+ assert_eq!(a.decide(&metrics(5.0, 50.0, 0)), ScaleDecision::Hold(2));
	250	+ // After cooldown expires, healthy metrics can scale down.
	251	+ let d = a.decide(&metrics(5.0, 50.0, 0));
	252	+ assert_eq!(d, ScaleDecision::ScaleDown(1));
	253	+ }
	254	+}
0	255

▼crates/justapi-inference/src/control_plane.rs

+458-0

@@ -0,0 +1,458 @@
	1	+//! LLM Control Plane — model registry, versioning, and LoRA-aware resolution.
	2	+//!
	3	+//! Pure-data, always-compiled layer (no weights, no GPU). This is the
	4	+//! orchestration "brain" (Ray Serve / Dynamo / AIBrix equivalent) that decides,
	5	+//! for a given request, which concrete model version + adapter to serve and
	6	+//! with what runtime profile. The actual weight loading is performed by the
	7	+//! [`crate::Engine`] using the resolved [`WeightLocation`].
	8	+//!
	9	+//! Multi-replica supervision, KV-aware routing, and autoscaling build on top of
	10	+//! this registry (tracked for later in Phase 46).
	11	+
	12	+use std::collections::HashMap;
	13	+use std::path::PathBuf;
	14	+use std::sync::{Arc, RwLock};
	15	+
	16	+use crate::engine::EngineDevice;
	17	+use crate::real::quant::QuantMethod;
	18	+
	19	+/// Where a model version's weights physically live.
	20	+#[derive(Debug, Clone, PartialEq, Eq)]
	21	+pub enum WeightLocation {
	22	+ /// Local filesystem directory or file.
	23	+ Local(PathBuf),
	24	+ /// Object storage (S3 / GCS / MinIO compatible).
	25	+ S3 { bucket: String, key: String },
	26	+ /// HuggingFace Hub repo (optional revision / branch).
	27	+ HuggingFace {
	28	+ repo: String,
	29	+ revision: Option<String>,
	30	+ },
	31	+ /// A remote model registry served over HTTP(S).
	32	+ RemoteRegistry { url: String },
	33	+}
	34	+
	35	+/// How a model version should be served.
	36	+#[derive(Debug, Clone)]
	37	+pub struct RuntimeProfile {
	38	+ /// Device the model is pinned to.
	39	+ pub device: EngineDevice,
	40	+ /// Max concurrent sequences that may share the model.
	41	+ pub max_concurrency: usize,
	42	+ /// Max total tokens in a single scheduler batch (prefill + decode).
	43	+ pub max_batch_tokens: usize,
	44	+ /// Estimated GPU memory footprint in MiB (for placement / admission).
	45	+ pub gpu_memory_mib: usize,
	46	+ /// Quantization method this profile targets.
	47	+ pub quant_method: QuantMethod,
	48	+ /// Adapter names expected to be resident (for LoRA-dense serving).
	49	+ pub expected_adapters: Vec<String>,
	50	+}
	51	+
	52	+impl Default for RuntimeProfile {
	53	+ fn default() -> Self {
	54	+ Self {
	55	+ device: EngineDevice::Cpu,
	56	+ max_concurrency: 1,
	57	+ max_batch_tokens: 2048,
	58	+ gpu_memory_mib: 0,
	59	+ quant_method: QuantMethod::None,
	60	+ expected_adapters: Vec::new(),
	61	+ }
	62	+ }
	63	+}
	64	+
	65	+/// A single versioned model artifact.
	66	+#[derive(Debug, Clone)]
	67	+pub struct ModelVersion {
	68	+ /// Version identifier (e.g. "v1", "2024-01-01", or a content hash).
	69	+ pub version: String,
	70	+ /// Where the weights live.
	71	+ pub weight_location: WeightLocation,
	72	+ /// Runtime profile for this version.
	73	+ pub runtime_profile: RuntimeProfile,
	74	+ /// Creation time as Unix seconds (for `latest` resolution).
	75	+ pub created_at_unix: u64,
	76	+ /// Alternate names that also resolve to this version (e.g. "stable").
	77	+ pub aliases: Vec<String>,
	78	+}
	79	+
	80	+impl ModelVersion {
	81	+ /// Create a version with the default runtime profile.
	82	+ pub fn new(version: impl Into<String>, weight_location: WeightLocation) -> Self {
	83	+ Self {
	84	+ version: version.into(),
	85	+ weight_location,
	86	+ runtime_profile: RuntimeProfile::default(),
	87	+ created_at_unix: 0,
	88	+ aliases: Vec::new(),
	89	+ }
	90	+ }
	91	+
	92	+ /// Builder: attach a runtime profile.
	93	+ pub fn with_profile(mut self, profile: RuntimeProfile) -> Self {
	94	+ self.runtime_profile = profile;
	95	+ self
	96	+ }
	97	+
	98	+ /// Builder: set creation timestamp.
	99	+ pub fn with_created_at(mut self, unix: u64) -> Self {
	100	+ self.created_at_unix = unix;
	101	+ self
	102	+ }
	103	+
	104	+ /// Builder: attach aliases.
	105	+ pub fn with_aliases(mut self, aliases: Vec<String>) -> Self {
	106	+ self.aliases = aliases;
	107	+ self
	108	+ }
	109	+
	110	+ /// Whether `name` matches this version or one of its aliases.
	111	+ pub fn matches(&self, name: &str) -> bool {
	112	+ self.version == name \|\| self.aliases.iter().any(\|a\| a == name)
	113	+ }
	114	+}
	115	+
	116	+/// A registered model: a name with multiple versions + adapter routing.
	117	+#[derive(Debug, Clone)]
	118	+pub struct ModelRecord {
	119	+ /// Canonical model name (e.g. "llama-7b").
	120	+ pub name: String,
	121	+ /// All known versions of this model.
	122	+ pub versions: Vec<ModelVersion>,
	123	+ /// The version used when none is requested.
	124	+ pub default_version: String,
	125	+ /// LoRA routing key → adapter name (resolved at request time).
	126	+ pub adapter_routing: HashMap<String, String>,
	127	+}
	128	+
	129	+impl ModelRecord {
	130	+ /// Create an empty record (no versions yet).
	131	+ pub fn new(name: impl Into<String>) -> Self {
	132	+ Self {
	133	+ name: name.into(),
	134	+ versions: Vec::new(),
	135	+ default_version: String::new(),
	136	+ adapter_routing: HashMap::new(),
	137	+ }
	138	+ }
	139	+
	140	+ /// Add (or replace) a version.
	141	+ pub fn add_version(&mut self, version: ModelVersion) {
	142	+ if let Some(slot) = self
	143	+ .versions
	144	+ .iter_mut()
	145	+ .find(\|v\| v.version == version.version)
	146	+ {
	147	+ *slot = version;
	148	+ } else {
	149	+ self.versions.push(version);
	150	+ }
	151	+ if self.default_version.is_empty() {
	152	+ self.default_version = self.versions.last().unwrap().version.clone();
	153	+ }
	154	+ }
	155	+
	156	+ /// Set the default version (must already exist).
	157	+ pub fn set_default(&mut self, version: &str) -> bool {
	158	+ if self.versions.iter().any(\|v\| v.matches(version)) {
	159	+ self.default_version = version.to_string();
	160	+ true
	161	+ } else {
	162	+ false
	163	+ }
	164	+ }
	165	+
	166	+ /// Map a routing key to an adapter name for this model.
	167	+ pub fn route_adapter(&mut self, routing_key: impl Into<String>, adapter: impl Into<String>) {
	168	+ self.adapter_routing
	169	+ .insert(routing_key.into(), adapter.into());
	170	+ }
	171	+
	172	+ /// Resolve a version-or-alias string to a concrete version.
	173	+ pub fn resolve_version(&self, version_or_alias: &str) -> Option<&ModelVersion> {
	174	+ self.versions.iter().find(\|v\| v.matches(version_or_alias))
	175	+ }
	176	+
	177	+ /// The newest version by `created_at_unix` (ties broken by insertion order).
	178	+ pub fn latest_version(&self) -> Option<&ModelVersion> {
	179	+ self.versions.iter().max_by_key(\|v\| v.created_at_unix)
	180	+ }
	181	+
	182	+ /// Resolve a routing key to an adapter name.
	183	+ pub fn resolve_adapter(&self, routing_key: &str) -> Option<&str> {
	184	+ self.adapter_routing.get(routing_key).map(\|s\| s.as_str())
	185	+ }
	186	+}
	187	+
	188	+/// A fully resolved serving target produced by [`ControlPlane::resolve`].
	189	+#[derive(Debug, Clone)]
	190	+pub struct ResolvedModel {
	191	+ pub model_name: String,
	192	+ pub version: String,
	193	+ pub weight_location: WeightLocation,
	194	+ pub runtime_profile: RuntimeProfile,
	195	+ pub adapter: Option<String>,
	196	+}
	197	+
	198	+/// Thread-safe, multi-tenant model registry.
	199	+#[derive(Clone)]
	200	+pub struct ControlPlane {
	201	+ inner: Arc<RwLock<HashMap<String, ModelRecord>>>,
	202	+}
	203	+
	204	+impl ControlPlane {
	205	+ /// Create an empty control plane.
	206	+ pub fn new() -> Self {
	207	+ Self {
	208	+ inner: Arc::new(RwLock::new(HashMap::new())),
	209	+ }
	210	+ }
	211	+
	212	+ /// Register (or replace) a whole model record.
	213	+ pub fn register(&self, record: ModelRecord) {
	214	+ self.inner
	215	+ .write()
	216	+ .unwrap()
	217	+ .insert(record.name.clone(), record);
	218	+ }
	219	+
	220	+ /// Register a single version under an existing (or new) model.
	221	+ pub fn register_version(&self, model: &str, version: ModelVersion) {
	222	+ let mut inner = self.inner.write().unwrap();
	223	+ let record = inner
	224	+ .entry(model.to_string())
	225	+ .or_insert_with(\|\| ModelRecord::new(model));
	226	+ record.add_version(version);
	227	+ }
	228	+
	229	+ /// Fetch a model record by name.
	230	+ pub fn get(&self, name: &str) -> Option<ModelRecord> {
	231	+ self.inner.read().unwrap().get(name).cloned()
	232	+ }
	233	+
	234	+ /// List all registered model names.
	235	+ pub fn list(&self) -> Vec<String> {
	236	+ self.inner.read().unwrap().keys().cloned().collect()
	237	+ }
	238	+
	239	+ /// Number of registered models.
	240	+ pub fn len(&self) -> usize {
	241	+ self.inner.read().unwrap().len()
	242	+ }
	243	+
	244	+ /// Whether the registry is empty.
	245	+ pub fn is_empty(&self) -> bool {
	246	+ self.inner.read().unwrap().is_empty()
	247	+ }
	248	+
	249	+ /// Unregister a model, returning it if present.
	250	+ pub fn unregister(&self, name: &str) -> Option<ModelRecord> {
	251	+ self.inner.write().unwrap().remove(name)
	252	+ }
	253	+
	254	+ /// Resolve a model + optional version/alias + optional LoRA routing key into
	255	+ /// a concrete serving target.
	256	+ ///
	257	+ /// Resolution order for the version:
	258	+ /// - If `version_or_alias` is given, it must match a known version or
	259	+ /// alias (strict — a wrong explicit request fails rather than silently
	260	+ /// serving a different version).
	261	+ /// - If omitted, fall back to the model's `default_version`, then to the
	262	+ /// latest version by `created_at_unix`.
	263	+ ///
	264	+ /// The adapter is resolved from the routing key against the model's
	265	+ /// `adapter_routing` map (LoRA-aware routing).
	266	+ pub fn resolve(
	267	+ &self,
	268	+ model: &str,
	269	+ version_or_alias: Option<&str>,
	270	+ routing_key: Option<&str>,
	271	+ ) -> Option<ResolvedModel> {
	272	+ let inner = self.inner.read().unwrap();
	273	+ let record = inner.get(model)?;
	274	+
	275	+ let version = match version_or_alias {
	276	+ Some(v) => record.resolve_version(v)?, // strict: explicit miss → None
	277	+ None => record
	278	+ .resolve_version(&record.default_version)
	279	+ .or_else(\|\| record.latest_version())?,
	280	+ };
	281	+
	282	+ let adapter = routing_key
	283	+ .and_then(\|k\| record.resolve_adapter(k))
	284	+ .map(String::from);
	285	+
	286	+ Some(ResolvedModel {
	287	+ model_name: record.name.clone(),
	288	+ version: version.version.clone(),
	289	+ weight_location: version.weight_location.clone(),
	290	+ runtime_profile: version.runtime_profile.clone(),
	291	+ adapter,
	292	+ })
	293	+ }
	294	+}
	295	+
	296	+impl Default for ControlPlane {
	297	+ fn default() -> Self {
	298	+ Self::new()
	299	+ }
	300	+}
	301	+
	302	+// ---------------------------------------------------------------------------
	303	+// Tests
	304	+// ---------------------------------------------------------------------------
	305	+
	306	+#[cfg(test)]
	307	+mod tests {
	308	+ use super::*;
	309	+
	310	+ fn version(v: &str, at: u64) -> ModelVersion {
	311	+ ModelVersion::new(
	312	+ v,
	313	+ WeightLocation::Local(PathBuf::from(format!("/models/{v}"))),
	314	+ )
	315	+ .with_created_at(at)
	316	+ }
	317	+
	318	+ #[test]
	319	+ fn register_and_get_model() {
	320	+ let cp = ControlPlane::new();
	321	+ assert!(cp.is_empty());
	322	+ cp.register_version("llama-7b", version("v1", 100));
	323	+ cp.register_version("llama-7b", version("v2", 200));
	324	+ assert_eq!(cp.len(), 1);
	325	+ assert_eq!(cp.list(), vec!["llama-7b"]);
	326	+ assert!(cp.get("llama-7b").is_some());
	327	+ }
	328	+
	329	+ #[test]
	330	+ fn default_version_is_first_registered() {
	331	+ let cp = ControlPlane::new();
	332	+ cp.register_version("m", version("v1", 100));
	333	+ cp.register_version("m", version("v2", 200));
	334	+ let r = cp.get("m").unwrap();
	335	+ assert_eq!(r.default_version, "v1");
	336	+ }
	337	+
	338	+ #[test]
	339	+ fn set_default_version() {
	340	+ let cp = ControlPlane::new();
	341	+ cp.register_version("m", version("v1", 100));
	342	+ cp.register_version("m", version("v2", 200));
	343	+ let mut r = cp.get("m").unwrap();
	344	+ assert!(r.set_default("v2"));
	345	+ assert!(!r.set_default("nope"));
	346	+ }
	347	+
	348	+ #[test]
	349	+ fn resolve_by_version_then_alias_then_default_then_latest() {
	350	+ let cp = ControlPlane::new();
	351	+ cp.register_version("m", version("v1", 100).with_aliases(vec!["stable".into()]));
	352	+ cp.register_version("m", version("v2", 200));
	353	+
	354	+ // explicit version
	355	+ let r = cp.resolve("m", Some("v2"), None).unwrap();
	356	+ assert_eq!(r.version, "v2");
	357	+
	358	+ // alias
	359	+ let r = cp.resolve("m", Some("stable"), None).unwrap();
	360	+ assert_eq!(r.version, "v1");
	361	+
	362	+ // default (v1, since it was registered first)
	363	+ let r = cp.resolve("m", None, None).unwrap();
	364	+ assert_eq!(r.version, "v1");
	365	+
	366	+ // missing model → None
	367	+ assert!(cp.resolve("missing", None, None).is_none());
	368	+ }
	369	+
	370	+ #[test]
	371	+ fn resolve_explicit_miss_is_strict_none() {
	372	+ let cp = ControlPlane::new();
	373	+ cp.register_version("m", version("v1", 100));
	374	+ cp.register_version("m", version("v2", 200));
	375	+ // A wrong explicit version must fail loudly, not silently fall back.
	376	+ assert!(cp.resolve("m", Some("ghost"), None).is_none());
	377	+ }
	378	+
	379	+ #[test]
	380	+ fn resolve_no_version_falls_back_to_default_then_latest() {
	381	+ let cp = ControlPlane::new();
	382	+ // Build a record whose default_version is empty so resolution must fall
	383	+ // through to the latest version by created_at.
	384	+ let mut record = ModelRecord::new("m");
	385	+ record.add_version(version("v1", 100));
	386	+ record.add_version(version("v3", 200));
	387	+ record.add_version(version("v2", 300));
	388	+ record.default_version = String::new();
	389	+ cp.register(record);
	390	+
	391	+ let r = cp.resolve("m", None, None).unwrap();
	392	+ assert_eq!(r.version, "v2"); // latest by created_at
	393	+ }
	394	+
	395	+ #[test]
	396	+ fn lora_aware_routing_resolves_adapter() {
	397	+ let cp = ControlPlane::new();
	398	+ cp.register_version("m", version("v1", 100));
	399	+ let mut record = cp.get("m").unwrap();
	400	+ record.route_adapter("domain:math", "math-adapter");
	401	+ record.route_adapter("domain:code", "code-adapter");
	402	+ cp.register(record);
	403	+
	404	+ // routing key matches → adapter resolved
	405	+ let r = cp.resolve("m", None, Some("domain:math")).unwrap();
	406	+ assert_eq!(r.adapter.as_deref(), Some("math-adapter"));
	407	+
	408	+ // routing key matches → different adapter
	409	+ let r = cp.resolve("m", None, Some("domain:code")).unwrap();
	410	+ assert_eq!(r.adapter.as_deref(), Some("code-adapter"));
	411	+
	412	+ // no routing key → no adapter
	413	+ let r = cp.resolve("m", None, None).unwrap();
	414	+ assert!(r.adapter.is_none());
	415	+
	416	+ // unknown routing key → no adapter
	417	+ let r = cp.resolve("m", None, Some("domain:unknown")).unwrap();
	418	+ assert!(r.adapter.is_none());
	419	+ }
	420	+
	421	+ #[test]
	422	+ fn resolve_carries_runtime_profile_and_location() {
	423	+ let cp = ControlPlane::new();
	424	+ let profile = RuntimeProfile {
	425	+ device: EngineDevice::Cuda(0),
	426	+ max_concurrency: 8,
	427	+ max_batch_tokens: 4096,
	428	+ gpu_memory_mib: 14_000,
	429	+ quant_method: QuantMethod::Gguf,
	430	+ expected_adapters: vec!["math-adapter".into()],
	431	+ };
	432	+ cp.register_version(
	433	+ "m",
	434	+ version("v1", 100)
	435	+ .with_profile(profile.clone())
	436	+ .with_aliases(vec!["latest".into()]),
	437	+ );
	438	+ let r = cp.resolve("m", Some("latest"), None).unwrap();
	439	+ assert_eq!(r.runtime_profile.device, EngineDevice::Cuda(0));
	440	+ assert_eq!(r.runtime_profile.max_concurrency, 8);
	441	+ assert_eq!(r.runtime_profile.quant_method, QuantMethod::Gguf);
	442	+ assert_eq!(
	443	+ r.weight_location,
	444	+ WeightLocation::Local(PathBuf::from("/models/v1"))
	445	+ );
	446	+ assert_eq!(r.runtime_profile.expected_adapters, vec!["math-adapter"]);
	447	+ }
	448	+
	449	+ #[test]
	450	+ fn unregister_removes_model() {
	451	+ let cp = ControlPlane::new();
	452	+ cp.register_version("m", version("v1", 100));
	453	+ let removed = cp.unregister("m");
	454	+ assert!(removed.is_some());
	455	+ assert!(cp.is_empty());
	456	+ assert!(cp.unregister("m").is_none());
	457	+ }
	458	+}
0	459

▼crates/justapi-inference/src/engine.rs

+257-0

@@ -0,0 +1,257 @@
	1	+//! The streaming [`Engine`]: device selection, an in-memory [`ModelRegistry`],
	2	+//! and token streaming via a tokio channel. Generation runs on a dedicated OS
	3	+//! thread so the hot path stays GIL-free.
	4	+
	5	+use std::collections::HashMap;
	6	+use std::path::Path;
	7	+use std::sync::{Arc, Mutex};
	8	+
	9	+use candle_core::Device;
	10	+use tokio::sync::mpsc;
	11	+
	12	+use crate::model::{GeneratedToken, MockModel, Model, ModelError, SamplingParams};
	13	+use crate::spec_decode::{SpeculativeConfig, SpeculativeModel};
	14	+use crate::spec_decode_tree::TreeSpeculativeModel;
	15	+
	16	+/// Compute device a model runs on. Mirrors Candle's [`Device`] minus the ROCm
	17	+/// variant (candle 0.11 enables AMD GPUs via the `cuda` feature on a ROCm
	18	+/// toolchain).
	19	+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	20	+pub enum EngineDevice {
	21	+ Cpu,
	22	+ Cuda(usize),
	23	+ Metal(usize),
	24	+}
	25	+
	26	+impl EngineDevice {
	27	+ /// Convert to a Candle [`Device`]. GPU variants are constructed via
	28	+ /// `new_cuda`/`new_metal` (which validate the ordinal exists).
	29	+ pub fn to_candle(self) -> candle_core::Result<Device> {
	30	+ match self {
	31	+ EngineDevice::Cpu => Ok(Device::Cpu),
	32	+ EngineDevice::Cuda(i) => Device::new_cuda(i),
	33	+ EngineDevice::Metal(i) => Device::new_metal(i),
	34	+ }
	35	+ }
	36	+
	37	+ /// Best-effort discovery of available devices.
	38	+ ///
	39	+ /// CPU is always present. GPUs are probed by attempting to construct the
	40	+ /// device (cheap on a properly configured host; errors are ignored).
	41	+ pub fn discover() -> Vec<EngineDevice> {
	42	+ let mut devices = vec![EngineDevice::Cpu];
	43	+ for i in 0..8 {
	44	+ if Device::new_cuda(i).is_ok() {
	45	+ devices.push(EngineDevice::Cuda(i));
	46	+ }
	47	+ }
	48	+ for i in 0..8 {
	49	+ if Device::new_metal(i).is_ok() {
	50	+ devices.push(EngineDevice::Metal(i));
	51	+ }
	52	+ }
	53	+ devices
	54	+ }
	55	+}
	56	+
	57	+/// Where a registered model came from.
	58	+#[derive(Debug, Clone)]
	59	+pub enum ModelSource {
	60	+ /// Weight-free model (tests, demos).
	61	+ Mock,
	62	+ /// Loaded from a directory of weights (path recorded).
	63	+ Real(std::path::PathBuf),
	64	+}
	65	+
	66	+struct RegistryEntry {
	67	+ model: Arc<dyn Model>,
	68	+ source: ModelSource,
	69	+}
	70	+
	71	+/// In-memory collection of loaded models, keyed by name.
	72	+#[derive(Default)]
	73	+pub struct ModelRegistry {
	74	+ models: HashMap<String, RegistryEntry>,
	75	+}
	76	+
	77	+impl ModelRegistry {
	78	+ /// Register a model under `name`, replacing any existing entry.
	79	+ pub fn register(&mut self, name: &str, model: Arc<dyn Model>, source: ModelSource) {
	80	+ self.models
	81	+ .insert(name.to_string(), RegistryEntry { model, source });
	82	+ }
	83	+
	84	+ /// Fetch a shared handle to a registered model.
	85	+ pub fn get(&self, name: &str) -> Option<Arc<dyn Model>> {
	86	+ self.models.get(name).map(\|e\| e.model.clone())
	87	+ }
	88	+
	89	+ /// List registered model names.
	90	+ pub fn list(&self) -> Vec<String> {
	91	+ self.models.keys().cloned().collect()
	92	+ }
	93	+
	94	+ /// Number of registered models.
	95	+ pub fn len(&self) -> usize {
	96	+ self.models.len()
	97	+ }
	98	+
	99	+ /// Whether the registry is empty.
	100	+ pub fn is_empty(&self) -> bool {
	101	+ self.models.is_empty()
	102	+ }
	103	+}
	104	+
	105	+/// The inference engine: owns the device and the model registry, and streams
	106	+/// generation output back to callers over a channel.
	107	+pub struct Engine {
	108	+ device: Device,
	109	+ registry: Arc<Mutex<ModelRegistry>>,
	110	+}
	111	+
	112	+impl Engine {
	113	+ /// Create an engine bound to `device`.
	114	+ pub fn new(device: EngineDevice) -> Result<Self, ModelError> {
	115	+ let device = device
	116	+ .to_candle()
	117	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	118	+ Ok(Self {
	119	+ device,
	120	+ registry: Arc::new(Mutex::new(ModelRegistry::default())),
	121	+ })
	122	+ }
	123	+
	124	+ /// The Candle device this engine runs models on.
	125	+ pub fn device(&self) -> &Device {
	126	+ &self.device
	127	+ }
	128	+
	129	+ /// Register an arbitrary model implementation.
	130	+ pub fn register(&self, name: &str, model: Arc<dyn Model>) {
	131	+ self.registry
	132	+ .lock()
	133	+ .unwrap()
	134	+ .register(name, model, ModelSource::Mock);
	135	+ }
	136	+
	137	+ /// Fetch a shared handle to a registered model by name.
	138	+ pub fn get(&self, name: &str) -> Option<Arc<dyn Model>> {
	139	+ self.registry.lock().unwrap().get(name)
	140	+ }
	141	+
	142	+ /// Register a [`MockModel`] and return a handle to it.
	143	+ pub fn register_mock(&self, name: &str) -> Arc<MockModel> {
	144	+ let model = Arc::new(MockModel::new(crate::DEFAULT_MOCK_VOCAB));
	145	+ self.registry
	146	+ .lock()
	147	+ .unwrap()
	148	+ .register(name, model.clone(), ModelSource::Mock);
	149	+ model
	150	+ }
	151	+
	152	+ /// Register a speculative-decode model: a `target` model served with a
	153	+ /// `draft` model proposing `gamma` candidate tokens per step. The wrapped
	154	+ /// [`SpeculativeModel`] is registered under `name` and served through the
	155	+ /// normal [`Engine::generate`] path transparently.
	156	+ pub fn register_speculative(
	157	+ &self,
	158	+ name: &str,
	159	+ target: Arc<dyn Model>,
	160	+ draft: Arc<dyn Model>,
	161	+ gamma: usize,
	162	+ seed: u64,
	163	+ ) {
	164	+ let spec = SpeculativeModel::new(target, SpeculativeConfig::new(gamma, draft, seed));
	165	+ self.registry
	166	+ .lock()
	167	+ .unwrap()
	168	+ .register(name, Arc::new(spec), ModelSource::Mock);
	169	+ }
	170	+
	171	+ /// Register a tree-based speculative-decode model (Medusa/EAGLE-style): a
	172	+ /// `draft` model proposes a tree of `branch` candidates at each of `gamma`
	173	+ /// positions, and the `target` verifies the longest matching path. The
	174	+ /// wrapped [`TreeSpeculativeModel`] is registered under `name` and served
	175	+ /// through the normal [`Engine::generate`] path transparently — callers get
	176	+ /// the same token stream as plain target decode, just faster.
	177	+ pub fn register_tree_speculative(
	178	+ &self,
	179	+ name: &str,
	180	+ target: Arc<dyn Model>,
	181	+ draft: Arc<dyn Model>,
	182	+ gamma: usize,
	183	+ branch: usize,
	184	+ seed: u64,
	185	+ ) {
	186	+ let spec = TreeSpeculativeModel::new(target, draft, gamma, branch, seed);
	187	+ self.registry
	188	+ .lock()
	189	+ .unwrap()
	190	+ .register(name, Arc::new(spec), ModelSource::Mock);
	191	+ }
	192	+
	193	+ /// Load a real model from `model_dir`.
	194	+ ///
	195	+ /// Behind the `real` feature (and with weights + a GPU toolkit) this will
	196	+ /// build the candle-transformers forward pass + KV cache (Phase 42). Until
	197	+ /// then it returns a clear error so callers don't silently get a no-op.
	198	+ pub fn load(&self, name: &str, model_dir: &Path) -> Result<(), ModelError> {
	199	+ #[cfg(feature = "real")]
	200	+ {
	201	+ use crate::real::model::RealModel;
	202	+ let model = RealModel::load(model_dir, self.device.clone())
	203	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	204	+ self.registry.lock().unwrap().register(
	205	+ name,
	206	+ Arc::new(model),
	207	+ ModelSource::Real(model_dir.to_path_buf()),
	208	+ );
	209	+ Ok(())
	210	+ }
	211	+ #[cfg(not(feature = "real"))]
	212	+ {
	213	+ let _ = (name, model_dir);
	214	+ Err(ModelError::FeatureRequired("real"))
	215	+ }
	216	+ }
	217	+
	218	+ /// Generate tokens for `prompt` with `params`, streaming them to the
	219	+ /// returned receiver as they are produced.
	220	+ ///
	221	+ /// Generation runs on a dedicated OS thread (Candle compute is synchronous);
	222	+ /// the caller receives tokens asynchronously with zero GIL involvement.
	223	+ pub fn generate(
	224	+ &self,
	225	+ name: &str,
	226	+ prompt: &[u32],
	227	+ params: SamplingParams,
	228	+ ) -> Result<mpsc::UnboundedReceiver<GeneratedToken>, ModelError> {
	229	+ let model = {
	230	+ let reg = self.registry.lock().unwrap();
	231	+ reg.get(name)
	232	+ .ok_or_else(\|\| ModelError::NotFound(name.to_string()))?
	233	+ };
	234	+
	235	+ let (tx, rx) = mpsc::unbounded_channel();
	236	+ let prompt = prompt.to_vec();
	237	+ std::thread::spawn(move \|\| {
	238	+ let _ = model.generate(&prompt, &params, &\|token\| tx.send(token).is_ok());
	239	+ });
	240	+ Ok(rx)
	241	+ }
	242	+
	243	+ /// List registered model names.
	244	+ pub fn list_models(&self) -> Vec<String> {
	245	+ self.registry.lock().unwrap().list()
	246	+ }
	247	+
	248	+ /// Return the [`ModelSource`] a registered model was loaded from, if known.
	249	+ pub fn model_source(&self, name: &str) -> Option<ModelSource> {
	250	+ self.registry
	251	+ .lock()
	252	+ .unwrap()
	253	+ .models
	254	+ .get(name)
	255	+ .map(\|e\| e.source.clone())
	256	+ }
	257	+}
0	258

▼crates/justapi-inference/src/gateway.rs

+302-0

@@ -0,0 +1,302 @@
	1	+//! K8s AI inference gateway — KV-aware + LoRA-aware routing inside the gateway.
	2	+//!
	3	+//! Composes the control-plane ([`ControlPlane`]) and the request router
	4	+//! ([`Router`]) into a deployable gateway unit: incoming inference requests are
	5	+//! resolved to a concrete model version + adapter, then steered to the best
	6	+//! serving replica. The gateway adds the Kubernetes-specific concerns the raw
	7	+//! router does not own:
	8	+//!
	9	+//! - KV-aware default strategy — the gateway routes by *lowest KV-cache
	10	+//! pressure* by default, so hot replicas (large prefix caches, long in-flight
	11	+//! decodes) are not overloaded. This is the "KV-aware routing inside the
	12	+//! gateway" property (Dynamo / AIBrix behaviour).
	13	+//! - Readiness gating — a pod must be K8s-Ready to receive traffic; a
	14	+//! not-Ready pod is excluded even if its liveness probe is fine. Readiness is
	15	+//! reflected into the router's `healthy` flag (standard K8s mapping:
	16	+//! not-Ready ⇒ not a routing candidate).
	17	+//! - Endpoint resolution — the routed replica is turned into a Kubernetes
	18	+//! service DNS name (`<replica>.<namespace>.svc.cluster.local` by default) the
	19	+//! gateway forwards the request to. The template is configurable and supports
	20	+//! `{replica}`, `{model}`, `{version}`, `{namespace}`.
	21	+//!
	22	+//! This is pure decision logic (no network, no GPU) — the same structural
	23	+//! approach as the rest of the control plane. A K8s executor fulfils the
	24	+//! emitted routing decisions and reports pod readiness/load back via
	25	+//! [`InferenceGateway::set_ready`] / [`InferenceGateway::update_load`].
	26	+
	27	+use crate::control_plane::{ControlPlane, ResolvedModel};
	28	+use crate::router::{Replica, RouteDecision, RouteRequest, Router, RoutingStrategy};
	29	+
	30	+/// Configuration for the K8s inference gateway.
	31	+#[derive(Debug, Clone)]
	32	+pub struct GatewayConfig {
	33	+ /// Kubernetes namespace the gateway serves (used in endpoint resolution and
	34	+ /// for isolating per-namespace routing).
	35	+ pub namespace: String,
	36	+ /// Service-DNS template. Supports `{replica}`, `{model}`, `{version}`,
	37	+ /// `{namespace}`. Defaults to `"{replica}.{namespace}.svc.cluster.local"`.
	38	+ pub service_template: String,
	39	+ /// Replica selection strategy. Defaults to [`RoutingStrategy::LowestKvPressure`]
	40	+ /// (the gateway's KV-aware default).
	41	+ pub strategy: RoutingStrategy,
	42	+ /// Whether K8s readiness gates routing (not-Ready pods excluded). Default true.
	43	+ pub require_ready: bool,
	44	+}
	45	+
	46	+impl Default for GatewayConfig {
	47	+ fn default() -> Self {
	48	+ Self {
	49	+ namespace: "justapi".to_string(),
	50	+ service_template: "{replica}.{namespace}.svc.cluster.local".to_string(),
	51	+ strategy: RoutingStrategy::LowestKvPressure,
	52	+ require_ready: true,
	53	+ }
	54	+ }
	55	+}
	56	+
	57	+/// The outcome of gateway routing: which replica + its Kubernetes endpoint.
	58	+#[derive(Debug, Clone)]
	59	+pub struct GatewayDecision {
	60	+ /// Chosen replica id (maps to a pod/deployment in the cluster).
	61	+ pub replica: String,
	62	+ /// Fully resolved model target (version + adapter).
	63	+ pub resolved: ResolvedModel,
	64	+ /// Kubernetes service endpoint the gateway should forward the request to.
	65	+ pub endpoint: String,
	66	+}
	67	+
	68	+/// The K8s inference gateway: control-plane resolution + KV-aware/LoRA-aware
	69	+/// routing with readiness gating and endpoint resolution.
	70	+pub struct InferenceGateway {
	71	+ cp: ControlPlane,
	72	+ router: Router,
	73	+ config: GatewayConfig,
	74	+}
	75	+
	76	+impl InferenceGateway {
	77	+ /// Create a gateway bound to `config`.
	78	+ pub fn new(config: GatewayConfig) -> Self {
	79	+ let cp = ControlPlane::new();
	80	+ let router = Router::new(config.strategy);
	81	+ Self { cp, router, config }
	82	+ }
	83	+
	84	+ /// Borrow the embedded control plane (to register model versions/adapters).
	85	+ pub fn control_plane(&self) -> &ControlPlane {
	86	+ &self.cp
	87	+ }
	88	+
	89	+ /// Borrow the embedded router (for advanced inspection).
	90	+ pub fn router(&self) -> &Router {
	91	+ &self.router
	92	+ }
	93	+
	94	+ /// Register (or replace) a serving replica with the gateway.
	95	+ pub fn register_replica(&self, replica: Replica) {
	96	+ self.router.register_replica(replica);
	97	+ }
	98	+
	99	+ /// Report a replica's live load back from the cluster.
	100	+ pub fn update_load(
	101	+ &self,
	102	+ id: &str,
	103	+ active_sequences: usize,
	104	+ kv_pressure_pct: f32,
	105	+ ready: bool,
	106	+ ) {
	107	+ // Readiness is folded into `healthy` so the router's candidate filter
	108	+ // enforces it; a not-Ready pod is not a routing candidate.
	109	+ let healthy = if self.config.require_ready {
	110	+ ready
	111	+ } else {
	112	+ true
	113	+ };
	114	+ self.router
	115	+ .update_load(id, active_sequences, kv_pressure_pct, healthy);
	116	+ }
	117	+
	118	+ /// K8s readiness-probe callback: mark a pod Ready/NotReady for traffic.
	119	+ pub fn set_ready(&self, id: &str, ready: bool) {
	120	+ if let Some(r) = self.router.replicas().into_iter().find(\|r\| r.id == id) {
	121	+ self.update_load(id, r.active_sequences, r.kv_pressure_pct, ready);
	122	+ }
	123	+ }
	124	+
	125	+ /// Remove a replica (pod deleted / drained).
	126	+ pub fn remove_replica(&self, id: &str) {
	127	+ self.router.remove_replica(id);
	128	+ }
	129	+
	130	+ /// Route an inference request to the best serving replica.
	131	+ ///
	132	+ /// Resolves the model via the control plane, selects a replica with the
	133	+ /// configured (KV-aware / LoRA-aware) strategy, and returns the replica plus
	134	+ /// its Kubernetes endpoint. Returns `None` if the model is unknown or no
	135	+ /// replica can serve it (no capacity / not Ready).
	136	+ pub fn route(&self, request: &RouteRequest) -> Option<GatewayDecision> {
	137	+ let decision: RouteDecision = self.router.route(&self.cp, request)?;
	138	+ let endpoint = self.endpoint_for(&decision.replica, &decision.resolved);
	139	+ Some(GatewayDecision {
	140	+ replica: decision.replica,
	141	+ resolved: decision.resolved,
	142	+ endpoint,
	143	+ })
	144	+ }
	145	+
	146	+ /// Resolve a replica id + resolved model into a Kubernetes service endpoint.
	147	+ fn endpoint_for(&self, replica: &str, resolved: &ResolvedModel) -> String {
	148	+ self.config
	149	+ .service_template
	150	+ .replace("{replica}", replica)
	151	+ .replace("{model}", &resolved.model_name)
	152	+ .replace("{version}", &resolved.version)
	153	+ .replace("{namespace}", &self.config.namespace)
	154	+ }
	155	+}
	156	+
	157	+// ---------------------------------------------------------------------------
	158	+// Tests
	159	+// ---------------------------------------------------------------------------
	160	+
	161	+#[cfg(test)]
	162	+mod tests {
	163	+ use super::*;
	164	+ use crate::control_plane::{ModelVersion, WeightLocation};
	165	+ use crate::engine::EngineDevice;
	166	+ use std::path::PathBuf;
	167	+
	168	+ fn gateway() -> InferenceGateway {
	169	+ let g = InferenceGateway::new(GatewayConfig {
	170	+ namespace: "llm-prod".to_string(),
	171	+ ..Default::default()
	172	+ });
	173	+ // Register a model with two adapter routes.
	174	+ g.control_plane().register_version(
	175	+ "llama-7b",
	176	+ ModelVersion::new("v1", WeightLocation::Local(PathBuf::from("/m/v1"))),
	177	+ );
	178	+ let mut record = g.control_plane().get("llama-7b").unwrap();
	179	+ record.route_adapter("domain:math", "math-adapter");
	180	+ record.route_adapter("domain:code", "code-adapter");
	181	+ g.control_plane().register(record);
	182	+ g
	183	+ }
	184	+
	185	+ fn replica(id: &str, active: usize, kv: f32, adapters: &[&str]) -> Replica {
	186	+ Replica {
	187	+ id: id.to_string(),
	188	+ model: "llama-7b".to_string(),
	189	+ version: "v1".to_string(),
	190	+ device: EngineDevice::Cuda(0),
	191	+ max_concurrency: 4,
	192	+ active_sequences: active,
	193	+ kv_pressure_pct: kv,
	194	+ healthy: true,
	195	+ loaded_adapters: adapters.iter().map(\|s\| s.to_string()).collect(),
	196	+ }
	197	+ }
	198	+
	199	+ #[test]
	200	+ fn route_resolves_kubernetes_endpoint() {
	201	+ let g = gateway();
	202	+ g.register_replica(replica("llama-7b-v1-0", 0, 10.0, &[]));
	203	+ let d = g
	204	+ .route(&RouteRequest {
	205	+ model: "llama-7b".into(),
	206	+ version: None,
	207	+ routing_key: None,
	208	+ })
	209	+ .unwrap();
	210	+ assert_eq!(d.replica, "llama-7b-v1-0");
	211	+ assert_eq!(d.endpoint, "llama-7b-v1-0.llm-prod.svc.cluster.local");
	212	+ assert_eq!(d.resolved.version, "v1");
	213	+ }
	214	+
	215	+ #[test]
	216	+ fn namespace_appears_in_endpoint() {
	217	+ let g = InferenceGateway::new(GatewayConfig {
	218	+ namespace: "team-a".to_string(),
	219	+ service_template: "justapi-{replica}.{namespace}.svc.cluster.local".to_string(),
	220	+ ..Default::default()
	221	+ });
	222	+ g.control_plane().register_version(
	223	+ "llama-7b",
	224	+ ModelVersion::new("v1", WeightLocation::Local(PathBuf::from("/m"))),
	225	+ );
	226	+ g.register_replica(replica("llama-7b-v1-0", 0, 10.0, &[]));
	227	+ let d = g
	228	+ .route(&RouteRequest {
	229	+ model: "llama-7b".into(),
	230	+ ..Default::default()
	231	+ })
	232	+ .unwrap();
	233	+ assert!(d.endpoint.contains("team-a"));
	234	+ assert!(d.endpoint.starts_with("justapi-llama-7b-v1-0"));
	235	+ }
	236	+
	237	+ #[test]
	238	+ fn readiness_gating_excludes_not_ready_pod() {
	239	+ let g = gateway();
	240	+ g.register_replica(replica("hot", 0, 10.0, &[]));
	241	+ g.register_replica(replica("cool", 0, 10.0, &[]));
	242	+ // Drain "hot" (not Ready) → traffic must go to "cool".
	243	+ g.set_ready("hot", false);
	244	+ let d = g
	245	+ .route(&RouteRequest {
	246	+ model: "llama-7b".into(),
	247	+ ..Default::default()
	248	+ })
	249	+ .unwrap();
	250	+ assert_eq!(d.replica, "cool");
	251	+
	252	+ // Drain both → no capacity.
	253	+ g.set_ready("cool", false);
	254	+ assert!(g
	255	+ .route(&RouteRequest {
	256	+ model: "llama-7b".into(),
	257	+ ..Default::default()
	258	+ })
	259	+ .is_none());
	260	+ }
	261	+
	262	+ #[test]
	263	+ fn kv_pressure_default_routes_to_coolest() {
	264	+ // Gateway default strategy is LowestKvPressure.
	265	+ let g = gateway();
	266	+ g.register_replica(replica("hot", 1, 90.0, &[]));
	267	+ g.register_replica(replica("cool", 1, 20.0, &[]));
	268	+ let d = g
	269	+ .route(&RouteRequest {
	270	+ model: "llama-7b".into(),
	271	+ ..Default::default()
	272	+ })
	273	+ .unwrap();
	274	+ assert_eq!(d.replica, "cool");
	275	+ }
	276	+
	277	+ #[test]
	278	+ fn lora_aware_routing_inside_gateway() {
	279	+ let g = gateway();
	280	+ // Only r2 holds the math-adapter → math routing key must land there.
	281	+ g.register_replica(replica("r1", 0, 10.0, &[]));
	282	+ g.register_replica(replica("r2", 0, 10.0, &["math-adapter"]));
	283	+ let d = g
	284	+ .route(&RouteRequest {
	285	+ model: "llama-7b".into(),
	286	+ routing_key: Some("domain:math".into()),
	287	+ ..Default::default()
	288	+ })
	289	+ .unwrap();
	290	+ assert_eq!(d.replica, "r2");
	291	+ assert_eq!(d.resolved.adapter.as_deref(), Some("math-adapter"));
	292	+
	293	+ // code-adapter held by nobody → no capacity.
	294	+ assert!(g
	295	+ .route(&RouteRequest {
	296	+ model: "llama-7b".into(),
	297	+ routing_key: Some("domain:code".into()),
	298	+ ..Default::default()
	299	+ })
	300	+ .is_none());
	301	+ }
	302	+}
0	303

▼crates/justapi-inference/src/kv_cache.rs

+773-0

@@ -0,0 +1,773 @@
	1	+//! Paged KV-cache manager for PagedAttention.
	2	+//!
	3	+//! Lean, pure-Rust data-structure layer — no Candle tensors, no GPU memory.
	4	+//! The model implementation (behind the `real` feature) maps [`BlockId`]s to
	5	+//! the actual tensor storage.
	6	+//!
	7	+//! ## Architecture
	8	+//!
	9	+//! - [`KvBlockPool`] — fixed-size block allocator with clock-sweep eviction.
	10	+//! - [`Sequence`] — a generation session holding a contiguous run of blocks.
	11	+//! - [`PrefixCache`] — hash-based prefix index for cross-request KV reuse.
	12	+//!
	13	+//! Blocks live in one of three states:
	14	+//!
	15	+//! 1. Free — available for allocation (in `free_list`).
	16	+//! 2. In-use — `ref_count > 0`, referenced by at least one [`Sequence`].
	17	+//! 3. Cached — `ref_count == 0` but still tracked by the [`PrefixCache`]
	18	+//! and reclaimable via eviction.
	19	+
	20	+use std::collections::HashMap;
	21	+
	22	+// ---------------------------------------------------------------------------
	23	+// Constants
	24	+// ---------------------------------------------------------------------------
	25	+
	26	+/// Number of token positions per KV-cache block (PagedAttention standard).
	27	+pub const BLOCK_SIZE: usize = 16;
	28	+
	29	+/// Maximum addressable blocks.
	30	+pub const MAX_BLOCKS: usize = 1 << 24;
	31	+
	32	+// ---------------------------------------------------------------------------
	33	+// Identifiers
	34	+// ---------------------------------------------------------------------------
	35	+
	36	+/// Opaque identifier for a single KV-cache block in the pool.
	37	+pub type BlockId = u32;
	38	+
	39	+// ---------------------------------------------------------------------------
	40	+// Statistics
	41	+// ---------------------------------------------------------------------------
	42	+
	43	+/// Snapshot of pool usage.
	44	+#[derive(Debug, Clone, Copy)]
	45	+pub struct PoolStats {
	46	+ pub total: usize,
	47	+ pub allocated: usize,
	48	+ pub cached: usize,
	49	+ pub free: usize,
	50	+ pub pressure_pct: f32,
	51	+ pub evictions: u64,
	52	+}
	53	+
	54	+/// Snapshot of prefix-cache effectiveness.
	55	+#[derive(Debug, Clone, Copy)]
	56	+pub struct PrefixCacheStats {
	57	+ pub entries: usize,
	58	+ pub hits: u64,
	59	+ pub misses: u64,
	60	+}
	61	+
	62	+// ---------------------------------------------------------------------------
	63	+// Per-block metadata
	64	+// ---------------------------------------------------------------------------
	65	+
	66	+#[derive(Debug, Clone)]
	67	+struct BlockMeta {
	68	+ ref_count: u32,
	69	+ pinned: bool,
	70	+ /// Token-content hash used by the prefix cache. Zero means no cached
	71	+ /// content (so the block is either free or in-use with uncacheable data).
	72	+ hash: u64,
	73	+}
	74	+
	75	+impl BlockMeta {
	76	+ const fn new() -> Self {
	77	+ Self {
	78	+ ref_count: 0,
	79	+ pinned: false,
	80	+ hash: 0,
	81	+ }
	82	+ }
	83	+}
	84	+
	85	+// ---------------------------------------------------------------------------
	86	+// KvBlockPool
	87	+// ---------------------------------------------------------------------------
	88	+
	89	+/// Fixed-size pool of KV-cache blocks with clock-sweep eviction.
	90	+///
	91	+/// When the free list is exhausted, [`alloc`](KvBlockPool::alloc) triggers a
	92	+/// clock-sweep scan to evict cacheable (ref\_count == 0) blocks.
	93	+pub struct KvBlockPool {
	94	+ metas: Vec<BlockMeta>,
	95	+ free: Vec<BlockId>,
	96	+ clock_hand: usize,
	97	+ eviction_count: u64,
	98	+ allocated_count: usize,
	99	+ cached_count: usize,
	100	+}
	101	+
	102	+impl KvBlockPool {
	103	+ /// Create a pool with `capacity` pre-allocated blocks.
	104	+ pub fn new(capacity: usize) -> Self {
	105	+ let capacity = capacity.clamp(1, MAX_BLOCKS);
	106	+ let metas = vec![BlockMeta::new(); capacity];
	107	+ let free: Vec<BlockId> = (0..capacity as BlockId).rev().collect();
	108	+ Self {
	109	+ metas,
	110	+ free,
	111	+ clock_hand: 0,
	112	+ eviction_count: 0,
	113	+ allocated_count: 0,
	114	+ cached_count: 0,
	115	+ }
	116	+ }
	117	+
	118	+ /// Total number of blocks the pool manages.
	119	+ pub fn capacity(&self) -> usize {
	120	+ self.metas.len()
	121	+ }
	122	+
	123	+ /// Number of blocks currently allocated (in-use + cached).
	124	+ pub fn allocated(&self) -> usize {
	125	+ self.allocated_count
	126	+ }
	127	+
	128	+ /// Number of blocks currently in-use (ref\_count > 0).
	129	+ pub fn in_use(&self) -> usize {
	130	+ self.allocated_count.saturating_sub(self.cached_count)
	131	+ }
	132	+
	133	+ /// Number of cacheable (ref\_count == 0) but still tracked blocks.
	134	+ pub fn cached(&self) -> usize {
	135	+ self.cached_count
	136	+ }
	137	+
	138	+ /// Number of blocks immediately available (free list).
	139	+ pub fn available(&self) -> usize {
	140	+ self.free.len()
	141	+ }
	142	+
	143	+ /// Memory pressure as a percentage of total blocks.
	144	+ pub fn pressure_pct(&self) -> f32 {
	145	+ 100.0 - (self.free.len() as f32 / self.metas.len() as f32 * 100.0)
	146	+ }
	147	+
	148	+ /// Snapshot of current pool statistics.
	149	+ pub fn stats(&self) -> PoolStats {
	150	+ PoolStats {
	151	+ total: self.metas.len(),
	152	+ allocated: self.allocated_count,
	153	+ cached: self.cached_count,
	154	+ free: self.free.len(),
	155	+ pressure_pct: self.pressure_pct(),
	156	+ evictions: self.eviction_count,
	157	+ }
	158	+ }
	159	+
	160	+ /// Allocate one block from the pool.
	161	+ ///
	162	+ /// If the free list is empty, tries to evict a cacheable block. Returns
	163	+ /// `None` if no blocks can be freed (all pinned or in-use).
	164	+ pub fn alloc(&mut self) -> Option<BlockId> {
	165	+ // Fast path: free list has capacity.
	166	+ if let Some(block) = self.free.pop() {
	167	+ let meta = &mut self.metas[block as usize];
	168	+ debug_assert_eq!(meta.ref_count, 0);
	169	+ meta.ref_count = 1;
	170	+ meta.hash = 0;
	171	+ self.allocated_count += 1;
	172	+ return Some(block);
	173	+ }
	174	+
	175	+ // Slow path: try clock-sweep eviction.
	176	+ // evict_one already decrements allocated_count and cached_count.
	177	+ let evicted = self.evict_one()?;
	178	+ let meta = &mut self.metas[evicted as usize];
	179	+ debug_assert_eq!(meta.ref_count, 0);
	180	+ meta.ref_count = 1;
	181	+ meta.hash = 0;
	182	+ self.allocated_count += 1;
	183	+ Some(evicted)
	184	+ }
	185	+
	186	+ /// Increment the reference count of an allocated (in-use) block.
	187	+ ///
	188	+ /// # Panics
	189	+ ///
	190	+ /// Panics if `block` is out of range or not currently in-use.
	191	+ pub fn retain(&mut self, block: BlockId) {
	192	+ let meta = &mut self.metas[block as usize];
	193	+ assert!(meta.ref_count > 0, "retain on non-allocated block {block}");
	194	+ meta.ref_count = meta.ref_count.saturating_add(1);
	195	+ }
	196	+
	197	+ /// Transition a cached block (ref\_count == 0) back to in-use.
	198	+ ///
	199	+ /// Called when a new sequence adopts blocks from the prefix cache.
	200	+ ///
	201	+ /// # Panics
	202	+ ///
	203	+ /// Panics if `block` is out of range or already in-use.
	204	+ pub fn reclaim(&mut self, block: BlockId) {
	205	+ let meta = &mut self.metas[block as usize];
	206	+ assert_eq!(meta.ref_count, 0, "reclaim on in-use block {block}");
	207	+ meta.ref_count = 1;
	208	+ self.cached_count = self.cached_count.saturating_sub(1);
	209	+ }
	210	+
	211	+ /// Release a reference to a block.
	212	+ ///
	213	+ /// When `ref_count` reaches zero the block transitions to the cached
	214	+ /// state (it stays tracked until evicted). Pass the token-content `hash`
	215	+ /// so the prefix cache can match against it later.
	216	+ pub fn release(&mut self, block: BlockId, hash: u64) {
	217	+ let meta = &mut self.metas[block as usize];
	218	+ assert!(meta.ref_count > 0, "release on non-allocated block {block}");
	219	+ meta.ref_count = meta.ref_count.saturating_sub(1);
	220	+ if meta.ref_count == 0 {
	221	+ meta.hash = hash;
	222	+ self.cached_count += 1;
	223	+ }
	224	+ }
	225	+
	226	+ /// Immediately return a cached (ref\_count == 0) block to the free list,
	227	+ /// undoing the cacheable state. Unlike [`KvBlockPool::release`], this makes
	228	+ /// the block reusable on the very next [`alloc`](KvBlockPool::alloc) without
	229	+ /// waiting for a clock-sweep — used by the scheduler's prefix-cache
	230	+ /// eviction path (`RadixPrefixCache::evict_filter`) to reclaim physical
	231	+ /// blocks the instant they leave the prefix tree.
	232	+ ///
	233	+ /// # Panics
	234	+ ///
	235	+ /// Panics if `block` is currently in-use (`ref_count > 0`).
	236	+ pub fn free_cached(&mut self, block: BlockId) {
	237	+ let meta = &mut self.metas[block as usize];
	238	+ assert_eq!(meta.ref_count, 0, "free_cached on in-use block {block}");
	239	+ meta.pinned = false;
	240	+ meta.hash = 0;
	241	+ self.cached_count = self.cached_count.saturating_sub(1);
	242	+ self.free.push(block);
	243	+ }
	244	+
	245	+ /// Mark a block as pinned (never evicted by the clock-sweep).
	246	+ ///
	247	+ /// A pinned block is protected from eviction regardless of its reference
	248	+ /// count — this is how the scheduler keeps cached KV blocks (which are in
	249	+ /// the `cached` state, `ref_count == 0`) resident while they are still
	250	+ /// tracked by the prefix cache. Eviction of pinned blocks is driven
	251	+ /// explicitly by [`KvBlockPool::free_cached`].
	252	+ pub fn pin(&mut self, block: BlockId) {
	253	+ let meta = &mut self.metas[block as usize];
	254	+ meta.pinned = true;
	255	+ }
	256	+
	257	+ /// Unpin a block.
	258	+ pub fn unpin(&mut self, block: BlockId) {
	259	+ self.metas[block as usize].pinned = false;
	260	+ }
	261	+
	262	+ /// Returns `true` if the block is in-use by at least one reference.
	263	+ pub fn is_in_use(&self, block: BlockId) -> bool {
	264	+ self.metas[block as usize].ref_count > 0
	265	+ }
	266	+
	267	+ /// Token-content hash stored for a cached block (0 if uncacheable/free).
	268	+ pub fn block_hash(&self, block: BlockId) -> u64 {
	269	+ self.metas[block as usize].hash
	270	+ }
	271	+
	272	+ // -- eviction -----------------------------------------------------------
	273	+
	274	+ /// Evict a single cacheable block via clock-sweep.
	275	+ ///
	276	+ /// Scans from `clock_hand` until it finds a block with `ref_count == 0`
	277	+ /// and `pinned == false`. That block is moved to the free list.
	278	+ fn evict_one(&mut self) -> Option<BlockId> {
	279	+ let n = self.metas.len();
	280	+ let start = self.clock_hand;
	281	+ for offset in 0..n {
	282	+ let idx = (start + offset) % n;
	283	+ self.clock_hand = (idx + 1) % n;
	284	+ let meta = &self.metas[idx];
	285	+ if meta.ref_count > 0 \|\| meta.pinned {
	286	+ continue;
	287	+ }
	288	+ // Found a cacheable block — evict it.
	289	+ self.metas[idx] = BlockMeta::new();
	290	+ self.free.push(idx as BlockId);
	291	+ self.eviction_count += 1;
	292	+ self.allocated_count = self.allocated_count.saturating_sub(1);
	293	+ self.cached_count = self.cached_count.saturating_sub(1);
	294	+ return Some(idx as BlockId);
	295	+ }
	296	+ None
	297	+ }
	298	+
	299	+ /// Try to evict `target` cacheable blocks (best-effort).
	300	+ pub fn evict(&mut self, target: usize) -> usize {
	301	+ let mut evicted = 0;
	302	+ for _ in 0..target {
	303	+ if self.evict_one().is_some() {
	304	+ evicted += 1;
	305	+ } else {
	306	+ break;
	307	+ }
	308	+ }
	309	+ evicted
	310	+ }
	311	+}
	312	+
	313	+// ---------------------------------------------------------------------------
	314	+// Sequence
	315	+// ---------------------------------------------------------------------------
	316	+
	317	+/// A generation session tracking its contiguous range of KV-cache blocks.
	318	+///
	319	+/// Each [`Sequence`] is created by the scheduler and lives for the duration of
	320	+/// a single request (or a chunk of a request in chunked prefill).
	321	+pub struct Sequence {
	322	+ id: u64,
	323	+ blocks: Vec<BlockId>,
	324	+ num_tokens: usize,
	325	+}
	326	+
	327	+impl Sequence {
	328	+ /// Create a new empty sequence.
	329	+ pub fn new(id: u64) -> Self {
	330	+ Self {
	331	+ id,
	332	+ blocks: Vec::new(),
	333	+ num_tokens: 0,
	334	+ }
	335	+ }
	336	+
	337	+ /// Opaque sequence identifier.
	338	+ pub fn id(&self) -> u64 {
	339	+ self.id
	340	+ }
	341	+
	342	+ /// Allocated block IDs in order.
	343	+ pub fn blocks(&self) -> &[BlockId] {
	344	+ &self.blocks
	345	+ }
	346	+
	347	+ /// Number of allocated blocks.
	348	+ pub fn num_blocks(&self) -> usize {
	349	+ self.blocks.len()
	350	+ }
	351	+
	352	+ /// Number of tokens generated so far.
	353	+ pub fn num_tokens(&self) -> usize {
	354	+ self.num_tokens
	355	+ }
	356	+
	357	+ /// Allocate one additional block and append it to this sequence.
	358	+ ///
	359	+ /// Returns the new block's ID, or `None` if the pool is exhausted.
	360	+ pub fn grow(&mut self, pool: &mut KvBlockPool) -> Option<BlockId> {
	361	+ let block = pool.alloc()?;
	362	+ self.blocks.push(block);
	363	+ Some(block)
	364	+ }
	365	+
	366	+ /// Record that `n` tokens have been generated (advances the logical
	367	+ /// cursor without allocating new blocks).
	368	+ pub fn advance(&mut self, n: usize) {
	369	+ self.num_tokens += n;
	370	+ }
	371	+
	372	+ /// Set the token count directly (used when restoring from a cached prefix).
	373	+ pub fn set_num_tokens(&mut self, n: usize) {
	374	+ self.num_tokens = n;
	375	+ }
	376	+
	377	+ /// Restore blocks and token count from a prefix-cache hit.
	378	+ ///
	379	+ /// Calls [`KvBlockPool::reclaim`] for each block to transition it from
	380	+ /// cacheable to in-use.
	381	+ pub fn restore_from_cache(
	382	+ &mut self,
	383	+ blocks: &[BlockId],
	384	+ num_tokens: usize,
	385	+ pool: &mut KvBlockPool,
	386	+ ) {
	387	+ self.blocks = blocks.to_vec();
	388	+ self.num_tokens = num_tokens;
	389	+ for &b in &self.blocks {
	390	+ pool.reclaim(b);
	391	+ }
	392	+ }
	393	+}
	394	+
	395	+// ---------------------------------------------------------------------------
	396	+// Prefix cache
	397	+// ---------------------------------------------------------------------------
	398	+
	399	+/// FNV-1a hash of a token slice.
	400	+fn hash_tokens(tokens: &[u32]) -> u64 {
	401	+ let mut h = 0xcbf29ce484222325u64;
	402	+ for &t in tokens {
	403	+ h ^= t as u64;
	404	+ h = h.wrapping_mul(0x100000001b3);
	405	+ }
	406	+ h
	407	+}
	408	+
	409	+/// Hash-based prefix cache for KV-block reuse.
	410	+///
	411	+/// The cache stores every prefix (first N blocks) of completed sequences.
	412	+/// When a new prompt arrives the cache is queried from longest to shortest
	413	+/// prefix until a match is found.
	414	+///
	415	+/// This uses the observation that LLM serving has high prefix overlap
	416	+/// (system prompt, chat history, few-shot examples). The hash-map approach
	417	+/// is simple, fast for lookup, and handles partial reuse correctly — the
	418	+/// downside is that it cannot match arbitrary shared substrings (only
	419	+/// aligned prefixes). A radix-tree upgrade is tracked for a future
	420	+/// optimization pass.
	421	+pub struct PrefixCache {
	422	+ /// Maps `prefix_hash(tokens, n_blocks)` → `Vec<BlockId>`.
	423	+ entries: HashMap<u64, Vec<BlockId>>,
	424	+ /// Maps `prefix_hash` → number of tokens the prefix covers.
	425	+ prefix_lengths: HashMap<u64, usize>,
	426	+ hits: std::cell::Cell<u64>,
	427	+ misses: std::cell::Cell<u64>,
	428	+}
	429	+
	430	+impl PrefixCache {
	431	+ /// Create an empty prefix cache.
	432	+ pub fn new() -> Self {
	433	+ Self {
	434	+ entries: HashMap::new(),
	435	+ prefix_lengths: HashMap::new(),
	436	+ hits: std::cell::Cell::new(0),
	437	+ misses: std::cell::Cell::new(0),
	438	+ }
	439	+ }
	440	+
	441	+ /// Compute the rolling hash for the first `num_tokens` of `tokens`.
	442	+ fn prefix_hash(tokens: &[u32], num_tokens: usize) -> u64 {
	443	+ let end = tokens.len().min(num_tokens);
	444	+ hash_tokens(&tokens[..end])
	445	+ }
	446	+
	447	+ /// Number of entries in the cache.
	448	+ pub fn len(&self) -> usize {
	449	+ self.entries.len()
	450	+ }
	451	+
	452	+ /// Total KV-block references stored across all prefix entries. Because each
	453	+ /// prefix entry stores the full block list for its length, a block that
	454	+ /// appears in many shared prefixes is counted once per entry — so this
	455	+ /// grows super-linearly for nested (chat-history-style) prompts. Compare
	456	+ /// with [`RadixPrefixCache::cached_tokens`], which stores each block once.
	457	+ pub fn total_stored_blocks(&self) -> usize {
	458	+ self.entries.values().map(\|blocks\| blocks.len()).sum()
	459	+ }
	460	+
	461	+ pub fn is_empty(&self) -> bool {
	462	+ self.entries.is_empty()
	463	+ }
	464	+
	465	+ /// Cache statistics.
	466	+ pub fn stats(&self) -> PrefixCacheStats {
	467	+ PrefixCacheStats {
	468	+ entries: self.entries.len(),
	469	+ hits: self.hits.get(),
	470	+ misses: self.misses.get(),
	471	+ }
	472	+ }
	473	+
	474	+ /// Insert the blocks for a completed sequence into the cache.
	475	+ ///
	476	+ /// Every prefix (first 1 block, first 2 blocks, …, all blocks) is stored
	477	+ /// under its token-content hash.
	478	+ ///
	479	+ /// The caller must call [`KvBlockPool::release`] separately to
	480	+ /// transition the blocks to the cacheable state.
	481	+ pub fn insert(&mut self, tokens: &[u32], blocks: &[BlockId]) {
	482	+ for n in 1..=blocks.len() {
	483	+ let num_tokens = n * BLOCK_SIZE;
	484	+ let h = Self::prefix_hash(tokens, num_tokens);
	485	+ self.entries.entry(h).or_insert_with(\|\| {
	486	+ self.prefix_lengths.insert(h, num_tokens);
	487	+ blocks[..n].to_vec()
	488	+ });
	489	+ }
	490	+ }
	491	+
	492	+ /// Look up the longest matching prefix.
	493	+ ///
	494	+ /// Returns `(num_tokens_matched, block_ids)` for the longest prefix that
	495	+ /// exists in the cache, or `None` if no prefix matches.
	496	+ pub fn lookup(&self, tokens: &[u32]) -> Option<(usize, &[BlockId])> {
	497	+ let max_blocks = tokens.len().div_ceil(BLOCK_SIZE);
	498	+ // Search from longest to shortest so we get the best match.
	499	+ for n in (1..=max_blocks).rev() {
	500	+ let num_tokens = n * BLOCK_SIZE;
	501	+ let h = Self::prefix_hash(tokens, num_tokens);
	502	+ if let Some(blocks) = self.entries.get(&h) {
	503	+ self.hits.set(self.hits.get() + 1);
	504	+ return Some((num_tokens, blocks.as_slice()));
	505	+ }
	506	+ }
	507	+ self.misses.set(self.misses.get() + 1);
	508	+ None
	509	+ }
	510	+
	511	+ /// Remove an entry from the cache. Returns the blocks that were cached.
	512	+ ///
	513	+ /// Called by the eviction path in [`KvBlockPool`].
	514	+ pub fn remove(&mut self, hash: u64) -> Option<Vec<BlockId>> {
	515	+ self.prefix_lengths.remove(&hash);
	516	+ self.entries.remove(&hash)
	517	+ }
	518	+}
	519	+
	520	+impl Default for PrefixCache {
	521	+ fn default() -> Self {
	522	+ Self::new()
	523	+ }
	524	+}
	525	+
	526	+// ---------------------------------------------------------------------------
	527	+// Tests
	528	+// ---------------------------------------------------------------------------
	529	+
	530	+#[cfg(test)]
	531	+mod tests {
	532	+ use super::*;
	533	+
	534	+ // -- KvBlockPool tests --------------------------------------------------
	535	+
	536	+ #[test]
	537	+ fn pool_new_has_capacity() {
	538	+ let pool = KvBlockPool::new(64);
	539	+ assert_eq!(pool.capacity(), 64);
	540	+ assert_eq!(pool.available(), 64);
	541	+ assert_eq!(pool.allocated(), 0);
	542	+ assert!(pool.pressure_pct() < 1.0);
	543	+ }
	544	+
	545	+ #[test]
	546	+ fn pool_alloc_and_release() {
	547	+ let mut pool = KvBlockPool::new(4);
	548	+ let b0 = pool.alloc().unwrap();
	549	+ let b1 = pool.alloc().unwrap();
	550	+ assert_ne!(b0, b1);
	551	+ assert_eq!(pool.allocated(), 2);
	552	+ assert_eq!(pool.available(), 2);
	553	+ pool.release(b0, 42);
	554	+ assert_eq!(pool.in_use(), 1);
	555	+ assert_eq!(pool.cached(), 1);
	556	+ pool.release(b1, 99);
	557	+ assert_eq!(pool.in_use(), 0);
	558	+ assert_eq!(pool.cached(), 2);
	559	+ }
	560	+
	561	+ #[test]
	562	+ fn pool_retain_multiple_references() {
	563	+ let mut pool = KvBlockPool::new(4);
	564	+ let b = pool.alloc().unwrap();
	565	+ pool.retain(b); // ref_count: 1 → 2
	566	+ assert!(pool.is_in_use(b));
	567	+ pool.release(b, 0); // 2 → 1
	568	+ assert!(pool.is_in_use(b));
	569	+ pool.release(b, 0); // 1 → 0
	570	+ assert!(!pool.is_in_use(b));
	571	+ assert_eq!(pool.cached(), 1);
	572	+ }
	573	+
	574	+ #[test]
	575	+ fn pool_exhaustion_triggers_eviction() {
	576	+ let mut pool = KvBlockPool::new(2);
	577	+ let b0 = pool.alloc().unwrap();
	578	+ let _b1 = pool.alloc().unwrap();
	579	+ assert!(pool.alloc().is_none()); // free list empty, nothing to evict
	580	+
	581	+ // Release b0 to cache → now there is a cacheable block.
	582	+ pool.release(b0, 42);
	583	+ assert_eq!(pool.cached(), 1);
	584	+
	585	+ // Next alloc should evict b0 (b1 is still in-use and thus skipped).
	586	+ let b2 = pool.alloc().unwrap();
	587	+ assert_eq!(b2, b0); // recycled
	588	+ assert_eq!(pool.stats().evictions, 1);
	589	+ assert_eq!(pool.in_use(), 2); // b1 + b2
	590	+ }
	591	+
	592	+ #[test]
	593	+ fn pool_pinned_blocks_are_not_evicted() {
	594	+ let mut pool = KvBlockPool::new(2);
	595	+ let b0 = pool.alloc().unwrap();
	596	+ pool.pin(b0);
	597	+ let _b1 = pool.alloc().unwrap();
	598	+ pool.release(b0, 42); // b0 is cached, but pinned
	599	+ pool.release(_b1, 99); // b1 is cached, unpinned
	600	+
	601	+ // Evict should skip b0 (pinned) and evict b1.
	602	+ let evicted = pool.evict(1);
	603	+ assert_eq!(evicted, 1);
	604	+ assert_eq!(pool.stats().evictions, 1);
	605	+ // b0 is still allocated (pinned).
	606	+ assert_eq!(pool.allocated(), 1);
	607	+ }
	608	+
	609	+ #[test]
	610	+ fn pool_evict_up_to_target() {
	611	+ let mut pool = KvBlockPool::new(10);
	612	+ let blocks: Vec<_> = (0..10).map(\|_\| pool.alloc().unwrap()).collect();
	613	+ // Release all to cache.
	614	+ for (i, &b) in blocks.iter().enumerate() {
	615	+ pool.release(b, i as u64);
	616	+ }
	617	+ assert_eq!(pool.cached(), 10);
	618	+
	619	+ let evicted = pool.evict(4);
	620	+ assert_eq!(evicted, 4);
	621	+ assert_eq!(pool.available(), 4);
	622	+ assert_eq!(pool.cached(), 6);
	623	+ }
	624	+
	625	+ #[test]
	626	+ fn pool_stats_consistent() {
	627	+ let mut pool = KvBlockPool::new(8);
	628	+ for _ in 0..5 {
	629	+ pool.alloc().unwrap();
	630	+ }
	631	+ let s = pool.stats();
	632	+ assert_eq!(s.total, 8);
	633	+ assert_eq!(s.allocated, 5);
	634	+ assert_eq!(s.free, 3);
	635	+ }
	636	+
	637	+ // -- PrefixCache tests --------------------------------------------------
	638	+
	639	+ #[test]
	640	+ fn prefix_cache_insert_and_lookup_exact() {
	641	+ let mut cache = PrefixCache::new();
	642	+ let tokens: Vec<u32> = (0..48).collect(); // 3 blocks
	643	+ let blocks = vec![10, 20, 30];
	644	+ cache.insert(&tokens, &blocks);
	645	+
	646	+ // Look up entire sequence.
	647	+ let (matched, cached) = cache.lookup(&tokens).unwrap();
	648	+ assert_eq!(matched, 48); // 3 blocks × 16
	649	+ assert_eq!(cached, &[10, 20, 30]);
	650	+ }
	651	+
	652	+ #[test]
	653	+ fn prefix_cache_partial_match() {
	654	+ let mut cache = PrefixCache::new();
	655	+ let tokens: Vec<u32> = (0..32).collect(); // 2 blocks
	656	+ let blocks = vec![10, 20];
	657	+ cache.insert(&tokens, &blocks);
	658	+
	659	+ // Same prefix but shorter.
	660	+ let short: Vec<u32> = (0..16).collect(); // 1 block
	661	+ let (matched, cached) = cache.lookup(&short).unwrap();
	662	+ assert_eq!(matched, 16);
	663	+ assert_eq!(cached, &[10]);
	664	+
	665	+ // Longer prompt that extends the prefix.
	666	+ let longer: Vec<u32> = (0..40).collect(); // 3 blocks, first 2 match
	667	+ let (matched, cached) = cache.lookup(&longer).unwrap();
	668	+ assert_eq!(matched, 32);
	669	+ assert_eq!(cached, &[10, 20]);
	670	+ }
	671	+
	672	+ #[test]
	673	+ fn prefix_cache_miss() {
	674	+ let mut cache = PrefixCache::new();
	675	+ cache.insert(&[0, 1, 2], &[1]);
	676	+
	677	+ let result = cache.lookup(&[99, 100, 101]);
	678	+ assert!(result.is_none());
	679	+ assert_eq!(cache.stats().misses, 1);
	680	+ }
	681	+
	682	+ #[test]
	683	+ fn prefix_cache_empty_is_miss() {
	684	+ let cache = PrefixCache::new();
	685	+ assert!(cache.is_empty());
	686	+ assert_eq!(cache.len(), 0);
	687	+ let result = cache.lookup(&[1, 2, 3]);
	688	+ assert!(result.is_none());
	689	+ }
	690	+
	691	+ #[test]
	692	+ fn prefix_cache_tracks_hits() {
	693	+ let mut cache = PrefixCache::new();
	694	+ let tokens: Vec<u32> = (0..16).collect();
	695	+ cache.insert(&tokens, &[7]);
	696	+
	697	+ // First lookup → hit.
	698	+ let _ = cache.lookup(&tokens);
	699	+ assert_eq!(cache.stats().hits, 1);
	700	+
	701	+ // Second lookup → hit.
	702	+ let _ = cache.lookup(&tokens);
	703	+ assert_eq!(cache.stats().hits, 2);
	704	+
	705	+ // Miss.
	706	+ let _ = cache.lookup(&[99]);
	707	+ assert_eq!(cache.stats().misses, 1);
	708	+ }
	709	+
	710	+ // -- Sequence tests -----------------------------------------------------
	711	+
	712	+ #[test]
	713	+ fn sequence_grow_allocates_blocks() {
	714	+ let mut pool = KvBlockPool::new(10);
	715	+ let mut seq = Sequence::new(1);
	716	+ assert_eq!(seq.num_blocks(), 0);
	717	+
	718	+ let b0 = seq.grow(&mut pool).unwrap();
	719	+ let b1 = seq.grow(&mut pool).unwrap();
	720	+ assert_eq!(seq.num_blocks(), 2);
	721	+ assert_eq!(seq.blocks(), &[b0, b1]);
	722	+ assert_eq!(pool.in_use(), 2);
	723	+ }
	724	+
	725	+ #[test]
	726	+ fn sequence_advance_tracks_tokens() {
	727	+ let mut seq = Sequence::new(42);
	728	+ assert_eq!(seq.num_tokens(), 0);
	729	+ seq.advance(16);
	730	+ assert_eq!(seq.num_tokens(), 16);
	731	+ seq.advance(8);
	732	+ assert_eq!(seq.num_tokens(), 24);
	733	+ }
	734	+
	735	+ #[test]
	736	+ fn sequence_grow_exhausted_pool_returns_none() {
	737	+ let mut pool = KvBlockPool::new(1);
	738	+ let mut seq = Sequence::new(1);
	739	+ assert!(seq.grow(&mut pool).is_some());
	740	+ assert!(seq.grow(&mut pool).is_none()); // no blocks to evict
	741	+ }
	742	+
	743	+ // -- Integration: pool + prefix cache + sequence ------------------------
	744	+
	745	+ #[test]
	746	+ fn end_to_end_prefix_caching() {
	747	+ let mut pool = KvBlockPool::new(16);
	748	+ let mut cache = PrefixCache::new();
	749	+
	750	+ // First request: process tokens 0..48, allocating 3 blocks.
	751	+ let tokens: Vec<u32> = (0..48).collect();
	752	+ let mut seq = Sequence::new(1);
	753	+ let blocks: Vec<BlockId> = (0..3).map(\|_\| seq.grow(&mut pool).unwrap()).collect();
	754	+
	755	+ // Insert into prefix cache and release blocks.
	756	+ cache.insert(&tokens, &blocks);
	757	+ for &b in &blocks {
	758	+ pool.release(b, hash_tokens(&tokens[..16])); // simplified
	759	+ }
	760	+ assert_eq!(pool.cached(), 3);
	761	+
	762	+ // Second request with same prefix: cache hit.
	763	+ let result = cache.lookup(&tokens);
	764	+ assert!(result.is_some());
	765	+ let (_, cached_blocks) = result.unwrap();
	766	+ assert_eq!(cached_blocks, &blocks[..]);
	767	+
	768	+ // The new sequence can skip the first 3 blocks.
	769	+ let mut seq2 = Sequence::new(2);
	770	+ seq2.set_num_tokens(48); // restored from cache
	771	+ assert_eq!(seq2.num_tokens(), 48);
	772	+ }
	773	+}
0	774

▼crates/justapi-inference/src/lib.rs

+74-0

@@ -0,0 +1,74 @@
	1	+//! # justapi-inference
	2	+//!
	3	+//! Native Rust inference engine for JustAPI — the "own the GPU path" layer that
	4	+//! makes JustAPI a better model-serving framework than FastAPI (which delegates
	5	+//! all GPU work to `torch` under the GIL).
	6	+//!
	7	+//! Phase 41 (foundation): crate + Candle integration, device management,
	8	+//! a `Model` trait, a streaming `Engine` with an in-memory model registry, and a
	9	+//! CPU-runnable `MockModel` so the whole pipeline is exercised without weights.
	10	+//!
	11	+//! Real weight loading + the KV-cache-backed forward pass land in Phase 42
	12	+//! (behind the `real` feature, gated on having model weights + a GPU toolkit).
	13	+//!
	14	+//! Design notes:
	15	+//! - Generation runs on a dedicated OS thread (Candle compute is synchronous);
	16	+//! tokens are streamed back over a tokio `mpsc` channel. The hot path is
	17	+//! GIL-free Rust — exactly the property FastAPI cannot offer.
	18	+//! - The `Model` trait is the extension point: `MockModel` for tests/demos,
	19	+//! `LlamaModel` (Phase 42) for real inference.
	20	+
	21	+mod autoscaler;
	22	+mod control_plane;
	23	+mod engine;
	24	+mod gateway;
	25	+mod kv_cache;
	26	+mod model;
	27	+pub mod openai;
	28	+pub mod pd;
	29	+pub mod radix_cache;
	30	+pub mod real;
	31	+mod router;
	32	+mod scheduler;
	33	+mod scheduler_engine;
	34	+mod spec_decode;
	35	+mod spec_decode_tree;
	36	+mod supervisor;
	37	+
	38	+pub use autoscaler::{Autoscaler, AutoscalerConfig, LlmMetrics, ScaleDecision};
	39	+pub use control_plane::{
	40	+ ControlPlane, ModelRecord, ModelVersion, ResolvedModel, RuntimeProfile, WeightLocation,
	41	+};
	42	+pub use engine::{Engine, EngineDevice, ModelRegistry, ModelSource};
	43	+pub use gateway::{GatewayConfig, GatewayDecision, InferenceGateway};
	44	+pub use kv_cache::{
	45	+ BlockId, KvBlockPool, PoolStats, PrefixCache, PrefixCacheStats, Sequence, BLOCK_SIZE,
	46	+ MAX_BLOCKS,
	47	+};
	48	+pub use model::{FinishReason, GeneratedToken, MockModel, Model, ModelError, SamplingParams};
	49	+pub use openai::{
	50	+ ApiError, ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage,
	51	+ ChatMessageOut, CompletionChoice, CompletionRequest, CompletionResponse, Delta, EmbeddingInput,
	52	+ EmbeddingObject, EmbeddingRequest, EmbeddingResponse, ErrorResponse, ModelList, ModelObject,
	53	+ ModelPermission, Usage,
	54	+};
	55	+pub use pd::{PdScheduler, PdStats};
	56	+pub use radix_cache::{RadixPrefixCache, RadixPrefixCacheStats};
	57	+pub use real::lora::{LoraAdapter, LoraConfig, LoraRegistry};
	58	+pub use real::quant::gguf::GgufType as RawGgufType;
	59	+pub use real::quant::{LayerQuantConfig, QuantConfig, QuantMethod};
	60	+pub use router::{Replica, RouteDecision, RouteRequest, Router, RoutingStrategy};
	61	+pub use scheduler::TransferableSequence;
	62	+pub use scheduler::{
	63	+ NewRequest, PrefillStep, Schedule, Scheduler, SchedulerConfig, SchedulerStats,
	64	+};
	65	+pub use scheduler_engine::SchedulerEngine;
	66	+pub use spec_decode::{speculative_generate, AcceptanceStats, SpeculativeConfig, SpeculativeModel};
	67	+pub use spec_decode_tree::{
	68	+ build_draft_tree, speculative_generate_tree, top_k_tokens, verify_tree, DraftTree, TreeNode,
	69	+ TreeSpeculativeModel,
	70	+};
	71	+pub use supervisor::{LiveReplica, Supervisor, SupervisorAction, SupervisorConfig};
	72	+
	73	+/// Default vocabulary size used by [`MockModel`] when none is supplied.
	74	+pub const DEFAULT_MOCK_VOCAB: usize = 32_000;
0	75

▼crates/justapi-inference/src/model.rs

+187-0

@@ -0,0 +1,187 @@
	1	+//! Model trait, sampling parameters, generated-token stream items, and the
	2	+//! CPU-runnable [`MockModel`] used for tests and offline demos.
	3	+
	4	+/// Stopping / sampling configuration for a generation request.
	5	+///
	6	+/// Mirrors the knobs vLLM/SGLang expose; the real forward pass (Phase 42) will
	7	+/// honor `temperature`/`top_p`/`top_k` when sampling logits.
	8	+#[derive(Debug, Clone)]
	9	+pub struct SamplingParams {
	10	+ /// Maximum number of new tokens to generate.
	11	+ pub max_tokens: usize,
	12	+ /// Sampling temperature (0.0 = greedy). Real models only; MockModel ignores it.
	13	+ pub temperature: f32,
	14	+ /// Nucleus sampling cutoff.
	15	+ pub top_p: f32,
	16	+ /// Top-k cutoff (0 = disabled).
	17	+ pub top_k: usize,
	18	+ /// Token ids that, when produced, end generation with `FinishReason::Stop`.
	19	+ pub stop_tokens: Vec<u32>,
	20	+ /// Substrings that, when produced, end generation (real tokenizers only).
	21	+ pub stop_sequences: Vec<String>,
	22	+}
	23	+
	24	+impl Default for SamplingParams {
	25	+ fn default() -> Self {
	26	+ Self {
	27	+ max_tokens: 256,
	28	+ temperature: 1.0,
	29	+ top_p: 1.0,
	30	+ top_k: 0,
	31	+ stop_tokens: Vec::new(),
	32	+ stop_sequences: Vec::new(),
	33	+ }
	34	+ }
	35	+}
	36	+
	37	+/// A single streamed token produced during generation.
	38	+#[derive(Debug, Clone)]
	39	+pub struct GeneratedToken {
	40	+ /// Vocabulary id of the token.
	41	+ pub id: u32,
	42	+ /// Decoded text for the token (may be empty/partial for sub-word pieces).
	43	+ pub text: String,
	44	+ /// Log-probability of the token (0.0 for [`MockModel`]).
	45	+ pub logprob: f32,
	46	+ /// Set on the final token of the stream.
	47	+ pub finish_reason: Option<FinishReason>,
	48	+}
	49	+
	50	+/// Why a generation stream ended.
	51	+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	52	+pub enum FinishReason {
	53	+ /// Hit `max_tokens`.
	54	+ Length,
	55	+ /// Produced a stop token / sequence.
	56	+ Stop,
	57	+ /// Downstream consumer stopped reading (e.g. client disconnected).
	58	+ Cancelled,
	59	+}
	60	+
	61	+/// Errors produced by the inference engine and models.
	62	+#[derive(Debug, thiserror::Error)]
	63	+pub enum ModelError {
	64	+ #[error("model '{0}' not found in registry")]
	65	+ NotFound(String),
	66	+ #[error("feature '{0}' is required for this operation (enable it in the crate features)")]
	67	+ FeatureRequired(&'static str),
	68	+ #[error("real model loading is not implemented yet (lands with the KV-cache in Phase 42)")]
	69	+ NotImplemented,
	70	+ #[error("generation failed: {0}")]
	71	+ Generation(String),
	72	+ #[error(transparent)]
	73	+ Io(#[from] std::io::Error),
	74	+}
	75	+
	76	+/// A loadable / runnable model behind the engine.
	77	+///
	78	+/// Implementors must be `Send + Sync` so they can be shared across the registry
	79	+/// and driven from a generation thread.
	80	+pub trait Model: Send + Sync {
	81	+ /// Size of the model vocabulary.
	82	+ fn vocab_size(&self) -> usize;
	83	+
	84	+ /// Run generation for `prompt` (token ids), streaming each produced token to
	85	+ /// `sink`. Return `false` from `sink` to abort early (e.g. client gone).
	86	+ /// Returns the reason the stream ended.
	87	+ fn generate(
	88	+ &self,
	89	+ prompt: &[u32],
	90	+ params: &SamplingParams,
	91	+ sink: &dyn Fn(GeneratedToken) -> bool,
	92	+ ) -> Result<FinishReason, ModelError>;
	93	+
	94	+ /// Return the unnormalized next-token logits for `context` (the raw score
	95	+ /// vector over the vocabulary at the position after `context`).
	96	+ ///
	97	+ /// This is the single-step primitive required for speculative decoding: a
	98	+ /// draft model proposes tokens and the target model scores them to accept
	99	+ /// or reject. The default implementation derives a degenerate distribution
	100	+ /// from [`Model::generate`]'s behavior, but real models should compute true
	101	+ /// logits from a single forward pass.
	102	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError>;
	103	+}
	104	+
	105	+/// Deterministic, weight-free model for tests, demos, and CI on CPU-only boxes.
	106	+///
	107	+/// It echoes the last prompt token and then walks the vocabulary one id at a
	108	+/// time (`(prev + 1) % vocab`), decoding each id to a lowercase letter so the
	109	+/// stream is human-readable. It is NOT a real LLM — it exists to exercise the
	110	+/// full engine/streaming/registry path end-to-end.
	111	+pub struct MockModel {
	112	+ vocab_size: usize,
	113	+}
	114	+
	115	+impl MockModel {
	116	+ /// Create a mock model with the given vocabulary size (clamped to >= 1).
	117	+ pub fn new(vocab_size: usize) -> Self {
	118	+ Self {
	119	+ vocab_size: vocab_size.max(1),
	120	+ }
	121	+ }
	122	+
	123	+ /// Trivial tokenization: each byte of `text` becomes a token id.
	124	+ pub fn tokenize(text: &str) -> Vec<u32> {
	125	+ text.bytes().map(\|b\| b as u32).collect()
	126	+ }
	127	+
	128	+ /// Trivial detokenization: low byte of each id -> char (lossy for >255).
	129	+ pub fn detokenize(ids: &[u32]) -> String {
	130	+ ids.iter().map(\|&id\| (id as u8) as char).collect()
	131	+ }
	132	+}
	133	+
	134	+impl Model for MockModel {
	135	+ fn vocab_size(&self) -> usize {
	136	+ self.vocab_size
	137	+ }
	138	+
	139	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError> {
	140	+ let vocab = self.vocab_size;
	141	+ let prev = context.last().copied().unwrap_or(0) % vocab as u32;
	142	+ let next = (prev.wrapping_add(1)) % vocab as u32;
	143	+ // Degenerate-but-smooth distribution: a sharp peak at the next token so
	144	+ // greedy decoding reproduces `generate`'s `(prev + 1) % vocab` rule,
	145	+ // with a small uniform tail so temperature/top_p sampling still works.
	146	+ let mut logits = vec![-5.0f32; vocab];
	147	+ logits[next as usize] = 5.0;
	148	+ let r = (prev.wrapping_add(3)) % vocab as u32;
	149	+ logits[r as usize] = 0.0;
	150	+ Ok(logits)
	151	+ }
	152	+
	153	+ fn generate(
	154	+ &self,
	155	+ prompt: &[u32],
	156	+ params: &SamplingParams,
	157	+ sink: &dyn Fn(GeneratedToken) -> bool,
	158	+ ) -> Result<FinishReason, ModelError> {
	159	+ let mut prev = prompt.last().copied().unwrap_or(0);
	160	+ for _ in 0..params.max_tokens {
	161	+ let next = (prev.wrapping_add(1)) % self.vocab_size as u32;
	162	+ let text = decode_mock(next);
	163	+ let stop = params.stop_tokens.contains(&next);
	164	+ let finish_reason = if stop { Some(FinishReason::Stop) } else { None };
	165	+ let token = GeneratedToken {
	166	+ id: next,
	167	+ text,
	168	+ logprob: 0.0,
	169	+ finish_reason,
	170	+ };
	171	+ if !sink(token) {
	172	+ return Ok(FinishReason::Cancelled);
	173	+ }
	174	+ if stop {
	175	+ return Ok(FinishReason::Stop);
	176	+ }
	177	+ prev = next;
	178	+ }
	179	+ Ok(FinishReason::Length)
	180	+ }
	181	+}
	182	+
	183	+/// Map a token id to a single lowercase letter so mock streams are readable.
	184	+fn decode_mock(id: u32) -> String {
	185	+ let c = (b'a' + (id % 26) as u8) as char;
	186	+ c.to_string()
	187	+}
0	188

▼crates/justapi-inference/src/openai.rs

+1064-0

@@ -0,0 +1,1064 @@
	1	+//! OpenAI-compatible request/response protocol for JustAPI inference.
	2	+//!
	3	+//! Provides the data types and streaming/non-streaming drivers that make a
	4	+//! [`Engine`](crate::engine::Engine) speak the OpenAI Chat Completions,
	5	+//! Completions, Embeddings, and Models APIs. The hot path (token generation)
	6	+//! runs inside the engine on a dedicated OS thread; this module only formats
	7	+//! the I/O — it never touches the GIL.
	8	+//!
	9	+//! The HTTP-wiring (mounting these as routes) lives in `justapi-core`'s
	10	+//! `openai` module behind the `inference` feature, so this crate stays
	11	+//! dependency-light and testable on its own.
	12	+
	13	+use std::sync::Arc;
	14	+use std::time::{SystemTime, UNIX_EPOCH};
	15	+
	16	+use serde::{Deserialize, Serialize};
	17	+
	18	+use crate::engine::Engine;
	19	+use crate::model::{FinishReason, MockModel, ModelError, SamplingParams};
	20	+use crate::scheduler_engine::SchedulerEngine;
	21	+
	22	+// ---------------------------------------------------------------------------
	23	+// Request types
	24	+// ---------------------------------------------------------------------------
	25	+
	26	+/// A single chat message.
	27	+#[derive(Debug, Clone, Deserialize, Serialize)]
	28	+pub struct ChatMessage {
	29	+ pub role: String,
	30	+ pub content: String,
	31	+ #[serde(default)]
	32	+ pub name: Option<String>,
	33	+}
	34	+
	35	+fn default_max_tokens() -> usize {
	36	+ 256
	37	+}
	38	+
	39	+/// `POST /v1/chat/completions` request.
	40	+#[derive(Debug, Clone, Deserialize, Serialize)]
	41	+pub struct ChatCompletionRequest {
	42	+ /// Model registry name (the name the model was registered under).
	43	+ pub model: String,
	44	+ pub messages: Vec<ChatMessage>,
	45	+ #[serde(default)]
	46	+ pub temperature: Option<f32>,
	47	+ #[serde(default = "default_max_tokens")]
	48	+ pub max_tokens: usize,
	49	+ /// Newer OpenAI clients send this instead of `max_tokens`.
	50	+ #[serde(default)]
	51	+ pub max_completion_tokens: Option<usize>,
	52	+ #[serde(default)]
	53	+ pub top_p: Option<f32>,
	54	+ /// Stream tokens as Server-Sent Events.
	55	+ #[serde(default)]
	56	+ pub stream: bool,
	57	+ /// Stop strings (ignored by the mock model; honored by real models).
	58	+ #[serde(default)]
	59	+ pub stop: Option<Vec<String>>,
	60	+ #[serde(default)]
	61	+ pub user: Option<String>,
	62	+}
	63	+
	64	+impl ChatCompletionRequest {
	65	+ /// Resolve the effective max-token count across both fields.
	66	+ pub fn effective_max_tokens(&self) -> usize {
	67	+ self.max_completion_tokens.unwrap_or(self.max_tokens)
	68	+ }
	69	+}
	70	+
	71	+/// `POST /v1/completions` request.
	72	+#[derive(Debug, Clone, Deserialize)]
	73	+pub struct CompletionRequest {
	74	+ pub model: String,
	75	+ pub prompt: String,
	76	+ #[serde(default)]
	77	+ pub temperature: Option<f32>,
	78	+ #[serde(default = "default_max_tokens")]
	79	+ pub max_tokens: usize,
	80	+ #[serde(default)]
	81	+ pub top_p: Option<f32>,
	82	+ #[serde(default)]
	83	+ pub stream: bool,
	84	+ #[serde(default)]
	85	+ pub stop: Option<Vec<String>>,
	86	+}
	87	+
	88	+/// `POST /v1/embeddings` request.
	89	+#[derive(Debug, Clone, Deserialize)]
	90	+pub struct EmbeddingRequest {
	91	+ pub model: String,
	92	+ #[serde(default)]
	93	+ pub input: EmbeddingInput,
	94	+ #[serde(default)]
	95	+ pub encoding_format: Option<String>,
	96	+}
	97	+
	98	+/// Embedding input: a single string or a list of strings.
	99	+#[derive(Debug, Clone, Deserialize)]
	100	+#[serde(untagged)]
	101	+pub enum EmbeddingInput {
	102	+ Single(String),
	103	+ Many(Vec<String>),
	104	+}
	105	+
	106	+impl Default for EmbeddingInput {
	107	+ fn default() -> Self {
	108	+ EmbeddingInput::Single(String::new())
	109	+ }
	110	+}
	111	+
	112	+impl EmbeddingInput {
	113	+ fn into_texts(self) -> Vec<String> {
	114	+ match self {
	115	+ EmbeddingInput::Single(s) => vec![s],
	116	+ EmbeddingInput::Many(v) => v,
	117	+ }
	118	+ }
	119	+}
	120	+
	121	+// ---------------------------------------------------------------------------
	122	+// Response types
	123	+// ---------------------------------------------------------------------------
	124	+
	125	+/// An OpenAI `model` object.
	126	+#[derive(Debug, Clone, Serialize, Default)]
	127	+pub struct ModelObject {
	128	+ pub id: String,
	129	+ pub object: &'static str,
	130	+ pub created: u64,
	131	+ pub owned_by: String,
	132	+ #[serde(default)]
	133	+ pub permission: Vec<ModelPermission>,
	134	+ #[serde(default)]
	135	+ pub root: Option<String>,
	136	+ #[serde(default)]
	137	+ pub parent: Option<String>,
	138	+}
	139	+
	140	+/// OpenAI model permission entry (placeholder).
	141	+#[derive(Debug, Clone, Serialize, Default)]
	142	+pub struct ModelPermission {
	143	+ #[serde(default)]
	144	+ pub id: String,
	145	+ #[serde(default)]
	146	+ pub object: &'static str,
	147	+ #[serde(default)]
	148	+ pub created: u64,
	149	+ #[serde(default)]
	150	+ pub allow_create_engine: bool,
	151	+ #[serde(default)]
	152	+ pub allow_sampling: bool,
	153	+ #[serde(default)]
	154	+ pub allow_logprobs: bool,
	155	+ #[serde(default)]
	156	+ pub allow_search_indices: bool,
	157	+ #[serde(default)]
	158	+ pub allow_view: bool,
	159	+ #[serde(default)]
	160	+ pub organization: String,
	161	+ #[serde(default)]
	162	+ pub group: Option<String>,
	163	+ #[serde(default)]
	164	+ pub is_blocking: bool,
	165	+}
	166	+
	167	+/// `GET /v1/models` response.
	168	+#[derive(Debug, Clone, Serialize)]
	169	+pub struct ModelList {
	170	+ pub object: &'static str,
	171	+ pub data: Vec<ModelObject>,
	172	+}
	173	+
	174	+/// Token usage accounting.
	175	+#[derive(Debug, Clone, Copy, Serialize)]
	176	+pub struct Usage {
	177	+ pub prompt_tokens: usize,
	178	+ pub completion_tokens: usize,
	179	+ pub total_tokens: usize,
	180	+}
	181	+
	182	+/// `POST /v1/chat/completions` (non-streaming) response.
	183	+#[derive(Debug, Clone, Serialize)]
	184	+pub struct ChatCompletionResponse {
	185	+ pub id: String,
	186	+ pub object: &'static str,
	187	+ pub created: u64,
	188	+ pub model: String,
	189	+ pub choices: Vec<ChatCompletionChoice>,
	190	+ pub usage: Usage,
	191	+}
	192	+
	193	+/// A single chat completion choice.
	194	+#[derive(Debug, Clone, Serialize)]
	195	+pub struct ChatCompletionChoice {
	196	+ pub index: usize,
	197	+ pub message: ChatMessageOut,
	198	+ pub finish_reason: String,
	199	+}
	200	+
	201	+/// An assistant message in a completion response.
	202	+#[derive(Debug, Clone, Serialize)]
	203	+pub struct ChatMessageOut {
	204	+ pub role: String,
	205	+ pub content: String,
	206	+ #[serde(skip_serializing_if = "Option::is_none")]
	207	+ pub finish_reason: Option<String>,
	208	+}
	209	+
	210	+/// A streaming chat-completion chunk.
	211	+#[derive(Debug, Clone, Serialize)]
	212	+pub struct ChatCompletionChunk {
	213	+ pub id: String,
	214	+ pub object: &'static str,
	215	+ pub created: u64,
	216	+ pub model: String,
	217	+ pub choices: Vec<ChunkChoice>,
	218	+}
	219	+
	220	+/// A single choice delta within a streaming chunk.
	221	+#[derive(Debug, Clone, Serialize)]
	222	+pub struct ChunkChoice {
	223	+ pub index: usize,
	224	+ pub delta: Delta,
	225	+ #[serde(skip_serializing_if = "Option::is_none")]
	226	+ pub finish_reason: Option<String>,
	227	+}
	228	+
	229	+/// The incremental content of a streaming chunk.
	230	+#[derive(Debug, Clone, Serialize, Default)]
	231	+pub struct Delta {
	232	+ #[serde(skip_serializing_if = "Option::is_none")]
	233	+ pub role: Option<String>,
	234	+ #[serde(skip_serializing_if = "Option::is_none")]
	235	+ pub content: Option<String>,
	236	+}
	237	+
	238	+/// `POST /v1/completions` (non-streaming) response.
	239	+#[derive(Debug, Clone, Serialize)]
	240	+pub struct CompletionResponse {
	241	+ pub id: String,
	242	+ pub object: &'static str,
	243	+ pub created: u64,
	244	+ pub model: String,
	245	+ pub choices: Vec<CompletionChoice>,
	246	+ pub usage: Usage,
	247	+}
	248	+
	249	+/// A single text-completion choice.
	250	+#[derive(Debug, Clone, Serialize)]
	251	+pub struct CompletionChoice {
	252	+ pub text: String,
	253	+ pub index: usize,
	254	+ pub finish_reason: String,
	255	+}
	256	+
	257	+/// A single embedding vector entry.
	258	+#[derive(Debug, Clone, Serialize)]
	259	+pub struct EmbeddingObject {
	260	+ pub object: &'static str,
	261	+ pub index: usize,
	262	+ pub embedding: Vec<f32>,
	263	+}
	264	+
	265	+/// `POST /v1/embeddings` response.
	266	+#[derive(Debug, Clone, Serialize)]
	267	+pub struct EmbeddingResponse {
	268	+ pub object: &'static str,
	269	+ pub data: Vec<EmbeddingObject>,
	270	+ pub model: String,
	271	+ pub usage: Usage,
	272	+}
	273	+
	274	+/// OpenAI-style error envelope.
	275	+#[derive(Debug, Clone, Serialize)]
	276	+pub struct ErrorResponse {
	277	+ pub error: ApiError,
	278	+}
	279	+
	280	+/// OpenAI-style error body.
	281	+#[derive(Debug, Clone, Serialize)]
	282	+pub struct ApiError {
	283	+ pub message: String,
	284	+ #[serde(rename = "type")]
	285	+ pub error_type: String,
	286	+ #[serde(default)]
	287	+ pub param: Option<String>,
	288	+ #[serde(default)]
	289	+ pub code: Option<String>,
	290	+}
	291	+
	292	+// ---------------------------------------------------------------------------
	293	+// Helpers
	294	+// ---------------------------------------------------------------------------
	295	+
	296	+const SSE_DONE: &str = "data: [DONE]\n\n";
	297	+
	298	+/// Map a finish reason to its OpenAI string.
	299	+pub fn finish_reason_str(r: FinishReason) -> &'static str {
	300	+ match r {
	301	+ FinishReason::Length => "length",
	302	+ FinishReason::Stop => "stop",
	303	+ FinishReason::Cancelled => "stop",
	304	+ }
	305	+}
	306	+
	307	+/// Flatten chat messages into a single prompt string.
	308	+pub fn prompt_from_messages(messages: &[ChatMessage]) -> String {
	309	+ let mut s = String::new();
	310	+ for m in messages {
	311	+ s.push_str(&m.role);
	312	+ s.push_str(": ");
	313	+ s.push_str(&m.content);
	314	+ s.push('\n');
	315	+ }
	316	+ s
	317	+}
	318	+
	319	+fn now_secs() -> u64 {
	320	+ SystemTime::now()
	321	+ .duration_since(UNIX_EPOCH)
	322	+ .map(\|d\| d.as_secs())
	323	+ .unwrap_or(0)
	324	+}
	325	+
	326	+fn short_id() -> String {
	327	+ let nanos = SystemTime::now()
	328	+ .duration_since(UNIX_EPOCH)
	329	+ .map(\|d\| d.as_nanos())
	330	+ .unwrap_or(0);
	331	+ format!("{:x}", nanos % 0xfffffffffffffff)
	332	+}
	333	+
	334	+/// Format a chunk as an SSE `data:` event.
	335	+pub fn sse_event(chunk: &ChatCompletionChunk) -> String {
	336	+ let json = serde_json::to_string(chunk).expect("chunk serializes");
	337	+ format!("data: {json}\n\n")
	338	+}
	339	+
	340	+/// Build an error SSE event from an API error.
	341	+pub fn sse_error(err: &ApiError) -> String {
	342	+ let env = ErrorResponse { error: err.clone() };
	343	+ let json = serde_json::to_string(&env).expect("error serializes");
	344	+ format!("data: {json}\n\n")
	345	+}
	346	+
	347	+/// Build a JSON error response string.
	348	+pub fn error_json(message: &str, error_type: &str, status: u16) -> (u16, String) {
	349	+ let env = ErrorResponse {
	350	+ error: ApiError {
	351	+ message: message.to_string(),
	352	+ error_type: error_type.to_string(),
	353	+ param: None,
	354	+ code: None,
	355	+ },
	356	+ };
	357	+ (status, serde_json::to_string(&env).unwrap_or_default())
	358	+}
	359	+
	360	+// ---------------------------------------------------------------------------
	361	+// Drivers
	362	+// ---------------------------------------------------------------------------
	363	+
	364	+/// Stream chat-completion tokens as OpenAI SSE events.
	365	+///
	366	+/// The returned stream is `'static` (it owns an `Arc<Engine>`), so it can be
	367	+/// handed straight to a streaming HTTP response body.
	368	+pub fn chat_completion_stream(
	369	+ engine: Arc<Engine>,
	370	+ req: ChatCompletionRequest,
	371	+) -> impl futures::Stream<Item = Result<bytes::Bytes, anyhow::Error>> + 'static {
	372	+ let model = req.model.clone();
	373	+ let prompt = prompt_from_messages(&req.messages);
	374	+ let tokens = MockModel::tokenize(&prompt);
	375	+ let params = SamplingParams {
	376	+ max_tokens: req.effective_max_tokens(),
	377	+ temperature: req.temperature.unwrap_or(1.0),
	378	+ top_p: req.top_p.unwrap_or(1.0),
	379	+ ..Default::default()
	380	+ };
	381	+ let id = format!("chatcmpl-{}", short_id());
	382	+ let created = now_secs();
	383	+
	384	+ async_stream::stream! {
	385	+ // First event carries the assistant role.
	386	+ let role_chunk = ChatCompletionChunk {
	387	+ id: id.clone(),
	388	+ object: "chat.completion.chunk",
	389	+ created,
	390	+ model: model.clone(),
	391	+ choices: vec![ChunkChoice {
	392	+ index: 0,
	393	+ delta: Delta { role: Some("assistant".into()), content: None },
	394	+ finish_reason: None,
	395	+ }],
	396	+ };
	397	+ yield Ok(bytes::Bytes::from(sse_event(&role_chunk)));
	398	+
	399	+ let mut rx = match engine.generate(&model, &tokens, params) {
	400	+ Ok(rx) => rx,
	401	+ Err(e) => {
	402	+ let err = ApiError {
	403	+ message: e.to_string(),
	404	+ error_type: "server_error".into(),
	405	+ param: None,
	406	+ code: None,
	407	+ };
	408	+ yield Ok(bytes::Bytes::from(sse_error(&err)));
	409	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	410	+ return;
	411	+ }
	412	+ };
	413	+
	414	+ let mut finished = false;
	415	+ while let Some(tok) = rx.recv().await {
	416	+ let finish_reason = tok.finish_reason.map(finish_reason_str).map(String::from);
	417	+ finished = finish_reason.is_some();
	418	+ let chunk = ChatCompletionChunk {
	419	+ id: id.clone(),
	420	+ object: "chat.completion.chunk",
	421	+ created,
	422	+ model: model.clone(),
	423	+ choices: vec![ChunkChoice {
	424	+ index: 0,
	425	+ delta: Delta { role: None, content: Some(tok.text) },
	426	+ finish_reason,
	427	+ }],
	428	+ };
	429	+ yield Ok(bytes::Bytes::from(sse_event(&chunk)));
	430	+ }
	431	+
	432	+ // Final event signals completion. If a stop token already ended the
	433	+ // stream, no extra finish chunk is needed.
	434	+ if !finished {
	435	+ let done_chunk = ChatCompletionChunk {
	436	+ id: id.clone(),
	437	+ object: "chat.completion.chunk",
	438	+ created,
	439	+ model: model.clone(),
	440	+ choices: vec![ChunkChoice {
	441	+ index: 0,
	442	+ delta: Delta { role: None, content: None },
	443	+ finish_reason: Some("length".into()),
	444	+ }],
	445	+ };
	446	+ yield Ok(bytes::Bytes::from(sse_event(&done_chunk)));
	447	+ }
	448	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	449	+ }
	450	+}
	451	+
	452	+/// Produce a full (non-streaming) chat completion.
	453	+pub async fn chat_completion(
	454	+ engine: Arc<Engine>,
	455	+ req: ChatCompletionRequest,
	456	+) -> Result<ChatCompletionResponse, ModelError> {
	457	+ let model = req.model.clone();
	458	+ let prompt = prompt_from_messages(&req.messages);
	459	+ let tokens = MockModel::tokenize(&prompt);
	460	+ let prompt_tokens = tokens.len();
	461	+ let params = SamplingParams {
	462	+ max_tokens: req.effective_max_tokens(),
	463	+ temperature: req.temperature.unwrap_or(1.0),
	464	+ top_p: req.top_p.unwrap_or(1.0),
	465	+ ..Default::default()
	466	+ };
	467	+
	468	+ let mut rx = engine.generate(&model, &tokens, params)?;
	469	+ let mut content = String::new();
	470	+ let mut finish = "length";
	471	+ let mut completion_tokens = 0usize;
	472	+ while let Some(tok) = rx.recv().await {
	473	+ content.push_str(&tok.text);
	474	+ completion_tokens += 1;
	475	+ if let Some(r) = tok.finish_reason {
	476	+ finish = finish_reason_str(r);
	477	+ }
	478	+ }
	479	+
	480	+ Ok(ChatCompletionResponse {
	481	+ id: format!("chatcmpl-{}", short_id()),
	482	+ object: "chat.completion",
	483	+ created: now_secs(),
	484	+ model,
	485	+ choices: vec![ChatCompletionChoice {
	486	+ index: 0,
	487	+ message: ChatMessageOut {
	488	+ role: "assistant".into(),
	489	+ content,
	490	+ finish_reason: None,
	491	+ },
	492	+ finish_reason: finish.into(),
	493	+ }],
	494	+ usage: Usage {
	495	+ prompt_tokens,
	496	+ completion_tokens,
	497	+ total_tokens: prompt_tokens + completion_tokens,
	498	+ },
	499	+ })
	500	+}
	501	+
	502	+/// Stream text-completion tokens as OpenAI SSE events.
	503	+pub fn completion_stream(
	504	+ engine: Arc<Engine>,
	505	+ req: CompletionRequest,
	506	+) -> impl futures::Stream<Item = Result<bytes::Bytes, anyhow::Error>> + 'static {
	507	+ let model = req.model.clone();
	508	+ let tokens = MockModel::tokenize(&req.prompt);
	509	+ let params = SamplingParams {
	510	+ max_tokens: req.max_tokens,
	511	+ temperature: req.temperature.unwrap_or(1.0),
	512	+ top_p: req.top_p.unwrap_or(1.0),
	513	+ ..Default::default()
	514	+ };
	515	+ let id = format!("cmpl-{}", short_id());
	516	+ let created = now_secs();
	517	+
	518	+ async_stream::stream! {
	519	+ let mut rx = match engine.generate(&model, &tokens, params) {
	520	+ Ok(rx) => rx,
	521	+ Err(e) => {
	522	+ let err = ApiError {
	523	+ message: e.to_string(),
	524	+ error_type: "server_error".into(),
	525	+ param: None,
	526	+ code: None,
	527	+ };
	528	+ yield Ok(bytes::Bytes::from(sse_error(&err)));
	529	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	530	+ return;
	531	+ }
	532	+ };
	533	+ let mut finished = false;
	534	+ while let Some(tok) = rx.recv().await {
	535	+ let finish_reason = tok.finish_reason.map(finish_reason_str).map(String::from);
	536	+ finished = finish_reason.is_some();
	537	+ let chunk = ChatCompletionChunk {
	538	+ id: id.clone(),
	539	+ object: "chat.completion.chunk",
	540	+ created,
	541	+ model: model.clone(),
	542	+ choices: vec![ChunkChoice {
	543	+ index: 0,
	544	+ delta: Delta { role: None, content: Some(tok.text) },
	545	+ finish_reason,
	546	+ }],
	547	+ };
	548	+ yield Ok(bytes::Bytes::from(sse_event(&chunk)));
	549	+ }
	550	+ if !finished {
	551	+ let done_chunk = ChatCompletionChunk {
	552	+ id: id.clone(),
	553	+ object: "chat.completion.chunk",
	554	+ created,
	555	+ model: model.clone(),
	556	+ choices: vec![ChunkChoice {
	557	+ index: 0,
	558	+ delta: Delta { role: None, content: None },
	559	+ finish_reason: Some("length".into()),
	560	+ }],
	561	+ };
	562	+ yield Ok(bytes::Bytes::from(sse_event(&done_chunk)));
	563	+ }
	564	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	565	+ }
	566	+}
	567	+
	568	+/// Produce a full (non-streaming) text completion.
	569	+pub async fn completion(
	570	+ engine: Arc<Engine>,
	571	+ req: CompletionRequest,
	572	+) -> Result<CompletionResponse, ModelError> {
	573	+ let model = req.model.clone();
	574	+ let tokens = MockModel::tokenize(&req.prompt);
	575	+ let prompt_tokens = tokens.len();
	576	+ let params = SamplingParams {
	577	+ max_tokens: req.max_tokens,
	578	+ temperature: req.temperature.unwrap_or(1.0),
	579	+ top_p: req.top_p.unwrap_or(1.0),
	580	+ ..Default::default()
	581	+ };
	582	+
	583	+ let mut rx = engine.generate(&model, &tokens, params)?;
	584	+ let mut text = String::new();
	585	+ let mut finish = "length";
	586	+ let mut completion_tokens = 0usize;
	587	+ while let Some(tok) = rx.recv().await {
	588	+ text.push_str(&tok.text);
	589	+ completion_tokens += 1;
	590	+ if let Some(r) = tok.finish_reason {
	591	+ finish = finish_reason_str(r);
	592	+ }
	593	+ }
	594	+
	595	+ Ok(CompletionResponse {
	596	+ id: format!("cmpl-{}", short_id()),
	597	+ object: "text_completion",
	598	+ created: now_secs(),
	599	+ model,
	600	+ choices: vec![CompletionChoice {
	601	+ text,
	602	+ index: 0,
	603	+ finish_reason: finish.into(),
	604	+ }],
	605	+ usage: Usage {
	606	+ prompt_tokens,
	607	+ completion_tokens,
	608	+ total_tokens: prompt_tokens + completion_tokens,
	609	+ },
	610	+ })
	611	+}
	612	+
	613	+/// Produce embeddings for the request inputs (mock: deterministic from tokens).
	614	+pub fn embeddings(engine: &Engine, req: EmbeddingRequest) -> Result<EmbeddingResponse, ModelError> {
	615	+ let _ = engine; // embeddings are model-independent in the mock path
	616	+ let texts = req.input.into_texts();
	617	+ let mut total_tokens = 0usize;
	618	+ let data = texts
	619	+ .into_iter()
	620	+ .enumerate()
	621	+ .map(\|(i, text)\| {
	622	+ let tokens = MockModel::tokenize(&text);
	623	+ total_tokens += tokens.len();
	624	+ // Deterministic pseudo-embedding: normalize token ids.
	625	+ let embedding: Vec<f32> = tokens.iter().map(\|&t\| (t as f32) / 1000.0 - 0.5).collect();
	626	+ EmbeddingObject {
	627	+ object: "embedding",
	628	+ index: i,
	629	+ embedding,
	630	+ }
	631	+ })
	632	+ .collect();
	633	+
	634	+ Ok(EmbeddingResponse {
	635	+ object: "list",
	636	+ data,
	637	+ model: req.model,
	638	+ usage: Usage {
	639	+ prompt_tokens: total_tokens,
	640	+ completion_tokens: 0,
	641	+ total_tokens,
	642	+ },
	643	+ })
	644	+}
	645	+
	646	+// ---------------------------------------------------------------------------
	647	+// Scheduler-backed drivers
	648	+// ---------------------------------------------------------------------------
	649	+
	650	+/// Produce a full (non-streaming) chat completion via the scheduler.
	651	+pub async fn scheduled_chat_completion(
	652	+ engine: Arc<SchedulerEngine>,
	653	+ req: ChatCompletionRequest,
	654	+) -> Result<ChatCompletionResponse, ModelError> {
	655	+ let model = req.model.clone();
	656	+ let prompt = prompt_from_messages(&req.messages);
	657	+ let tokens = MockModel::tokenize(&prompt);
	658	+ let prompt_tokens = tokens.len();
	659	+ let params = SamplingParams {
	660	+ max_tokens: req.effective_max_tokens(),
	661	+ temperature: req.temperature.unwrap_or(1.0),
	662	+ top_p: req.top_p.unwrap_or(1.0),
	663	+ ..Default::default()
	664	+ };
	665	+
	666	+ let mut rx = engine.generate(&model, &tokens, params)?;
	667	+ let mut content = String::new();
	668	+ let mut finish = "length";
	669	+ let mut completion_tokens = 0usize;
	670	+ while let Some(tok) = rx.recv().await {
	671	+ content.push_str(&tok.text);
	672	+ completion_tokens += 1;
	673	+ if let Some(r) = tok.finish_reason {
	674	+ finish = finish_reason_str(r);
	675	+ }
	676	+ }
	677	+
	678	+ Ok(ChatCompletionResponse {
	679	+ id: format!("chatcmpl-{}", short_id()),
	680	+ object: "chat.completion",
	681	+ created: now_secs(),
	682	+ model,
	683	+ choices: vec![ChatCompletionChoice {
	684	+ index: 0,
	685	+ message: ChatMessageOut {
	686	+ role: "assistant".into(),
	687	+ content,
	688	+ finish_reason: None,
	689	+ },
	690	+ finish_reason: finish.into(),
	691	+ }],
	692	+ usage: Usage {
	693	+ prompt_tokens,
	694	+ completion_tokens,
	695	+ total_tokens: prompt_tokens + completion_tokens,
	696	+ },
	697	+ })
	698	+}
	699	+
	700	+/// Stream chat-completion tokens as OpenAI SSE events via the scheduler.
	701	+pub fn scheduled_chat_completion_stream(
	702	+ engine: Arc<SchedulerEngine>,
	703	+ req: ChatCompletionRequest,
	704	+) -> impl futures::Stream<Item = Result<bytes::Bytes, anyhow::Error>> + 'static {
	705	+ let model = req.model.clone();
	706	+ let prompt = prompt_from_messages(&req.messages);
	707	+ let tokens = MockModel::tokenize(&prompt);
	708	+ let params = SamplingParams {
	709	+ max_tokens: req.effective_max_tokens(),
	710	+ temperature: req.temperature.unwrap_or(1.0),
	711	+ top_p: req.top_p.unwrap_or(1.0),
	712	+ ..Default::default()
	713	+ };
	714	+ let id = format!("chatcmpl-{}", short_id());
	715	+ let created = now_secs();
	716	+
	717	+ async_stream::stream! {
	718	+ let role_chunk = ChatCompletionChunk {
	719	+ id: id.clone(),
	720	+ object: "chat.completion.chunk",
	721	+ created,
	722	+ model: model.clone(),
	723	+ choices: vec![ChunkChoice {
	724	+ index: 0,
	725	+ delta: Delta { role: Some("assistant".into()), content: None },
	726	+ finish_reason: None,
	727	+ }],
	728	+ };
	729	+ yield Ok(bytes::Bytes::from(sse_event(&role_chunk)));
	730	+
	731	+ let mut rx = match engine.generate(&model, &tokens, params) {
	732	+ Ok(rx) => rx,
	733	+ Err(e) => {
	734	+ let err = ApiError {
	735	+ message: e.to_string(),
	736	+ error_type: "server_error".into(),
	737	+ param: None,
	738	+ code: None,
	739	+ };
	740	+ yield Ok(bytes::Bytes::from(sse_error(&err)));
	741	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	742	+ return;
	743	+ }
	744	+ };
	745	+
	746	+ let mut finished = false;
	747	+ while let Some(tok) = rx.recv().await {
	748	+ let finish_reason = tok.finish_reason.map(finish_reason_str).map(String::from);
	749	+ if finish_reason.is_some() {
	750	+ finished = true;
	751	+ }
	752	+ let chunk = ChatCompletionChunk {
	753	+ id: id.clone(),
	754	+ object: "chat.completion.chunk",
	755	+ created,
	756	+ model: model.clone(),
	757	+ choices: vec![ChunkChoice {
	758	+ index: 0,
	759	+ delta: Delta { role: None, content: Some(tok.text) },
	760	+ finish_reason,
	761	+ }],
	762	+ };
	763	+ yield Ok(bytes::Bytes::from(sse_event(&chunk)));
	764	+ }
	765	+
	766	+ if !finished {
	767	+ let done_chunk = ChatCompletionChunk {
	768	+ id: id.clone(),
	769	+ object: "chat.completion.chunk",
	770	+ created,
	771	+ model: model.clone(),
	772	+ choices: vec![ChunkChoice {
	773	+ index: 0,
	774	+ delta: Delta { role: None, content: None },
	775	+ finish_reason: Some("length".into()),
	776	+ }],
	777	+ };
	778	+ yield Ok(bytes::Bytes::from(sse_event(&done_chunk)));
	779	+ }
	780	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	781	+ }
	782	+}
	783	+
	784	+/// Produce a full (non-streaming) text completion via the scheduler.
	785	+pub async fn scheduled_completion(
	786	+ engine: Arc<SchedulerEngine>,
	787	+ req: CompletionRequest,
	788	+) -> Result<CompletionResponse, ModelError> {
	789	+ let model = req.model.clone();
	790	+ let tokens = MockModel::tokenize(&req.prompt);
	791	+ let prompt_tokens = tokens.len();
	792	+ let params = SamplingParams {
	793	+ max_tokens: req.max_tokens,
	794	+ temperature: req.temperature.unwrap_or(1.0),
	795	+ top_p: req.top_p.unwrap_or(1.0),
	796	+ ..Default::default()
	797	+ };
	798	+
	799	+ let mut rx = engine.generate(&model, &tokens, params)?;
	800	+ let mut text = String::new();
	801	+ let mut finish = "length";
	802	+ let mut completion_tokens = 0usize;
	803	+ while let Some(tok) = rx.recv().await {
	804	+ text.push_str(&tok.text);
	805	+ completion_tokens += 1;
	806	+ if let Some(r) = tok.finish_reason {
	807	+ finish = finish_reason_str(r);
	808	+ }
	809	+ }
	810	+
	811	+ Ok(CompletionResponse {
	812	+ id: format!("cmpl-{}", short_id()),
	813	+ object: "text_completion",
	814	+ created: now_secs(),
	815	+ model,
	816	+ choices: vec![CompletionChoice {
	817	+ text,
	818	+ index: 0,
	819	+ finish_reason: finish.into(),
	820	+ }],
	821	+ usage: Usage {
	822	+ prompt_tokens,
	823	+ completion_tokens,
	824	+ total_tokens: prompt_tokens + completion_tokens,
	825	+ },
	826	+ })
	827	+}
	828	+
	829	+/// Stream text-completion tokens as OpenAI SSE events via the scheduler.
	830	+pub fn scheduled_completion_stream(
	831	+ engine: Arc<SchedulerEngine>,
	832	+ req: CompletionRequest,
	833	+) -> impl futures::Stream<Item = Result<bytes::Bytes, anyhow::Error>> + 'static {
	834	+ let model = req.model.clone();
	835	+ let tokens = MockModel::tokenize(&req.prompt);
	836	+ let params = SamplingParams {
	837	+ max_tokens: req.max_tokens,
	838	+ temperature: req.temperature.unwrap_or(1.0),
	839	+ top_p: req.top_p.unwrap_or(1.0),
	840	+ ..Default::default()
	841	+ };
	842	+ let id = format!("cmpl-{}", short_id());
	843	+ let created = now_secs();
	844	+
	845	+ async_stream::stream! {
	846	+ let mut rx = match engine.generate(&model, &tokens, params) {
	847	+ Ok(rx) => rx,
	848	+ Err(e) => {
	849	+ let err = ApiError {
	850	+ message: e.to_string(),
	851	+ error_type: "server_error".into(),
	852	+ param: None,
	853	+ code: None,
	854	+ };
	855	+ yield Ok(bytes::Bytes::from(sse_error(&err)));
	856	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	857	+ return;
	858	+ }
	859	+ };
	860	+ let mut finished = false;
	861	+ while let Some(tok) = rx.recv().await {
	862	+ let finish_reason = tok.finish_reason.map(finish_reason_str).map(String::from);
	863	+ if finish_reason.is_some() {
	864	+ finished = true;
	865	+ }
	866	+ let chunk = ChatCompletionChunk {
	867	+ id: id.clone(),
	868	+ object: "chat.completion.chunk",
	869	+ created,
	870	+ model: model.clone(),
	871	+ choices: vec![ChunkChoice {
	872	+ index: 0,
	873	+ delta: Delta { role: None, content: Some(tok.text) },
	874	+ finish_reason,
	875	+ }],
	876	+ };
	877	+ yield Ok(bytes::Bytes::from(sse_event(&chunk)));
	878	+ }
	879	+ if !finished {
	880	+ let done_chunk = ChatCompletionChunk {
	881	+ id: id.clone(),
	882	+ object: "chat.completion.chunk",
	883	+ created,
	884	+ model: model.clone(),
	885	+ choices: vec![ChunkChoice {
	886	+ index: 0,
	887	+ delta: Delta { role: None, content: None },
	888	+ finish_reason: Some("length".into()),
	889	+ }],
	890	+ };
	891	+ yield Ok(bytes::Bytes::from(sse_event(&done_chunk)));
	892	+ }
	893	+ yield Ok(bytes::Bytes::from_static(SSE_DONE.as_bytes()));
	894	+ }
	895	+}
	896	+
	897	+/// List models registered in the engine.
	898	+pub fn model_list(engine: &Engine) -> ModelList {
	899	+ let data = engine
	900	+ .list_models()
	901	+ .into_iter()
	902	+ .map(\|id\| ModelObject {
	903	+ id,
	904	+ object: "model",
	905	+ created: now_secs(),
	906	+ owned_by: "justapi".into(),
	907	+ permission: vec![ModelPermission::default()],
	908	+ root: None,
	909	+ parent: None,
	910	+ })
	911	+ .collect();
	912	+ ModelList {
	913	+ object: "list",
	914	+ data,
	915	+ }
	916	+}
	917	+
	918	+// ---------------------------------------------------------------------------
	919	+// Tests
	920	+// ---------------------------------------------------------------------------
	921	+
	922	+#[cfg(test)]
	923	+mod tests {
	924	+ use super::*;
	925	+ use crate::engine::EngineDevice;
	926	+ use crate::model::FinishReason;
	927	+ use futures::StreamExt;
	928	+
	929	+ fn test_engine() -> Arc<Engine> {
	930	+ let engine = Arc::new(Engine::new(EngineDevice::Cpu).unwrap());
	931	+ engine.register_mock("mock");
	932	+ engine
	933	+ }
	934	+
	935	+ fn chat_req(stream: bool) -> ChatCompletionRequest {
	936	+ ChatCompletionRequest {
	937	+ model: "mock".into(),
	938	+ messages: vec![ChatMessage {
	939	+ role: "user".into(),
	940	+ content: "hello world".into(),
	941	+ name: None,
	942	+ }],
	943	+ temperature: Some(0.0),
	944	+ max_tokens: 10,
	945	+ max_completion_tokens: None,
	946	+ top_p: None,
	947	+ stream,
	948	+ stop: None,
	949	+ user: None,
	950	+ }
	951	+ }
	952	+
	953	+ #[test]
	954	+ fn protocol_types_roundtrip() {
	955	+ let req = chat_req(false);
	956	+ let json = serde_json::to_string(&req).unwrap();
	957	+ let back: ChatCompletionRequest = serde_json::from_str(&json).unwrap();
	958	+ assert_eq!(back.model, "mock");
	959	+ assert_eq!(back.messages.len(), 1);
	960	+ }
	961	+
	962	+ #[tokio::test]
	963	+ async fn chat_completion_non_streaming() {
	964	+ let engine = test_engine();
	965	+ let resp = chat_completion(engine, chat_req(false)).await.unwrap();
	966	+ assert_eq!(resp.object, "chat.completion");
	967	+ assert_eq!(resp.choices.len(), 1);
	968	+ assert_eq!(resp.choices[0].message.role, "assistant");
	969	+ assert!(!resp.choices[0].message.content.is_empty());
	970	+ assert_eq!(resp.choices[0].finish_reason, "length");
	971	+ assert!(resp.usage.total_tokens > 0);
	972	+ }
	973	+
	974	+ #[tokio::test]
	975	+ async fn chat_completion_streaming_format() {
	976	+ let engine = test_engine();
	977	+ let mut stream = Box::pin(chat_completion_stream(engine, chat_req(true)));
	978	+
	979	+ // First event: role.
	980	+ let first = stream.next().await.unwrap().unwrap();
	981	+ let first = String::from_utf8(first.to_vec()).unwrap();
	982	+ assert!(first.starts_with("data: "));
	983	+ assert!(first.contains("\"role\":\"assistant\""));
	984	+
	985	+ // Subsequent events: content deltas, then finish, then [DONE].
	986	+ let mut saw_content = false;
	987	+ let mut saw_done = false;
	988	+ while let Some(item) = stream.next().await {
	989	+ let s = String::from_utf8(item.unwrap().to_vec()).unwrap();
	990	+ if s.contains("\"content\"") {
	991	+ saw_content = true;
	992	+ }
	993	+ if s == SSE_DONE {
	994	+ saw_done = true;
	995	+ }
	996	+ }
	997	+ assert!(saw_content);
	998	+ assert!(saw_done);
	999	+ }
	1000	+
	1001	+ #[tokio::test]
	1002	+ async fn text_completion_works() {
	1003	+ let engine = test_engine();
	1004	+ let req = CompletionRequest {
	1005	+ model: "mock".into(),
	1006	+ prompt: "hi".into(),
	1007	+ temperature: Some(0.0),
	1008	+ max_tokens: 8,
	1009	+ top_p: None,
	1010	+ stream: false,
	1011	+ stop: None,
	1012	+ };
	1013	+ let resp = completion(engine, req).await.unwrap();
	1014	+ assert_eq!(resp.object, "text_completion");
	1015	+ assert!(!resp.choices[0].text.is_empty());
	1016	+ }
	1017	+
	1018	+ #[test]
	1019	+ fn embeddings_work() {
	1020	+ let engine = test_engine();
	1021	+ let req = EmbeddingRequest {
	1022	+ model: "mock".into(),
	1023	+ input: EmbeddingInput::Many(vec!["a".into(), "bb".into()]),
	1024	+ encoding_format: None,
	1025	+ };
	1026	+ let resp = embeddings(&engine, req).unwrap();
	1027	+ assert_eq!(resp.data.len(), 2);
	1028	+ assert_eq!(resp.data[0].embedding.len(), 1);
	1029	+ assert_eq!(resp.data[1].embedding.len(), 2);
	1030	+ }
	1031	+
	1032	+ #[test]
	1033	+ fn model_list_reflects_registry() {
	1034	+ let engine = test_engine();
	1035	+ let list = model_list(&engine);
	1036	+ assert_eq!(list.object, "list");
	1037	+ assert!(list.data.iter().any(\|m\| m.id == "mock"));
	1038	+ }
	1039	+
	1040	+ #[test]
	1041	+ fn finish_reason_mapping() {
	1042	+ assert_eq!(finish_reason_str(FinishReason::Stop), "stop");
	1043	+ assert_eq!(finish_reason_str(FinishReason::Length), "length");
	1044	+ assert_eq!(finish_reason_str(FinishReason::Cancelled), "stop");
	1045	+ }
	1046	+
	1047	+ #[test]
	1048	+ fn prompt_from_messages_flattened() {
	1049	+ let prompt = prompt_from_messages(&[
	1050	+ ChatMessage {
	1051	+ role: "system".into(),
	1052	+ content: "be nice".into(),
	1053	+ name: None,
	1054	+ },
	1055	+ ChatMessage {
	1056	+ role: "user".into(),
	1057	+ content: "hi".into(),
	1058	+ name: None,
	1059	+ },
	1060	+ ]);
	1061	+ assert!(prompt.contains("system: be nice"));
	1062	+ assert!(prompt.contains("user: hi"));
	1063	+ }
	1064	+}
0	1065

▼crates/justapi-inference/src/pd.rs

+349-0

@@ -0,0 +1,349 @@
	1	+//! Disaggregated prefill/decode scheduler.
	2	+//!
	3	+//! Owns two independent [`Scheduler`] instances — one for prefill (compute-heavy,
	4	+//! prompt processing) and one for decode (memory-bandwidth-bound, token generation).
	5	+//! Sequences that finish prefill are transferred from the prefill pool to the
	6	+//! decode pool via the KV-cache transfer protocol (see
	7	+//! [`Scheduler::take_completed_prefill`] and [`Scheduler::admit_transferred`]).
	8	+//!
	9	+//! This enables independent scaling of prefill and decode capacity:
	10	+//! - Prefill pool: optimised for a few sequences at a time (high compute).
	11	+//! - Decode pool: many sequences batched together for high throughput.
	12	+
	13	+use crate::scheduler::{NewRequest, Schedule, Scheduler, SchedulerConfig, SchedulerStats};
	14	+
	15	+// ---------------------------------------------------------------------------
	16	+// Statistics
	17	+// ---------------------------------------------------------------------------
	18	+
	19	+/// Combined statistics for the disaggregated prefill/decode scheduler.
	20	+#[derive(Debug, Clone)]
	21	+pub struct PdStats {
	22	+ /// Stats from the prefill scheduler.
	23	+ pub prefill: SchedulerStats,
	24	+ /// Stats from the decode scheduler.
	25	+ pub decode: SchedulerStats,
	26	+ /// Total prompt tokens transferred to the decode pool across all
	27	+ /// completed sequences.
	28	+ pub total_transferred_tokens: usize,
	29	+}
	30	+
	31	+// ---------------------------------------------------------------------------
	32	+// PdScheduler
	33	+// ---------------------------------------------------------------------------
	34	+
	35	+/// Scheduler that owns independent prefill and decode pools.
	36	+///
	37	+/// The caller drives the disaggregated loop:
	38	+///
	39	+/// ```ignore
	40	+/// loop {
	41	+/// // 1. Schedule prefill work.
	42	+/// let pre_sched = pd.schedule_prefill();
	43	+/// // engine.forward(pre_sched) ...
	44	+/// for step in &pre_sched.prefill {
	45	+/// pd.on_prefill_complete(step.seq_id, 1);
	46	+/// }
	47	+///
	48	+/// // 2. Transfer fully-prefilled sequences to decode pool.
	49	+/// pd.transfer_completed();
	50	+///
	51	+/// // 3. Schedule decode work.
	52	+/// let dec_sched = pd.schedule_decode();
	53	+/// // engine.forward(dec_sched) ...
	54	+/// for seq_id in &dec_sched.decode {
	55	+/// pd.on_decode_complete(*seq_id, 1);
	56	+/// }
	57	+/// }
	58	+/// ```
	59	+pub struct PdScheduler {
	60	+ prefill: Scheduler,
	61	+ decode: Scheduler,
	62	+ total_transferred_tokens: usize,
	63	+}
	64	+
	65	+impl PdScheduler {
	66	+ /// Create a new disaggregated scheduler with independent prefill and decode
	67	+ /// configurations and block pools.
	68	+ pub fn new(
	69	+ prefill_config: SchedulerConfig,
	70	+ prefill_pool: crate::kv_cache::KvBlockPool,
	71	+ decode_config: SchedulerConfig,
	72	+ decode_pool: crate::kv_cache::KvBlockPool,
	73	+ ) -> Self {
	74	+ Self {
	75	+ prefill: Scheduler::new(prefill_config, prefill_pool),
	76	+ decode: Scheduler::new(decode_config, decode_pool),
	77	+ total_transferred_tokens: 0,
	78	+ }
	79	+ }
	80	+
	81	+ // -- accessors ----------------------------------------------------------
	82	+
	83	+ /// Immutable reference to the prefill scheduler.
	84	+ pub fn prefill(&self) -> &Scheduler {
	85	+ &self.prefill
	86	+ }
	87	+
	88	+ /// Immutable reference to the decode scheduler.
	89	+ pub fn decode(&self) -> &Scheduler {
	90	+ &self.decode
	91	+ }
	92	+
	93	+ /// Number of tokens transferred from prefill to decode pool so far.
	94	+ pub fn total_transferred_tokens(&self) -> usize {
	95	+ self.total_transferred_tokens
	96	+ }
	97	+
	98	+ /// Combined statistics snapshot.
	99	+ pub fn stats(&self) -> PdStats {
	100	+ PdStats {
	101	+ prefill: self.prefill.stats(),
	102	+ decode: self.decode.stats(),
	103	+ total_transferred_tokens: self.total_transferred_tokens,
	104	+ }
	105	+ }
	106	+
	107	+ /// Return `true` when both schedulers are idle (no running or waiting
	108	+ /// sequences).
	109	+ pub fn is_idle(&self) -> bool {
	110	+ self.prefill.is_idle() && self.decode.is_idle()
	111	+ }
	112	+
	113	+ // -- request admission --------------------------------------------------
	114	+
	115	+ /// Enqueue a new generation request into the prefill scheduler.
	116	+ pub fn add_request(&mut self, req: NewRequest) {
	117	+ self.prefill.add_request(req);
	118	+ }
	119	+
	120	+ // -- prefill scheduling -------------------------------------------------
	121	+
	122	+ /// Schedule prefill work for the next forward step.
	123	+ ///
	124	+ /// Delegates to the prefill [`Scheduler::schedule`].
	125	+ pub fn schedule_prefill(&mut self) -> Schedule {
	126	+ self.prefill.schedule()
	127	+ }
	128	+
	129	+ /// Mark a sequence's prefill step as complete.
	130	+ ///
	131	+ /// You must call [`PdScheduler::transfer_completed`] after all prefill
	132	+ /// steps have been reported to move fully-prefilled sequences into the
	133	+ /// decode pool.
	134	+ pub fn on_prefill_complete(&mut self, seq_id: u64, num_new_tokens: usize) {
	135	+ self.prefill.on_step_complete(seq_id, num_new_tokens);
	136	+ }
	137	+
	138	+ /// Transfer all sequences that have finished prefill from the prefill
	139	+ /// scheduler to the decode scheduler.
	140	+ ///
	141	+ /// Returns the number of sequences transferred.
	142	+ pub fn transfer_completed(&mut self) -> usize {
	143	+ let transferred = self.prefill.take_completed_prefill();
	144	+ let count = transferred.len();
	145	+ for t in &transferred {
	146	+ self.total_transferred_tokens = self
	147	+ .total_transferred_tokens
	148	+ .saturating_add(t.num_prefilled);
	149	+ }
	150	+ for t in transferred {
	151	+ self.decode.admit_transferred(t);
	152	+ }
	153	+ count
	154	+ }
	155	+
	156	+ // -- decode scheduling --------------------------------------------------
	157	+
	158	+ /// Schedule decode work for the next forward step.
	159	+ ///
	160	+ /// Delegates to the decode [`Scheduler::schedule`].
	161	+ pub fn schedule_decode(&mut self) -> Schedule {
	162	+ self.decode.schedule()
	163	+ }
	164	+
	165	+ /// Mark a sequence's decode step as complete.
	166	+ pub fn on_decode_complete(&mut self, seq_id: u64, num_new_tokens: usize) {
	167	+ self.decode.on_step_complete(seq_id, num_new_tokens);
	168	+ }
	169	+
	170	+ /// Immediately cancel a sequence from both schedulers (client
	171	+ /// disconnected, error, etc.).
	172	+ pub fn cancel(&mut self, seq_id: u64) {
	173	+ self.prefill.cancel(seq_id);
	174	+ self.decode.cancel(seq_id);
	175	+ }
	176	+}
	177	+
	178	+// ---------------------------------------------------------------------------
	179	+// Tests
	180	+// ---------------------------------------------------------------------------
	181	+
	182	+#[cfg(test)]
	183	+mod tests {
	184	+ use super::*;
	185	+ use crate::kv_cache::KvBlockPool;
	186	+ use crate::model::SamplingParams;
	187	+
	188	+ fn pool(capacity: usize) -> KvBlockPool {
	189	+ KvBlockPool::new(capacity)
	190	+ }
	191	+
	192	+ fn prefill_cfg() -> SchedulerConfig {
	193	+ SchedulerConfig {
	194	+ max_num_seqs: 1,
	195	+ max_num_batched_tokens: 64,
	196	+ max_seq_len: 128,
	197	+ chunked_prefill_size: None,
	198	+ }
	199	+ }
	200	+
	201	+ fn decode_cfg() -> SchedulerConfig {
	202	+ SchedulerConfig {
	203	+ max_num_seqs: 8,
	204	+ max_num_batched_tokens: 128,
	205	+ max_seq_len: 128,
	206	+ chunked_prefill_size: None,
	207	+ }
	208	+ }
	209	+
	210	+ fn short_req(id: u64) -> NewRequest {
	211	+ NewRequest {
	212	+ id,
	213	+ prompt: (0..8).collect(),
	214	+ sampling_params: SamplingParams {
	215	+ max_tokens: 16,
	216	+ ..Default::default()
	217	+ },
	218	+ prefix_cached_tokens: 0,
	219	+ cached_blocks: Vec::new(),
	220	+ }
	221	+ }
	222	+
	223	+ #[test]
	224	+ fn pd_starts_idle() {
	225	+ let pd = PdScheduler::new(prefill_cfg(), pool(64), decode_cfg(), pool(128));
	226	+ assert!(pd.is_idle());
	227	+ assert_eq!(pd.total_transferred_tokens(), 0);
	228	+ assert_eq!(pd.prefill().num_running(), 0);
	229	+ assert_eq!(pd.decode().num_running(), 0);
	230	+ }
	231	+
	232	+ #[test]
	233	+ fn pd_single_request_prefill_transfer_decode() {
	234	+ let mut pd = PdScheduler::new(prefill_cfg(), pool(64), decode_cfg(), pool(128));
	235	+ pd.add_request(short_req(1));
	236	+
	237	+ // ---- prefill step ----
	238	+ let s = pd.schedule_prefill();
	239	+ assert_eq!(s.prefill.len(), 1);
	240	+ assert!(s.decode.is_empty());
	241	+ pd.on_prefill_complete(s.prefill[0].seq_id, 1);
	242	+
	243	+ // ---- transfer ----
	244	+ let n = pd.transfer_completed();
	245	+ assert_eq!(n, 1);
	246	+ assert!(pd.prefill().is_idle());
	247	+ assert_eq!(pd.total_transferred_tokens(), 8);
	248	+
	249	+ // ---- decode step ----
	250	+ let s = pd.schedule_decode();
	251	+ assert!(s.prefill.is_empty());
	252	+ assert_eq!(s.decode.len(), 1);
	253	+ pd.on_decode_complete(s.decode[0], 1);
	254	+ }
	255	+
	256	+ #[test]
	257	+ fn pd_independent_pool_sizing() {
	258	+ // Prefill admits 1 at a time; decode can run 8 concurrently.
	259	+ let mut pd = PdScheduler::new(prefill_cfg(), pool(64), decode_cfg(), pool(128));
	260	+
	261	+ pd.add_request(short_req(10));
	262	+ pd.add_request(short_req(20));
	263	+ pd.add_request(short_req(30));
	264	+
	265	+ // ---- transfer 1st request ----
	266	+ let s = pd.schedule_prefill();
	267	+ assert_eq!(s.prefill.len(), 1);
	268	+ pd.on_prefill_complete(s.prefill[0].seq_id, 1);
	269	+ pd.transfer_completed();
	270	+
	271	+ // ---- transfer 2nd request ----
	272	+ let s = pd.schedule_prefill();
	273	+ assert_eq!(s.prefill.len(), 1);
	274	+ pd.on_prefill_complete(s.prefill[0].seq_id, 1);
	275	+ pd.transfer_completed();
	276	+
	277	+ // Decode pool has 2 sequences now.
	278	+ let s = pd.schedule_decode();
	279	+ assert_eq!(s.decode.len(), 2);
	280	+
	281	+ // 3rd request is still being prefilled (admitted one at a time).
	282	+ let s = pd.schedule_prefill();
	283	+ assert_eq!(s.prefill.len(), 1);
	284	+ }
	285	+
	286	+ #[test]
	287	+ fn pd_total_transferred_tokens_accounting() {
	288	+ let mut pd = PdScheduler::new(prefill_cfg(), pool(64), decode_cfg(), pool(128));
	289	+
	290	+ // 16-token prompt.
	291	+ pd.add_request(NewRequest {
	292	+ id: 1,
	293	+ prompt: (0..16).collect(),
	294	+ sampling_params: SamplingParams {
	295	+ max_tokens: 8,
	296	+ ..Default::default()
	297	+ },
	298	+ prefix_cached_tokens: 0,
	299	+ cached_blocks: Vec::new(),
	300	+ });
	301	+ let s = pd.schedule_prefill();
	302	+ pd.on_prefill_complete(s.prefill[0].seq_id, 1);
	303	+ pd.transfer_completed();
	304	+ assert_eq!(pd.total_transferred_tokens(), 16);
	305	+
	306	+ // 8-token prompt.
	307	+ pd.add_request(NewRequest {
	308	+ id: 2,
	309	+ prompt: (0..8).collect(),
	310	+ sampling_params: SamplingParams {
	311	+ max_tokens: 8,
	312	+ ..Default::default()
	313	+ },
	314	+ prefix_cached_tokens: 0,
	315	+ cached_blocks: Vec::new(),
	316	+ });
	317	+ let s = pd.schedule_prefill();
	318	+ pd.on_prefill_complete(s.prefill[0].seq_id, 1);
	319	+ pd.transfer_completed();
	320	+ assert_eq!(pd.total_transferred_tokens(), 24);
	321	+ }
	322	+
	323	+ #[test]
	324	+ fn pd_pool_block_allocation_release() {
	325	+ let prefill_pool = pool(256);
	326	+ let decode_pool = pool(256);
	327	+
	328	+ // Prefill pool: nothing in-use before.
	329	+ assert_eq!(prefill_pool.in_use(), 0);
	330	+ // Decode pool: nothing in-use before.
	331	+ assert_eq!(decode_pool.in_use(), 0);
	332	+
	333	+ let mut pd = PdScheduler::new(prefill_cfg(), prefill_pool, decode_cfg(), decode_pool);
	334	+ pd.add_request(short_req(1));
	335	+ let s = pd.schedule_prefill();
	336	+
	337	+ // At this point the prefill pool should have blocks in-use.
	338	+ assert!(pd.prefill().pool().in_use() > 0);
	339	+
	340	+ pd.on_prefill_complete(s.prefill[0].seq_id, 1);
	341	+ pd.transfer_completed();
	342	+
	343	+ // Prefill pool released all its blocks (now cached, not in-use).
	344	+ assert_eq!(pd.prefill().pool().in_use(), 0);
	345	+
	346	+ // Decode pool allocated blocks for the transferred sequence.
	347	+ assert!(pd.decode().pool().in_use() > 0);
	348	+ }
	349	+}
0	350

▼crates/justapi-inference/src/radix_cache.rs

+776-0

@@ -0,0 +1,776 @@
	1	+//! RadixAttention-style prefix caching (SGLang approach).
	2	+//!
	3	+//! A radix tree of token sequences where each node holds a block-aligned run
	4	+//! of KV-cache blocks. Sequences that share a prompt prefix share the internal
	5	+//! nodes of the tree, so the same KV blocks are stored (and reused) once. This
	6	+//! is strictly more memory-efficient than the flat hash-map [`crate::PrefixCache`]:
	7	+//! common prefixes are merged instead of duplicated, and eviction is LRU over
	8	+//! the leaf paths (least-recently-used prompt variants are dropped first).
	9	+//!
	10	+//! Nodes are block-aligned: a node's `tokens` length is always a multiple of
	11	+//! [`BLOCK_SIZE`], and `blocks` has one [`BlockId`] per block. The tree itself
	12	+//! never frees physical blocks — [`RadixPrefixCache::evict`] returns the
	13	+//! [`BlockId`]s of least-recently-used leaves as a hint; the caller (the
	14	+//! [`crate::KvBlockPool`]) makes the final recycle decision and skips any block
	15	+//! still held by a live sequence.
	16	+
	17	+use std::cell::Cell;
	18	+use std::collections::HashMap;
	19	+
	20	+use crate::kv_cache::{BlockId, BLOCK_SIZE};
	21	+
	22	+/// A node in the radix prefix tree.
	23	+#[derive(Debug, Clone)]
	24	+struct RadixNode {
	25	+ /// Edge label: KV-covered tokens for this node (multiple of `BLOCK_SIZE`).
	26	+ tokens: Vec<u32>,
	27	+ /// One KV block per `BLOCK_SIZE` tokens.
	28	+ blocks: Vec<BlockId>,
	29	+ /// Child nodes keyed by the first token of their edge.
	30	+ children: HashMap<u32, RadixNode>,
	31	+ /// Last-access timestamp (for LRU eviction).
	32	+ last_access: u64,
	33	+ /// Content hash of `tokens` (FNV-1a). The cache key: a lookup must match
	34	+ /// both the exact token run and this hash, so a corrupted or hash-only
	35	+ /// (colliding) match is rejected instead of returning the wrong KV blocks.
	36	+ token_hash: u64,
	37	+}
	38	+
	39	+/// FNV-1a 64-bit hash of a token slice. KV-cache content is determined solely
	40	+/// by the token ids, so this is the collision-resistant key for a prefix.
	41	+fn hash_tokens(tokens: &[u32]) -> u64 {
	42	+ let mut h = 0xcbf29ce484222325u64;
	43	+ for &t in tokens {
	44	+ h ^= t as u64;
	45	+ h = h.wrapping_mul(0x100000001b3);
	46	+ }
	47	+ h
	48	+}
	49	+
	50	+impl RadixNode {
	51	+ fn new(tokens: Vec<u32>, blocks: Vec<BlockId>, now: u64) -> Self {
	52	+ let token_hash = hash_tokens(&tokens);
	53	+ Self {
	54	+ tokens,
	55	+ blocks,
	56	+ children: HashMap::new(),
	57	+ last_access: now,
	58	+ token_hash,
	59	+ }
	60	+ }
	61	+
	62	+ fn is_leaf(&self) -> bool {
	63	+ self.children.is_empty()
	64	+ }
	65	+}
	66	+
	67	+/// Snapshot of radix-prefix-cache effectiveness.
	68	+#[derive(Debug, Clone, Copy, Default)]
	69	+pub struct RadixPrefixCacheStats {
	70	+ /// Number of nodes in the tree (excluding the root).
	71	+ pub nodes: usize,
	72	+ /// Number of successful prefix lookups.
	73	+ pub hits: u64,
	74	+ /// Number of failed prefix lookups.
	75	+ pub misses: u64,
	76	+ /// Number of lookups whose structural (token) match was rejected because
	77	+ /// the content hash disagreed — i.e. a collision or tree-corruption guard
	78	+ /// firing. Should stay at 0 in correct operation.
	79	+ pub hash_collisions: u64,
	80	+ /// Total tokens matched across all hits (the reuse benefit — each matched
	81	+ /// token is a KV block that did not need recomputation).
	82	+ pub tokens_saved: u64,
	83	+}
	84	+
	85	+/// Radix-tree prefix cache for cross-request KV-block reuse.
	86	+pub struct RadixPrefixCache {
	87	+ root: RadixNode,
	88	+ hits: Cell<u64>,
	89	+ misses: Cell<u64>,
	90	+ hash_collisions: Cell<u64>,
	91	+ tokens_saved: Cell<u64>,
	92	+ clock: Cell<u64>,
	93	+}
	94	+
	95	+impl RadixPrefixCache {
	96	+ /// Create an empty radix prefix cache.
	97	+ pub fn new() -> Self {
	98	+ Self {
	99	+ root: RadixNode::new(vec![], vec![], 0),
	100	+ hits: Cell::new(0),
	101	+ misses: Cell::new(0),
	102	+ hash_collisions: Cell::new(0),
	103	+ tokens_saved: Cell::new(0),
	104	+ clock: Cell::new(0),
	105	+ }
	106	+ }
	107	+
	108	+ fn tick(&self) -> u64 {
	109	+ let t = self.clock.get() + 1;
	110	+ self.clock.set(t);
	111	+ t
	112	+ }
	113	+
	114	+ /// Number of nodes in the tree (excluding the root).
	115	+ pub fn nodes(&self) -> usize {
	116	+ count_nodes(&self.root)
	117	+ }
	118	+
	119	+ /// Total tokens currently cached across all nodes.
	120	+ pub fn cached_tokens(&self) -> usize {
	121	+ cached_tokens(&self.root)
	122	+ }
	123	+
	124	+ /// Insert a completed sequence's `(tokens, blocks)` into the tree, sharing
	125	+ /// any common prefix with existing sequences. `tokens.len()` must equal
	126	+ /// `blocks.len() * BLOCK_SIZE`.
	127	+ pub fn insert(&mut self, tokens: &[u32], blocks: &[BlockId]) {
	128	+ assert_eq!(
	129	+ tokens.len(),
	130	+ blocks.len() * BLOCK_SIZE,
	131	+ "block misalignment"
	132	+ );
	133	+ let now = self.tick();
	134	+ insert_node(&mut self.root, tokens, blocks, now);
	135	+ }
	136	+
	137	+ /// Look up the longest matching prefix for `tokens`.
	138	+ ///
	139	+ /// Returns `(matched_tokens, block_ids)` for the longest prefix that exists
	140	+ /// in the tree, or `None` if no prefix matches. A match is only accepted
	141	+ /// when the accumulated content hash of the matched tokens equals a fresh
	142	+ /// recomputation — this is the collision-safe cache key: even if the tree
	143	+ /// were corrupted or a future change matched by hash alone, a hash
	144	+ /// disagreement rejects the match instead of returning the wrong KV blocks.
	145	+ /// Matched nodes have their last-access timestamp refreshed (LRU).
	146	+ pub fn lookup(&mut self, tokens: &[u32]) -> Option<(usize, Vec<BlockId>)> {
	147	+ let mut out = (0usize, Vec::new(), false);
	148	+ let now = self.tick();
	149	+ lookup_node(&mut self.root, tokens, &mut out, now);
	150	+ if out.2 {
	151	+ // A structural match was found but its content hash diverged from
	152	+ // the prompt tokens — a corruption or collision. Reject it.
	153	+ self.hash_collisions.set(self.hash_collisions.get() + 1);
	154	+ self.misses.set(self.misses.get() + 1);
	155	+ None
	156	+ } else if out.0 > 0 {
	157	+ self.hits.set(self.hits.get() + 1);
	158	+ self.tokens_saved
	159	+ .set(self.tokens_saved.get() + out.0 as u64);
	160	+ Some((out.0, out.1))
	161	+ } else {
	162	+ self.misses.set(self.misses.get() + 1);
	163	+ None
	164	+ }
	165	+ }
	166	+
	167	+ /// Walk the tree and assert every node's stored `token_hash` equals the hash
	168	+ /// of its tokens. Used by tests/debugging to prove the cache keys are
	169	+ /// internally consistent after inserts and splits.
	170	+ pub fn verify_hashes(&self) -> bool {
	171	+ fn check(node: &RadixNode) -> bool {
	172	+ if node.token_hash != hash_tokens(&node.tokens) {
	173	+ return false;
	174	+ }
	175	+ node.children.values().all(check)
	176	+ }
	177	+ check(&self.root)
	178	+ }
	179	+
	180	+ /// Evict up to `k` least-recently-used leaf blocks, returning the freed
	181	+ /// [`BlockId`]s. Leaf nodes are removed from the tree; the caller is
	182	+ /// responsible for actually recycling any blocks still held by a live
	183	+ /// sequence. Returns fewer than `k` ids if the tree is exhausted.
	184	+ pub fn evict(&mut self, k: usize) -> Vec<BlockId> {
	185	+ let mut freed = Vec::new();
	186	+ while freed.len() < k {
	187	+ match find_lru_leaf(&self.root) {
	188	+ Some((path, blocks)) => {
	189	+ remove_path(&mut self.root, &path);
	190	+ freed.extend(blocks);
	191	+ }
	192	+ None => break,
	193	+ }
	194	+ }
	195	+ freed
	196	+ }
	197	+
	198	+ /// Like [`RadixPrefixCache::evict`] but only evicts leaf paths whose blocks
	199	+ /// all satisfy `keep`. This lets the caller (the scheduler) skip blocks that
	200	+ /// are currently referenced by a live sequence, so eviction never frees a
	201	+ /// block still in use. Leaves whose blocks fail `keep` are left in the tree.
	202	+ ///
	203	+ /// Leaves are considered in least-recently-used order; eviction stops once
	204	+ /// `k` freeable blocks have been collected (or the tree is exhausted).
	205	+ pub fn evict_filter<F>(&mut self, k: usize, mut keep: F) -> Vec<BlockId>
	206	+ where
	207	+ F: FnMut(&[BlockId]) -> bool,
	208	+ {
	209	+ let mut leaves = Vec::new();
	210	+ collect_leaves(&self.root, &mut Vec::new(), &mut leaves);
	211	+ // Least-recently-used first.
	212	+ leaves.sort_by_key(\|(access, _, _)\| *access);
	213	+
	214	+ let mut freed = Vec::new();
	215	+ for (_, path, blocks) in leaves {
	216	+ if freed.len() >= k {
	217	+ break;
	218	+ }
	219	+ if keep(&blocks) {
	220	+ remove_path(&mut self.root, &path);
	221	+ freed.extend(blocks);
	222	+ }
	223	+ }
	224	+ freed
	225	+ }
	226	+
	227	+ /// Cache effectiveness snapshot.
	228	+ pub fn stats(&self) -> RadixPrefixCacheStats {
	229	+ RadixPrefixCacheStats {
	230	+ nodes: self.nodes(),
	231	+ hits: self.hits.get(),
	232	+ misses: self.misses.get(),
	233	+ hash_collisions: self.hash_collisions.get(),
	234	+ tokens_saved: self.tokens_saved.get(),
	235	+ }
	236	+ }
	237	+}
	238	+
	239	+impl Default for RadixPrefixCache {
	240	+ fn default() -> Self {
	241	+ Self::new()
	242	+ }
	243	+}
	244	+
	245	+// ---------------------------------------------------------------------------
	246	+// Internal helpers
	247	+// ---------------------------------------------------------------------------
	248	+
	249	+fn common_prefix_len(a: &[u32], b: &[u32]) -> usize {
	250	+ let mut i = 0;
	251	+ while i < a.len() && i < b.len() && a[i] == b[i] {
	252	+ i += 1;
	253	+ }
	254	+ i
	255	+}
	256	+
	257	+fn count_nodes(node: &RadixNode) -> usize {
	258	+ let mut n = node.children.len();
	259	+ for child in node.children.values() {
	260	+ n += count_nodes(child);
	261	+ }
	262	+ n
	263	+}
	264	+
	265	+fn cached_tokens(node: &RadixNode) -> usize {
	266	+ let mut t = node.tokens.len();
	267	+ for child in node.children.values() {
	268	+ t += cached_tokens(child);
	269	+ }
	270	+ t
	271	+}
	272	+
	273	+/// Recursively insert `(tokens, blocks)` under `node`, splitting existing
	274	+/// edges where the incoming path diverges.
	275	+fn insert_node(node: &mut RadixNode, tokens: &[u32], blocks: &[BlockId], now: u64) {
	276	+ if tokens.is_empty() {
	277	+ node.last_access = now;
	278	+ return;
	279	+ }
	280	+ let first = tokens[0];
	281	+
	282	+ // Does a child edge start with the same token?
	283	+ if let Some(child) = node.children.get_mut(&first) {
	284	+ let cp = common_prefix_len(tokens, &child.tokens);
	285	+ if cp == child.tokens.len() {
	286	+ // Whole child edge consumed — descend into it.
	287	+ insert_node(child, &tokens[cp..], &blocks[cp / BLOCK_SIZE..], now);
	288	+ } else {
	289	+ // Divergence mid-edge: split the child at `cp`.
	290	+ let split = RadixNode {
	291	+ tokens: child.tokens[cp..].to_vec(),
	292	+ blocks: child.blocks[cp / BLOCK_SIZE..].to_vec(),
	293	+ children: std::mem::take(&mut child.children),
	294	+ last_access: child.last_access,
	295	+ token_hash: hash_tokens(&child.tokens[cp..]),
	296	+ };
	297	+ child.tokens.truncate(cp);
	298	+ child.blocks.truncate(cp / BLOCK_SIZE);
	299	+ // The truncated edge now represents a shorter prefix → recompute
	300	+ // its content hash so the cache key stays consistent.
	301	+ child.token_hash = hash_tokens(&child.tokens);
	302	+ child.children.clear();
	303	+ child.children.insert(split.tokens[0], split);
	304	+ // Now `child` matches `tokens[..cp]`; descend with the remainder.
	305	+ insert_node(child, &tokens[cp..], &blocks[cp / BLOCK_SIZE..], now);
	306	+ }
	307	+ } else {
	308	+ // No matching child — create a leaf for the remaining path.
	309	+ node.children
	310	+ .insert(first, RadixNode::new(tokens.to_vec(), blocks.to_vec(), now));
	311	+ }
	312	+}
	313	+
	314	+/// Walk the longest matching prefix, accumulating matched `(tokens, blocks)`
	315	+/// and a `collision` flag. Refreshes `last_access` on every matched node to
	316	+/// `now` (LRU).
	317	+///
	318	+/// A full-edge match is only accepted if the node's stored `token_hash` equals
	319	+/// the hash of the tokens we actually matched (which equal `tokens[..cp]`) —
	320	+/// the collision-safe cache key. On divergence `collision` is set and the walk
	321	+/// stops, so the caller can reject instead of returning the wrong KV blocks.
	322	+fn lookup_node(
	323	+ node: &mut RadixNode,
	324	+ tokens: &[u32],
	325	+ out: &mut (usize, Vec<BlockId>, bool),
	326	+ now: u64,
	327	+) {
	328	+ if tokens.is_empty() {
	329	+ return;
	330	+ }
	331	+ let first = tokens[0];
	332	+ let cp = match node.children.get(&first) {
	333	+ Some(child) => common_prefix_len(tokens, &child.tokens),
	334	+ None => return,
	335	+ };
	336	+ let child = node.children.get_mut(&first).unwrap();
	337	+ if cp == child.tokens.len() {
	338	+ // Full edge match — accept only if the content hash agrees.
	339	+ if child.token_hash != hash_tokens(&tokens[..cp]) {
	340	+ // Structural match but hash divergence: a corruption or collision.
	341	+ out.2 = true;
	342	+ return;
	343	+ }
	344	+ out.0 += cp;
	345	+ out.1.extend_from_slice(&child.blocks);
	346	+ child.last_access = now;
	347	+ lookup_node(child, &tokens[cp..], out, now);
	348	+ } else if cp > 0 {
	349	+ // Partial edge match — the prompt is a strict prefix of this node's
	350	+ // tokens. The matched `cp` tokens equal `child.tokens[..cp]`, which is
	351	+ // already verified by the exact token comparison, so accumulate and stop.
	352	+ out.0 += cp;
	353	+ out.1.extend_from_slice(&child.blocks[..cp / BLOCK_SIZE]);
	354	+ child.last_access = now;
	355	+ }
	356	+}
	357	+
	358	+/// Find the globally least-recently-used unreferenced leaf, returning its
	359	+/// key path from the root and its block ids.
	360	+fn find_lru_leaf(node: &RadixNode) -> Option<(Vec<u32>, Vec<BlockId>)> {
	361	+ // Direct leaf children first.
	362	+ let mut direct: Vec<(u32, u64)> = node
	363	+ .children
	364	+ .iter()
	365	+ .filter(\|(_, c)\| c.is_leaf())
	366	+ .map(\|(k, c)\| (*k, c.last_access))
	367	+ .collect();
	368	+ direct.sort_by_key(\|&(_, a)\| a);
	369	+ if let Some((key, _)) = direct.first() {
	370	+ let child = &node.children[key];
	371	+ return Some((vec![*key], child.blocks.clone()));
	372	+ }
	373	+
	374	+ // Recurse into children, tracking the globally-oldest leaf.
	375	+ let mut best: Option<(Vec<u32>, Vec<BlockId>)> = None;
	376	+ let mut best_access = u64::MAX;
	377	+ for (k, child) in &node.children {
	378	+ if let Some((mut path, blocks)) = find_lru_leaf(child) {
	379	+ path.insert(0, *k);
	380	+ let access = child.last_access;
	381	+ if access < best_access {
	382	+ best_access = access;
	383	+ best = Some((path, blocks));
	384	+ }
	385	+ }
	386	+ }
	387	+ best
	388	+}
	389	+
	390	+/// Collect every leaf under `node` as `(last_access, key_path, block_ids)`,
	391	+/// where `key_path` is the sequence of child keys from the root to the leaf.
	392	+fn collect_leaves(
	393	+ node: &RadixNode,
	394	+ path: &mut Vec<u32>,
	395	+ out: &mut Vec<(u64, Vec<u32>, Vec<BlockId>)>,
	396	+) {
	397	+ for (k, child) in &node.children {
	398	+ path.push(*k);
	399	+ if child.is_leaf() {
	400	+ out.push((child.last_access, path.clone(), child.blocks.clone()));
	401	+ } else {
	402	+ collect_leaves(child, path, out);
	403	+ }
	404	+ path.pop();
	405	+ }
	406	+}
	407	+
	408	+/// Remove the leaf at `path` (a sequence of child keys from `node`).
	409	+fn remove_path(node: &mut RadixNode, path: &[u32]) -> bool {
	410	+ if path.is_empty() {
	411	+ return false;
	412	+ }
	413	+ let key = path[0];
	414	+ if let Some(child) = node.children.get_mut(&key) {
	415	+ if path.len() == 1 {
	416	+ node.children.remove(&key);
	417	+ true
	418	+ } else {
	419	+ remove_path(child, &path[1..])
	420	+ }
	421	+ } else {
	422	+ false
	423	+ }
	424	+}
	425	+
	426	+#[cfg(test)]
	427	+mod tests {
	428	+ use super::*;
	429	+
	430	+ fn blocks(n: usize) -> Vec<BlockId> {
	431	+ (0..n as BlockId).collect()
	432	+ }
	433	+
	434	+ #[test]
	435	+ fn insert_and_lookup_exact() {
	436	+ let mut cache = RadixPrefixCache::new();
	437	+ // 48 tokens = 3 blocks.
	438	+ let tokens: Vec<u32> = (0..48).collect();
	439	+ cache.insert(&tokens, &blocks(3));
	440	+
	441	+ let (matched, cached) = cache.lookup(&tokens).unwrap();
	442	+ assert_eq!(matched, 48);
	443	+ assert_eq!(cached, vec![0, 1, 2]);
	444	+ }
	445	+
	446	+ #[test]
	447	+ fn lookup_partial_prefix() {
	448	+ let mut cache = RadixPrefixCache::new();
	449	+ let tokens: Vec<u32> = (0..32).collect(); // 2 blocks
	450	+ cache.insert(&tokens, &blocks(2));
	451	+
	452	+ // Shorter prompt that is a prefix.
	453	+ let short: Vec<u32> = (0..16).collect();
	454	+ let (matched, cached) = cache.lookup(&short).unwrap();
	455	+ assert_eq!(matched, 16);
	456	+ assert_eq!(cached, vec![0]);
	457	+
	458	+ // Longer prompt extending the prefix (first 2 blocks match).
	459	+ let longer: Vec<u32> = (0..40).collect();
	460	+ let (matched, cached) = cache.lookup(&longer).unwrap();
	461	+ assert_eq!(matched, 32);
	462	+ assert_eq!(cached, vec![0, 1]);
	463	+ }
	464	+
	465	+ #[test]
	466	+ fn lookup_miss() {
	467	+ let mut cache = RadixPrefixCache::new();
	468	+ cache.insert(
	469	+ &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
	470	+ &[1],
	471	+ );
	472	+ assert!(cache.lookup(&[99, 100]).is_none());
	473	+ assert_eq!(cache.stats().misses, 1);
	474	+ }
	475	+
	476	+ #[test]
	477	+ fn empty_lookup_is_miss() {
	478	+ let mut cache = RadixPrefixCache::new();
	479	+ assert!(cache.lookup(&[1, 2, 3]).is_none());
	480	+ assert_eq!(cache.nodes(), 0);
	481	+ }
	482	+
	483	+ #[test]
	484	+ fn stats_track_hits_and_saved() {
	485	+ let mut cache = RadixPrefixCache::new();
	486	+ let tokens: Vec<u32> = (0..48).collect();
	487	+ cache.insert(&tokens, &blocks(3));
	488	+
	489	+ let _ = cache.lookup(&tokens);
	490	+ assert_eq!(cache.stats().hits, 1);
	491	+ assert_eq!(cache.stats().tokens_saved, 48);
	492	+
	493	+ let _ = cache.lookup(&tokens);
	494	+ assert_eq!(cache.stats().hits, 2);
	495	+ assert_eq!(cache.stats().tokens_saved, 96);
	496	+
	497	+ let _ = cache.lookup(&[999]);
	498	+ assert_eq!(cache.stats().misses, 1);
	499	+ }
	500	+
	501	+ // ---------- prefix sharing (the radix-tree win) ----------
	502	+
	503	+ #[test]
	504	+ fn shared_prefix_merges_nodes() {
	505	+ let mut cache = RadixPrefixCache::new();
	506	+ // Two sequences sharing the first 2 blocks (32 tokens).
	507	+ let a: Vec<u32> = (0..48).collect(); // 3 blocks
	508	+ let mut b: Vec<u32> = (0..32).collect(); // first 2 blocks
	509	+ b.extend_from_slice(&[
	510	+ 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115,
	511	+ ]); // 3rd block differs
	512	+
	513	+ cache.insert(&a, &blocks(3));
	514	+ cache.insert(&b, &blocks(3));
	515	+
	516	+ // Without sharing: 3 + 3 = 6 nodes would be needed. With sharing,
	517	+ // the common 2-block prefix is one node, plus 2 distinct leaves = 3.
	518	+ assert_eq!(cache.nodes(), 3, "common prefix should be merged");
	519	+
	520	+ // Both sequences still resolve their full prefixes.
	521	+ let (ma, _) = cache.lookup(&a).unwrap();
	522	+ assert_eq!(ma, 48);
	523	+ let (mb, _) = cache.lookup(&b).unwrap();
	524	+ assert_eq!(mb, 48);
	525	+ }
	526	+
	527	+ #[test]
	528	+ fn three_way_sharing() {
	529	+ let mut cache = RadixPrefixCache::new();
	530	+ // All share the first block [0..16], then diverge.
	531	+ let base: Vec<u32> = (0..16).collect();
	532	+ for suffix in [200u32, 300, 400] {
	533	+ let mut toks = base.clone();
	534	+ toks.extend_from_slice(&(suffix..suffix + 16).collect::<Vec<_>>());
	535	+ cache.insert(&toks, &blocks(2));
	536	+ }
	537	+ // Root's shared node [0..16] + 3 distinct leaves = 4 nodes (not 3×2 = 6).
	538	+ assert_eq!(cache.nodes(), 4);
	539	+ }
	540	+
	541	+ // ---------- splitting ----------
	542	+
	543	+ #[test]
	544	+ fn insert_splits_existing_edge() {
	545	+ let mut cache = RadixPrefixCache::new();
	546	+ // First insert a 2-block sequence.
	547	+ let a: Vec<u32> = (0..32).collect();
	548	+ cache.insert(&a, &blocks(2));
	549	+ assert_eq!(cache.nodes(), 1);
	550	+
	551	+ // Now insert a 1-block prefix of it — should split the existing node.
	552	+ let b: Vec<u32> = (0..16).collect();
	553	+ cache.insert(&b, &blocks(1));
	554	+ // Root + split node (1 block) + child (1 block) = 2 nodes.
	555	+ assert_eq!(cache.nodes(), 2);
	556	+
	557	+ // Both lookups resolve.
	558	+ let (ma, _) = cache.lookup(&a).unwrap();
	559	+ assert_eq!(ma, 32);
	560	+ let (mb, _) = cache.lookup(&b).unwrap();
	561	+ assert_eq!(mb, 16);
	562	+ }
	563	+
	564	+ // ---------- LRU eviction ----------
	565	+
	566	+ #[test]
	567	+ fn evict_removes_lru_leaf() {
	568	+ let mut cache = RadixPrefixCache::new();
	569	+ // Insert three independent 1-block sequences in order (oldest first).
	570	+ for start in [0u32, 100, 200] {
	571	+ let toks: Vec<u32> = (start..start + 16).collect();
	572	+ cache.insert(&toks, &[start / 100 as BlockId + 1]);
	573	+ }
	574	+ assert_eq!(cache.nodes(), 3);
	575	+
	576	+ // Evict 1 block — should remove the oldest (start=0) leaf.
	577	+ let freed = cache.evict(1);
	578	+ assert_eq!(freed, vec![1]);
	579	+ assert_eq!(cache.nodes(), 2);
	580	+
	581	+ // The evicted prefix no longer matches.
	582	+ let gone: Vec<u32> = (0..16).collect();
	583	+ assert!(cache.lookup(&gone).is_none());
	584	+ // The others still match.
	585	+ let still: Vec<u32> = (100..116).collect();
	586	+ assert!(cache.lookup(&still).is_some());
	587	+ }
	588	+
	589	+ #[test]
	590	+ fn evict_frees_multiple_blocks_from_one_leaf() {
	591	+ let mut cache = RadixPrefixCache::new();
	592	+ // One 3-block sequence.
	593	+ let toks: Vec<u32> = (0..48).collect();
	594	+ cache.insert(&toks, &[5, 6, 7]);
	595	+ assert_eq!(cache.nodes(), 1);
	596	+
	597	+ // Evicting 2 blocks from a single 3-block leaf frees all its blocks
	598	+ // (the whole leaf is removed).
	599	+ let freed = cache.evict(2);
	600	+ assert_eq!(freed, vec![5, 6, 7]);
	601	+ assert_eq!(cache.nodes(), 0);
	602	+ }
	603	+
	604	+ #[test]
	605	+ fn evict_handles_exhaustion() {
	606	+ let mut cache = RadixPrefixCache::new();
	607	+ let toks: Vec<u32> = (0..16).collect();
	608	+ cache.insert(&toks, &[1]);
	609	+ let freed = cache.evict(10); // ask for more than exists
	610	+ assert_eq!(freed, vec![1]);
	611	+ assert_eq!(cache.nodes(), 0);
	612	+ }
	613	+
	614	+ #[test]
	615	+ fn cached_tokens_accounts_for_shared_structure() {
	616	+ let mut cache = RadixPrefixCache::new();
	617	+ let a: Vec<u32> = (0..48).collect();
	618	+ let mut b: Vec<u32> = (0..32).collect();
	619	+ b.extend_from_slice(&[
	620	+ 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115,
	621	+ ]);
	622	+ cache.insert(&a, &blocks(3));
	623	+ cache.insert(&b, &blocks(3));
	624	+ // Shared prefix [0..32] stored once; a's tail [32..48] + b's tail
	625	+ // [100..116] stored distinctly. Total cached = 48 + 16 = 64.
	626	+ assert_eq!(cache.cached_tokens(), 64);
	627	+ }
	628	+
	629	+ // ---------- cache-key content hashing (collision safety) ----------
	630	+
	631	+ /// Test-only: flip the content hash of the first leaf node to simulate a
	632	+ /// corrupted / colliding cache entry, exercising the lookup guard.
	633	+ #[cfg(test)]
	634	+ fn corrupt_first_leaf_hash(cache: &mut RadixPrefixCache) {
	635	+ fn flip(node: &mut RadixNode) {
	636	+ if node.children.is_empty() {
	637	+ node.token_hash ^= 0xdead_beef;
	638	+ } else {
	639	+ let key = *node.children.keys().next().unwrap();
	640	+ flip(node.children.get_mut(&key).unwrap());
	641	+ }
	642	+ }
	643	+ flip(&mut cache.root);
	644	+ }
	645	+
	646	+ #[test]
	647	+ fn node_hash_matches_token_content() {
	648	+ let mut cache = RadixPrefixCache::new();
	649	+ let tokens: Vec<u32> = (0..48).collect();
	650	+ cache.insert(&tokens, &[1, 2, 3]);
	651	+ // Every node's stored hash must equal the hash of its tokens.
	652	+ assert!(cache.verify_hashes());
	653	+ }
	654	+
	655	+ #[test]
	656	+ fn hashes_consistent_after_split() {
	657	+ let mut cache = RadixPrefixCache::new();
	658	+ let a: Vec<u32> = (0..32).collect();
	659	+ cache.insert(&a, &[1, 2]);
	660	+ assert!(cache.verify_hashes());
	661	+ // Insert a strict prefix of `a` → splits the existing node. The
	662	+ // truncated edge must recompute its hash so keys stay consistent.
	663	+ let b: Vec<u32> = (0..16).collect();
	664	+ cache.insert(&b, &[3]);
	665	+ assert!(cache.verify_hashes());
	666	+ assert_eq!(cache.nodes(), 2);
	667	+ }
	668	+
	669	+ #[test]
	670	+ fn distinct_prefixes_do_not_collide() {
	671	+ let mut cache = RadixPrefixCache::new();
	672	+ let a: Vec<u32> = (0..48).collect();
	673	+ let mut b_tok: Vec<u32> = (0..48).collect();
	674	+ b_tok[0] = 999; // completely different first token → different hash
	675	+ cache.insert(&a, &[1, 2, 3]);
	676	+ cache.insert(&b_tok, &[4, 5, 6]);
	677	+
	678	+ let (ma, ba) = cache.lookup(&a).unwrap();
	679	+ assert_eq!(ma, 48);
	680	+ assert_eq!(ba, vec![1, 2, 3]);
	681	+ let (mb, bb) = cache.lookup(&b_tok).unwrap();
	682	+ assert_eq!(mb, 48);
	683	+ assert_eq!(bb, vec![4, 5, 6]);
	684	+ // Two structurally distinct prefixes never collide on the key.
	685	+ assert_eq!(cache.stats().hash_collisions, 0);
	686	+ }
	687	+
	688	+ #[test]
	689	+ fn lookup_rejects_corrupted_hash() {
	690	+ let mut cache = RadixPrefixCache::new();
	691	+ let tokens: Vec<u32> = (0..48).collect();
	692	+ cache.insert(&tokens, &[5, 6, 7]);
	693	+ assert!(cache.lookup(&tokens).is_some());
	694	+ assert!(cache.verify_hashes());
	695	+
	696	+ // Corrupt a node's content hash → the lookup guard must reject the
	697	+ // match rather than return the wrong KV blocks.
	698	+ corrupt_first_leaf_hash(&mut cache);
	699	+ assert!(!cache.verify_hashes());
	700	+ assert!(cache.lookup(&tokens).is_none());
	701	+ assert_eq!(cache.stats().hash_collisions, 1);
	702	+ }
	703	+
	704	+ // ---------- structural benchmark: radix vs flat hash-map ----------
	705	+
	706	+ /// Simulate a chat-history (nested prefix) workload: request `r` shares the
	707	+ /// chain `[b_0, b_1, …, b_r]` with prior requests, then appends one unique
	708	+ /// block `u_r`. Returns `(radix_block_slots, flat_block_refs)` where
	709	+ /// - `radix_block_slots` = `cached_tokens / BLOCK_SIZE` — each block
	710	+ /// stored once, so this is linear in `n_reqs`
	711	+ /// - `flat_block_refs` = `total_stored_blocks()` — the flat hash-map
	712	+ /// duplicates block IDs across distinct prefix entries, making this
	713	+ /// quadratic in `n_reqs`.
	714	+ fn bench_nested(n_reqs: usize) -> (usize, usize) {
	715	+ use crate::PrefixCache;
	716	+
	717	+ let block_size = BLOCK_SIZE;
	718	+
	719	+ let mut radix = RadixPrefixCache::new();
	720	+ let mut flat = PrefixCache::new();
	721	+
	722	+ for r in 0..n_reqs {
	723	+ // shared chain [b_0, b_1, ..., b_r] + unique [u_r]
	724	+ let mut toks: Vec<u32> = Vec::new();
	725	+ for k in 0..=r {
	726	+ let base = k as u32 * block_size as u32;
	727	+ toks.extend(base..base + block_size as u32);
	728	+ }
	729	+ let u_base = (n_reqs + r) as u32 * block_size as u32;
	730	+ toks.extend(u_base..u_base + block_size as u32);
	731	+
	732	+ let n_blocks = r + 2;
	733	+ let block_ids: Vec<BlockId> = (0..n_blocks as BlockId).collect();
	734	+ radix.insert(&toks, &block_ids);
	735	+ flat.insert(&toks, &block_ids);
	736	+ }
	737	+
	738	+ (
	739	+ radix.cached_tokens() / block_size,
	740	+ flat.total_stored_blocks(),
	741	+ )
	742	+ }
	743	+
	744	+ #[test]
	745	+ fn radix_linear_vs_flat_quadratic_nested_prefix() {
	746	+ // Chat-history workload: request `r` shares `r` blocks with prior
	747	+ // requests + 1 unique block. Flat stores each distinct prefix
	748	+ // separately, duplicating block IDs across entries → O(N²) block
	749	+ // references. Radix merges the shared chain → O(N) block slots.
	750	+ let (r20, f20) = bench_nested(20);
	751	+ let (r200, f200) = bench_nested(200);
	752	+
	753	+ // Radix scales linearly: 10× requests → ~10× block slots.
	754	+ assert!(r200 <= r20 * 12, "radix block-slots: {r200} vs {r20}");
	755	+
	756	+ // Flat scales quadratically: 10× requests → ~100× block refs.
	757	+ assert!(f200 >= f20 * 50, "flat block-refs: {f200} vs {f20}");
	758	+
	759	+ // At scale, radix uses a tiny fraction of flat's storage.
	760	+ assert!(r200 < f200 / 10, "radix={r200} flat={f200}");
	761	+ }
	762	+
	763	+ #[test]
	764	+ fn radix_memory_win_scales_with_nested_depth() {
	765	+ // More requests → deeper nested prefix → radix advantage grows.
	766	+ let (r_low, f_low) = bench_nested(50);
	767	+ let (r_big, f_big) = bench_nested(500);
	768	+
	769	+ // 10× requests → radix grows ~10×, flat grows ~100×.
	770	+ assert!(r_big <= r_low * 12, "radix: {r_big} vs {r_low}");
	771	+ assert!(f_big >= f_low * 50, "flat: {f_big} vs {f_low}");
	772	+
	773	+ // Radix dramatically outperforms at scale.
	774	+ assert!(r_big < f_big / 50, "radix={r_big} flat={f_big}");
	775	+ }
	776	+}
0	777

▼crates/justapi-inference/src/real/lora.rs

+334-0

@@ -0,0 +1,334 @@
	1	+//! LoRA adapter registry and routing.
	2	+//!
	3	+//! Pure data structures — no weight tensors. The registry decides which adapter
	4	+//! to activate for a given request and can resolve multi-LoRA routing keys.
	5	+//! Actual weight loading is handled by `model.rs` behind the `real` feature.
	6	+
	7	+use std::collections::HashMap;
	8	+use std::sync::{Arc, RwLock};
	9	+
	10	+/// LoRA rank configuration for a single adapter.
	11	+#[derive(Debug, Clone)]
	12	+pub struct LoraConfig {
	13	+ /// Adapter rank (typically 8–128).
	14	+ pub rank: usize,
	15	+ /// Alpha scaling factor (often equal to rank).
	16	+ pub alpha: usize,
	17	+ /// Dropout probability (0.0 = no dropout; used during training only).
	18	+ pub dropout: f32,
	19	+ /// Target module name patterns (e.g. "q_proj", "v_proj", "gate_proj").
	20	+ pub target_modules: Vec<String>,
	21	+}
	22	+
	23	+/// A loaded LoRA adapter (weights live in `model.rs` behind `real` feature).
	24	+#[derive(Debug, Clone)]
	25	+pub struct LoraAdapter {
	26	+ /// Unique adapter name (e.g. "adapter_v2_math").
	27	+ pub name: String,
	28	+ /// Optional base model hash — adapters are only valid for matching bases.
	29	+ pub base_model_hash: Option<String>,
	30	+ /// LoRA hyper-parameters.
	31	+ pub config: LoraConfig,
	32	+ /// Routing key used by the multi-LoRA scheduler. If set, requests matching
	33	+ /// this key are routed to this adapter.
	34	+ pub routing_key: Option<String>,
	35	+ /// Adapter is active and ready to serve.
	36	+ pub active: bool,
	37	+}
	38	+
	39	+impl LoraAdapter {
	40	+ /// Create a new adapter with sensible defaults.
	41	+ pub fn new(name: impl Into<String>, rank: usize, alpha: usize) -> Self {
	42	+ Self {
	43	+ name: name.into(),
	44	+ base_model_hash: None,
	45	+ config: LoraConfig {
	46	+ rank,
	47	+ alpha,
	48	+ dropout: 0.0,
	49	+ target_modules: Vec::new(),
	50	+ },
	51	+ routing_key: None,
	52	+ active: true,
	53	+ }
	54	+ }
	55	+
	56	+ /// Builder-style setter for routing key.
	57	+ pub fn with_routing_key(mut self, key: impl Into<String>) -> Self {
	58	+ self.routing_key = Some(key.into());
	59	+ self
	60	+ }
	61	+
	62	+ /// Builder-style setter for base model hash.
	63	+ pub fn with_base_model(mut self, hash: impl Into<String>) -> Self {
	64	+ self.base_model_hash = Some(hash.into());
	65	+ self
	66	+ }
	67	+
	68	+ /// Builder-style setter for target modules.
	69	+ pub fn with_targets(mut self, modules: Vec<String>) -> Self {
	70	+ self.config.target_modules = modules;
	71	+ self
	72	+ }
	73	+
	74	+ /// Effective scaling factor: `alpha / rank`.
	75	+ pub fn scaling(&self) -> f32 {
	76	+ self.config.alpha as f32 / self.config.rank.max(1) as f32
	77	+ }
	78	+}
	79	+
	80	+/// Thread-safe registry for managing multiple LoRA adapters.
	81	+#[derive(Clone)]
	82	+pub struct LoraRegistry {
	83	+ inner: Arc<RwLock<LoraRegistryInner>>,
	84	+}
	85	+
	86	+struct LoraRegistryInner {
	87	+ adapters: HashMap<String, LoraAdapter>,
	88	+ /// Routing-key → adapter-name lookup (built on register/unregister).
	89	+ routing: HashMap<String, String>,
	90	+}
	91	+
	92	+impl LoraRegistry {
	93	+ /// Create an empty registry.
	94	+ pub fn new() -> Self {
	95	+ Self {
	96	+ inner: Arc::new(RwLock::new(LoraRegistryInner {
	97	+ adapters: HashMap::new(),
	98	+ routing: HashMap::new(),
	99	+ })),
	100	+ }
	101	+ }
	102	+
	103	+ /// Register an adapter. Replaces any existing adapter with the same name.
	104	+ pub fn register(&self, adapter: LoraAdapter) {
	105	+ let mut inner = self.inner.write().unwrap();
	106	+ if let Some(key) = &adapter.routing_key {
	107	+ inner.routing.insert(key.clone(), adapter.name.clone());
	108	+ }
	109	+ inner.adapters.insert(adapter.name.clone(), adapter);
	110	+ }
	111	+
	112	+ /// Unregister an adapter by name.
	113	+ pub fn unregister(&self, name: &str) -> Option<LoraAdapter> {
	114	+ let mut inner = self.inner.write().unwrap();
	115	+ let adapter = inner.adapters.remove(name)?;
	116	+ if let Some(key) = &adapter.routing_key {
	117	+ inner.routing.remove(key);
	118	+ }
	119	+ Some(adapter)
	120	+ }
	121	+
	122	+ /// Look up an adapter by name.
	123	+ pub fn get(&self, name: &str) -> Option<LoraAdapter> {
	124	+ self.inner.read().unwrap().adapters.get(name).cloned()
	125	+ }
	126	+
	127	+ /// Look up an adapter by routing key.
	128	+ pub fn get_by_routing_key(&self, key: &str) -> Option<LoraAdapter> {
	129	+ let inner = self.inner.read().unwrap();
	130	+ let name = inner.routing.get(key)?;
	131	+ inner.adapters.get(name).cloned()
	132	+ }
	133	+
	134	+ /// List all active adapter names.
	135	+ pub fn active_adapters(&self) -> Vec<String> {
	136	+ self.inner
	137	+ .read()
	138	+ .unwrap()
	139	+ .adapters
	140	+ .values()
	141	+ .filter(\|a\| a.active)
	142	+ .map(\|a\| a.name.clone())
	143	+ .collect()
	144	+ }
	145	+
	146	+ /// List all adapter names (active + inactive).
	147	+ pub fn all_adapters(&self) -> Vec<String> {
	148	+ self.inner
	149	+ .read()
	150	+ .unwrap()
	151	+ .adapters
	152	+ .keys()
	153	+ .cloned()
	154	+ .collect()
	155	+ }
	156	+
	157	+ /// Number of registered adapters.
	158	+ pub fn len(&self) -> usize {
	159	+ self.inner.read().unwrap().adapters.len()
	160	+ }
	161	+
	162	+ /// Whether the registry is empty.
	163	+ pub fn is_empty(&self) -> bool {
	164	+ self.inner.read().unwrap().adapters.is_empty()
	165	+ }
	166	+
	167	+ /// Activate an adapter by name. Returns `false` if not found.
	168	+ pub fn activate(&self, name: &str) -> bool {
	169	+ let mut inner = self.inner.write().unwrap();
	170	+ if let Some(a) = inner.adapters.get_mut(name) {
	171	+ a.active = true;
	172	+ true
	173	+ } else {
	174	+ false
	175	+ }
	176	+ }
	177	+
	178	+ /// Deactivate an adapter by name. Returns `false` if not found.
	179	+ pub fn deactivate(&self, name: &str) -> bool {
	180	+ let mut inner = self.inner.write().unwrap();
	181	+ if let Some(a) = inner.adapters.get_mut(name) {
	182	+ a.active = false;
	183	+ true
	184	+ } else {
	185	+ false
	186	+ }
	187	+ }
	188	+
	189	+ /// Resolve which adapter to use for a request.
	190	+ ///
	191	+ /// Priority: explicit adapter name → routing key → `None` (base model).
	192	+ pub fn resolve(
	193	+ &self,
	194	+ adapter_name: Option<&str>,
	195	+ routing_key: Option<&str>,
	196	+ ) -> Option<LoraAdapter> {
	197	+ if let Some(name) = adapter_name {
	198	+ let adapter = self.get(name)?;
	199	+ return if adapter.active { Some(adapter) } else { None };
	200	+ }
	201	+ if let Some(key) = routing_key {
	202	+ let adapter = self.get_by_routing_key(key)?;
	203	+ return if adapter.active { Some(adapter) } else { None };
	204	+ }
	205	+ None
	206	+ }
	207	+
	208	+ /// Check if any registered adapter targets a given module.
	209	+ pub fn active_for_module(&self, module_name: &str) -> bool {
	210	+ self.inner
	211	+ .read()
	212	+ .unwrap()
	213	+ .adapters
	214	+ .values()
	215	+ .filter(\|a\| a.active)
	216	+ .any(\|a\| a.config.target_modules.iter().any(\|m\| m == module_name))
	217	+ }
	218	+}
	219	+
	220	+impl Default for LoraRegistry {
	221	+ fn default() -> Self {
	222	+ Self::new()
	223	+ }
	224	+}
	225	+
	226	+// ---------------------------------------------------------------------------
	227	+// Tests
	228	+// ---------------------------------------------------------------------------
	229	+
	230	+#[cfg(test)]
	231	+mod tests {
	232	+ use super::*;
	233	+
	234	+ #[test]
	235	+ fn adapter_new_and_scaling() {
	236	+ let a = LoraAdapter::new("math", 64, 128);
	237	+ assert_eq!(a.name, "math");
	238	+ assert_eq!(a.config.rank, 64);
	239	+ assert_eq!(a.config.alpha, 128);
	240	+ assert!((a.scaling() - 2.0).abs() < f32::EPSILON);
	241	+ assert!(a.active);
	242	+ }
	243	+
	244	+ #[test]
	245	+ fn adapter_builder_chain() {
	246	+ let a = LoraAdapter::new("code", 32, 32)
	247	+ .with_routing_key("task:code")
	248	+ .with_base_model("llama-7b-hash")
	249	+ .with_targets(vec!["q_proj".into(), "v_proj".into()]);
	250	+ assert_eq!(a.routing_key.as_deref(), Some("task:code"));
	251	+ assert_eq!(a.base_model_hash.as_deref(), Some("llama-7b-hash"));
	252	+ assert_eq!(a.config.target_modules, vec!["q_proj", "v_proj"]);
	253	+ }
	254	+
	255	+ #[test]
	256	+ fn registry_register_and_get() {
	257	+ let reg = LoraRegistry::new();
	258	+ assert!(reg.is_empty());
	259	+ reg.register(LoraAdapter::new("a", 8, 8));
	260	+ reg.register(LoraAdapter::new("b", 16, 16));
	261	+ assert_eq!(reg.len(), 2);
	262	+ assert!(reg.get("a").is_some());
	263	+ assert!(reg.get("c").is_none());
	264	+ }
	265	+
	266	+ #[test]
	267	+ fn registry_routing_key() {
	268	+ let reg = LoraRegistry::new();
	269	+ reg.register(LoraAdapter::new("math_adapter", 32, 32).with_routing_key("domain:math"));
	270	+ let found = reg.get_by_routing_key("domain:math");
	271	+ assert!(found.is_some());
	272	+ assert_eq!(found.unwrap().name, "math_adapter");
	273	+ assert!(reg.get_by_routing_key("domain:code").is_none());
	274	+ }
	275	+
	276	+ #[test]
	277	+ fn registry_activate_deactivate() {
	278	+ let reg = LoraRegistry::new();
	279	+ reg.register(LoraAdapter::new("a", 8, 8));
	280	+ assert!(reg.activate("a"));
	281	+ assert!(reg.deactivate("a"));
	282	+ assert!(!reg.get("a").unwrap().active);
	283	+ assert!(reg.active_adapters().is_empty());
	284	+ assert_eq!(reg.all_adapters(), vec!["a"]);
	285	+ }
	286	+
	287	+ #[test]
	288	+ fn registry_resolve_priority() {
	289	+ let reg = LoraRegistry::new();
	290	+ reg.register(LoraAdapter::new("explicit", 8, 8).with_routing_key("r1"));
	291	+ reg.register(LoraAdapter::new("routed", 8, 8).with_routing_key("r2"));
	292	+
	293	+ // Explicit name takes priority.
	294	+ let a = reg.resolve(Some("explicit"), Some("r2"));
	295	+ assert_eq!(a.unwrap().name, "explicit");
	296	+
	297	+ // Fallback to routing key.
	298	+ let a = reg.resolve(None, Some("r2"));
	299	+ assert_eq!(a.unwrap().name, "routed");
	300	+
	301	+ // No match → None.
	302	+ assert!(reg.resolve(None, None).is_none());
	303	+ }
	304	+
	305	+ #[test]
	306	+ fn registry_resolve_inactive_returns_none() {
	307	+ let reg = LoraRegistry::new();
	308	+ reg.register(LoraAdapter::new("a", 8, 8).with_routing_key("k"));
	309	+ reg.deactivate("a");
	310	+ assert!(reg.resolve(Some("a"), None).is_none());
	311	+ assert!(reg.resolve(None, Some("k")).is_none());
	312	+ }
	313	+
	314	+ #[test]
	315	+ fn registry_active_for_module() {
	316	+ let reg = LoraRegistry::new();
	317	+ reg.register(
	318	+ LoraAdapter::new("a", 8, 8).with_targets(vec!["q_proj".into(), "v_proj".into()]),
	319	+ );
	320	+ assert!(reg.active_for_module("q_proj"));
	321	+ assert!(reg.active_for_module("v_proj"));
	322	+ assert!(!reg.active_for_module("gate_proj"));
	323	+ }
	324	+
	325	+ #[test]
	326	+ fn registry_unregister() {
	327	+ let reg = LoraRegistry::new();
	328	+ reg.register(LoraAdapter::new("a", 8, 8).with_routing_key("k1"));
	329	+ let removed = reg.unregister("a");
	330	+ assert!(removed.is_some());
	331	+ assert!(reg.is_empty());
	332	+ assert!(reg.get_by_routing_key("k1").is_none());
	333	+ }
	334	+}
0	335

▼crates/justapi-inference/src/real/mod.rs

+17-0

@@ -0,0 +1,17 @@
	1	+//! Real-model support: quantization configs, LoRA adapters, weight loading.
	2	+//!
	3	+//! ## Compilation strategy
	4	+//!
	5	+//! - `quant.rs` + `lora.rs` are always compiled and tested in the default
	6	+//! gate (pure data structures, zero weight dependencies).
	7	+//! - `model.rs` + `tokenizer.rs` are gated behind `#[cfg(feature = "real")]`
	8	+//! and pull in `candle-transformers` + `tokenizers`.
	9	+
	10	+pub mod lora;
	11	+pub mod quant;
	12	+
	13	+#[cfg(feature = "real")]
	14	+pub mod model;
	15	+
	16	+#[cfg(feature = "real")]
	17	+pub mod tokenizer;
0	18

▼crates/justapi-inference/src/real/model.rs

+500-0

@@ -0,0 +1,500 @@
	1	+//! Real model loading and forward pass — GGUF and safetensors.
	2	+//!
	3	+//! Only compiled under `#[cfg(feature = "real")]`. Pulls in `candle-transformers`
	4	+//! for the actual Llama forward pass and `tokenizers` for decoding.
	5	+
	6	+use std::fs::File;
	7	+use std::path::{Path, PathBuf};
	8	+use std::sync::Mutex;
	9	+
	10	+use candle_core::quantized::gguf_file;
	11	+use candle_core::{DType, Device, Tensor};
	12	+use candle_nn::VarBuilder;
	13	+use candle_transformers::generation::{LogitsProcessor, Sampling};
	14	+use candle_transformers::models::llama::{Cache, Config as LlamaConfig, Llama};
	15	+use candle_transformers::models::quantized_llama::ModelWeights;
	16	+
	17	+use crate::model::{FinishReason, GeneratedToken, Model, ModelError, SamplingParams};
	18	+use crate::real::tokenizer::Tokenizer;
	19	+
	20	+/// Returns true if `path` looks like a valid GGUF file (correct magic header).
	21	+fn is_valid_gguf(path: &Path) -> bool {
	22	+ let mut buf = [0u8; 4];
	23	+ match std::fs::File::open(path).and_then(\|mut f\| std::io::Read::read_exact(&mut f, &mut buf)) {
	24	+ Ok(()) => &buf == b"GGUF",
	25	+ Err(_) => false,
	26	+ }
	27	+}
	28	+
	29	+/// Detect whether a model directory contains a valid `.gguf` file or safetensors.
	30	+pub fn detect_format(model_dir: &Path) -> Result<ModelFormat, ModelError> {
	31	+ let entries: Vec<_> = std::fs::read_dir(model_dir)
	32	+ .map_err(ModelError::Io)?
	33	+ .filter_map(\|e\| e.ok())
	34	+ .collect();
	35	+
	36	+ let has_gguf = entries
	37	+ .iter()
	38	+ .filter(\|e\| {
	39	+ e.path()
	40	+ .extension()
	41	+ .map(\|ext\| ext == "gguf")
	42	+ .unwrap_or(false)
	43	+ })
	44	+ .any(\|e\| is_valid_gguf(&e.path()));
	45	+
	46	+ let has_safetensors = entries.iter().any(\|e\| {
	47	+ e.path()
	48	+ .extension()
	49	+ .map(\|ext\| ext == "safetensors")
	50	+ .unwrap_or(false)
	51	+ });
	52	+
	53	+ match (has_gguf, has_safetensors) {
	54	+ (true, _) => Ok(ModelFormat::Gguf),
	55	+ (_, true) => Ok(ModelFormat::Safetensors),
	56	+ _ => Err(ModelError::Generation(format!(
	57	+ "no .gguf or .safetensors files found in {}",
	58	+ model_dir.display()
	59	+ ))),
	60	+ }
	61	+}
	62	+
	63	+/// Detected model weight format.
	64	+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	65	+pub enum ModelFormat {
	66	+ Gguf,
	67	+ Safetensors,
	68	+}
	69	+
	70	+/// Raw Llama config.json structure (safetensors models ship a separate config).
	71	+#[derive(Debug, serde::Deserialize)]
	72	+struct RawLlamaConfig {
	73	+ hidden_size: usize,
	74	+ intermediate_size: usize,
	75	+ num_attention_heads: usize,
	76	+ num_hidden_layers: usize,
	77	+ num_key_value_heads: usize,
	78	+ vocab_size: usize,
	79	+ rms_norm_eps: Option<f64>,
	80	+ rope_theta: Option<f32>,
	81	+ max_position_embeddings: Option<usize>,
	82	+ #[serde(default)]
	83	+ tie_word_embeddings: bool,
	84	+}
	85	+
	86	+impl RawLlamaConfig {
	87	+ fn to_candle_config(&self) -> LlamaConfig {
	88	+ LlamaConfig {
	89	+ hidden_size: self.hidden_size,
	90	+ intermediate_size: self.intermediate_size,
	91	+ num_attention_heads: self.num_attention_heads,
	92	+ num_hidden_layers: self.num_hidden_layers,
	93	+ num_key_value_heads: self.num_key_value_heads,
	94	+ vocab_size: self.vocab_size,
	95	+ rms_norm_eps: self.rms_norm_eps.unwrap_or(1e-5),
	96	+ rope_theta: self.rope_theta.unwrap_or(10000.0),
	97	+ max_position_embeddings: self.max_position_embeddings.unwrap_or(2048),
	98	+ tie_word_embeddings: self.tie_word_embeddings,
	99	+ bos_token_id: None,
	100	+ eos_token_id: None,
	101	+ rope_scaling: None,
	102	+ use_flash_attn: false,
	103	+ }
	104	+ }
	105	+}
	106	+
	107	+/// Mutable forward-pass state. Wrapped in a `Mutex` so `RealModel::generate`
	108	+/// can hold a `&self` borrow (required by the [`Model`] trait) while driving the
	109	+/// autoregressive loop. The GGUF path keeps its KV cache inside [`ModelWeights`];
	110	+/// the safetensors path keeps an explicit [`Cache`].
	111	+enum ForwardState {
	112	+ Gguf(ModelWeights),
	113	+ Safetensors { model: Llama, cache: Cache },
	114	+}
	115	+
	116	+impl ForwardState {
	117	+ /// Run one forward step over `tokens` at absolute `index_pos`.
	118	+ fn forward(
	119	+ &mut self,
	120	+ tokens: &[u32],
	121	+ index_pos: usize,
	122	+ device: &Device,
	123	+ ) -> Result<Tensor, ModelError> {
	124	+ let t = Tensor::from_slice(tokens, (1, tokens.len()), device)
	125	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	126	+ match self {
	127	+ ForwardState::Gguf(m) => m
	128	+ .forward(&t, index_pos)
	129	+ .map_err(\|e\| ModelError::Generation(e.to_string())),
	130	+ ForwardState::Safetensors { model, cache } => model
	131	+ .forward(&t, index_pos, cache)
	132	+ .map_err(\|e\| ModelError::Generation(e.to_string())),
	133	+ }
	134	+ }
	135	+
	136	+ /// Reset KV state before a fresh sequence.
	137	+ fn reset_cache(&mut self, config: &LlamaConfig, device: &Device) -> Result<(), ModelError> {
	138	+ match self {
	139	+ ForwardState::Gguf(m) => m.clear_kv_cache(),
	140	+ ForwardState::Safetensors { cache, .. } => {
	141	+ *cache = Cache::new(true, DType::F32, config, device)
	142	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	143	+ }
	144	+ }
	145	+ Ok(())
	146	+ }
	147	+}
	148	+
	149	+/// A loaded real model (GGUF or safetensors) implementing the `Model` trait.
	150	+pub struct RealModel {
	151	+ format: ModelFormat,
	152	+ config: LlamaConfig,
	153	+ device: Device,
	154	+ /// Tokenizer for decoding generated token IDs to text.
	155	+ tokenizer: Tokenizer,
	156	+ /// Forward-pass state (mutated during generation).
	157	+ state: Mutex<ForwardState>,
	158	+}
	159	+
	160	+impl RealModel {
	161	+ /// Load a real model from `model_dir`.
	162	+ ///
	163	+ /// Expects:
	164	+ /// - `tokenizer.json` — HuggingFace tokenizer file
	165	+ /// - Either `.gguf` (GGUF format) or `.safetensors` + `config.json`
	166	+ pub fn load(model_dir: &Path, device: Device) -> Result<Self, ModelError> {
	167	+ let format = detect_format(model_dir)?;
	168	+ let tokenizer_path = model_dir.join("tokenizer.json");
	169	+ let tokenizer = if tokenizer_path.exists() {
	170	+ Tokenizer::from_file(&tokenizer_path)?
	171	+ } else {
	172	+ // For GGUF files, try to extract the tokenizer from metadata.
	173	+ if format == ModelFormat::Gguf {
	174	+ // First load the GGUF metadata.
	175	+ let entries: Vec<_> = std::fs::read_dir(model_dir)
	176	+ .map_err(ModelError::Io)?
	177	+ .filter_map(\|e\| e.ok())
	178	+ .map(\|e\| e.path())
	179	+ .collect();
	180	+
	181	+ let gguf_path = entries
	182	+ .iter()
	183	+ .find(\|p\| {
	184	+ p.extension().map(\|e\| e == "gguf").unwrap_or(false) && is_valid_gguf(p)
	185	+ })
	186	+ .ok_or_else(\|\| {
	187	+ ModelError::Generation("no valid .gguf file found".to_string())
	188	+ })?;
	189	+
	190	+ let mut file = File::open(gguf_path).map_err(ModelError::Io)?;
	191	+ let ct = gguf_file::Content::read(&mut file)
	192	+ .map_err(\|e\| ModelError::Generation(format!("gguf read failed: {e}")))?;
	193	+
	194	+ Tokenizer::from_gguf_metadata(&ct.metadata)?
	195	+ } else {
	196	+ return Err(ModelError::Generation(
	197	+ "tokenizer.json not found — safetensors models require tokenizer.json"
	198	+ .to_string(),
	199	+ ));
	200	+ }
	201	+ };
	202	+
	203	+ let entries: Vec<_> = std::fs::read_dir(model_dir)
	204	+ .map_err(ModelError::Io)?
	205	+ .filter_map(\|e\| e.ok())
	206	+ .map(\|e\| e.path())
	207	+ .collect();
	208	+
	209	+ // Build the Llama config (needed by both formats for the safetensors
	210	+ // forward pass and as a fallback for GGUF).
	211	+ let config = load_config(model_dir, &entries)?;
	212	+
	213	+ let state = match format {
	214	+ ModelFormat::Gguf => {
	215	+ let gguf_path = entries
	216	+ .iter()
	217	+ .find(\|p\| {
	218	+ p.extension().map(\|e\| e == "gguf").unwrap_or(false) && is_valid_gguf(p)
	219	+ })
	220	+ .ok_or_else(\|\| {
	221	+ ModelError::Generation("no valid .gguf file found".to_string())
	222	+ })?;
	223	+ let mut file = File::open(gguf_path).map_err(ModelError::Io)?;
	224	+ let ct = gguf_file::Content::read(&mut file)
	225	+ .map_err(\|e\| ModelError::Generation(format!("gguf read failed: {e}")))?;
	226	+ let model = ModelWeights::from_gguf(ct, &mut file, &device)
	227	+ .map_err(\|e\| ModelError::Generation(format!("gguf load failed: {e}")))?;
	228	+ ForwardState::Gguf(model)
	229	+ }
	230	+ ModelFormat::Safetensors => {
	231	+ let st_paths: Vec<PathBuf> = entries
	232	+ .iter()
	233	+ .filter(\|p\| p.extension().map(\|e\| e == "safetensors").unwrap_or(false))
	234	+ .cloned()
	235	+ .collect();
	236	+ let vb =
	237	+ unsafe { VarBuilder::from_mmaped_safetensors(&st_paths, DType::F32, &device) }
	238	+ .map_err(\|e\| {
	239	+ ModelError::Generation(format!("safetensors load failed: {e}"))
	240	+ })?;
	241	+ let model = Llama::load(vb, &config)
	242	+ .map_err(\|e\| ModelError::Generation(format!("llama load failed: {e}")))?;
	243	+ let cache = Cache::new(true, DType::F32, &config, &device)
	244	+ .map_err(\|e\| ModelError::Generation(format!("cache init failed: {e}")))?;
	245	+ ForwardState::Safetensors { model, cache }
	246	+ }
	247	+ };
	248	+
	249	+ tracing::info!(
	250	+ "loaded {:?} model: hidden={}, layers={}, heads={}, kv_heads={}, vocab={}",
	251	+ format,
	252	+ config.hidden_size,
	253	+ config.num_hidden_layers,
	254	+ config.num_attention_heads,
	255	+ config.num_key_value_heads,
	256	+ config.vocab_size,
	257	+ );
	258	+
	259	+ Ok(Self {
	260	+ format,
	261	+ config,
	262	+ device,
	263	+ tokenizer,
	264	+ state: Mutex::new(state),
	265	+ })
	266	+ }
	267	+
	268	+ /// The detected format.
	269	+ pub fn format(&self) -> ModelFormat {
	270	+ self.format
	271	+ }
	272	+
	273	+ /// The loaded Llama config.
	274	+ pub fn config(&self) -> &LlamaConfig {
	275	+ &self.config
	276	+ }
	277	+
	278	+ /// Tokenize text into token IDs using the model's tokenizer.
	279	+ pub fn encode(&self, text: &str) -> Result<Vec<u32>, ModelError> {
	280	+ self.tokenizer.encode(text)
	281	+ }
	282	+}
	283	+
	284	+/// Load the Llama config: from `config.json` if present, else a sensible default.
	285	+fn load_config(model_dir: &Path, entries: &[PathBuf]) -> Result<LlamaConfig, ModelError> {
	286	+ let has_config = entries
	287	+ .iter()
	288	+ .any(\|p\| p.file_name().map(\|n\| n == "config.json").unwrap_or(false));
	289	+ if has_config {
	290	+ let config_str =
	291	+ std::fs::read_to_string(model_dir.join("config.json")).map_err(ModelError::Io)?;
	292	+ let raw: RawLlamaConfig = serde_json::from_str(&config_str)
	293	+ .map_err(\|e\| ModelError::Generation(format!("config.json parse error: {e}")))?;
	294	+ Ok(raw.to_candle_config())
	295	+ } else {
	296	+ // Minimal default — GGUF metadata overrides tensor dims at load time.
	297	+ Ok(LlamaConfig {
	298	+ hidden_size: 4096,
	299	+ intermediate_size: 11008,
	300	+ num_attention_heads: 32,
	301	+ num_hidden_layers: 32,
	302	+ num_key_value_heads: 32,
	303	+ vocab_size: 32000,
	304	+ rms_norm_eps: 1e-5,
	305	+ rope_theta: 10000.0,
	306	+ max_position_embeddings: 2048,
	307	+ tie_word_embeddings: false,
	308	+ bos_token_id: None,
	309	+ eos_token_id: None,
	310	+ rope_scaling: None,
	311	+ use_flash_attn: false,
	312	+ })
	313	+ }
	314	+}
	315	+
	316	+/// Build a [`LogitsProcessor`] honoring temperature / top_p / top_k.
	317	+fn build_processor(seed: u64, params: &SamplingParams) -> LogitsProcessor {
	318	+ if params.top_k > 0 {
	319	+ LogitsProcessor::from_sampling(
	320	+ seed,
	321	+ Sampling::TopKThenTopP {
	322	+ k: params.top_k,
	323	+ p: params.top_p as f64,
	324	+ temperature: params.temperature as f64,
	325	+ },
	326	+ )
	327	+ } else {
	328	+ let temperature = if params.temperature < 1e-7 {
	329	+ None
	330	+ } else {
	331	+ Some(params.temperature as f64)
	332	+ };
	333	+ let top_p = if params.top_p >= 1.0 {
	334	+ None
	335	+ } else {
	336	+ Some(params.top_p as f64)
	337	+ };
	338	+ LogitsProcessor::new(seed, temperature, top_p)
	339	+ }
	340	+}
	341	+
	342	+impl Model for RealModel {
	343	+ fn vocab_size(&self) -> usize {
	344	+ self.config.vocab_size
	345	+ }
	346	+
	347	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError> {
	348	+ let mut state = self.state.lock().unwrap();
	349	+ state.reset_cache(&self.config, &self.device)?;
	350	+ let logits = state.forward(context, 0, &self.device)?;
	351	+ // `to_vec1` returns the last position's logits (shape [vocab]).
	352	+ let vec = logits
	353	+ .to_vec1::<f32>()
	354	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	355	+ Ok(vec)
	356	+ }
	357	+
	358	+ #[allow(clippy::explicit_counter_loop)]
	359	+ fn generate(
	360	+ &self,
	361	+ prompt: &[u32],
	362	+ params: &SamplingParams,
	363	+ sink: &dyn Fn(GeneratedToken) -> bool,
	364	+ ) -> Result<FinishReason, ModelError> {
	365	+ // Non-deterministic seed so repeated calls vary (temperature/top_p).
	366	+ let seed = std::time::SystemTime::now()
	367	+ .duration_since(std::time::UNIX_EPOCH)
	368	+ .map(\|d\| d.as_nanos() as u64)
	369	+ .unwrap_or(299792458);
	370	+ let mut lp = build_processor(seed, params);
	371	+
	372	+ let mut state = self.state.lock().unwrap();
	373	+ state.reset_cache(&self.config, &self.device)?;
	374	+
	375	+ // Prefill: feed the whole prompt at once, get logits for the last position.
	376	+ // `index_pos` tracks the absolute position in the KV cache (prompt length
	377	+ // + number of tokens already decoded) — it is NOT a redundant copy of the
	378	+ // `step` loop counter, so the explicit counter is intentional.
	379	+ let logits = state.forward(prompt, 0, &self.device)?;
	380	+ let mut token = lp
	381	+ .sample(&logits)
	382	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	383	+ let mut index_pos = prompt.len();
	384	+
	385	+ for step in 0..params.max_tokens {
	386	+ let text = self
	387	+ .tokenizer
	388	+ .decode(&[token])
	389	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	390	+ let is_stop = params.stop_tokens.contains(&token);
	391	+ let finish = if is_stop {
	392	+ Some(FinishReason::Stop)
	393	+ } else {
	394	+ None
	395	+ };
	396	+ let produced = GeneratedToken {
	397	+ id: token,
	398	+ text,
	399	+ logprob: 0.0,
	400	+ finish_reason: finish,
	401	+ };
	402	+ if !sink(produced) {
	403	+ return Ok(FinishReason::Cancelled);
	404	+ }
	405	+ if is_stop {
	406	+ return Ok(FinishReason::Stop);
	407	+ }
	408	+ // Produced the final token → done.
	409	+ if step + 1 >= params.max_tokens {
	410	+ return Ok(FinishReason::Length);
	411	+ }
	412	+ // Decode the next token, extending the KV cache.
	413	+ let next_logits = state.forward(&[token], index_pos, &self.device)?;
	414	+ token = lp
	415	+ .sample(&next_logits)
	416	+ .map_err(\|e\| ModelError::Generation(e.to_string()))?;
	417	+ index_pos += 1;
	418	+ }
	419	+ Ok(FinishReason::Length)
	420	+ }
	421	+}
	422	+
	423	+// ---------------------------------------------------------------------------
	424	+// Tests (feature-gated `real`)
	425	+// ---------------------------------------------------------------------------
	426	+
	427	+#[cfg(test)]
	428	+mod tests {
	429	+ use super::*;
	430	+
	431	+ #[test]
	432	+ fn detect_format_gguf() {
	433	+ let dir = std::env::temp_dir().join("justapi_test_gguf");
	434	+ let _ = std::fs::create_dir_all(&dir);
	435	+ let _ = std::fs::write(dir.join("model.gguf"), b"GGUF");
	436	+ let fmt = detect_format(&dir).unwrap();
	437	+ assert_eq!(fmt, ModelFormat::Gguf);
	438	+ let _ = std::fs::remove_dir_all(&dir);
	439	+ }
	440	+
	441	+ #[test]
	442	+ fn detect_format_safetensors() {
	443	+ let dir = std::env::temp_dir().join("justapi_test_st");
	444	+ let _ = std::fs::create_dir_all(&dir);
	445	+ let _ = std::fs::write(dir.join("model.safetensors"), b"{}");
	446	+ let _ = std::fs::write(dir.join("config.json"), b"{}");
	447	+ let fmt = detect_format(&dir).unwrap();
	448	+ assert_eq!(fmt, ModelFormat::Safetensors);
	449	+ let _ = std::fs::remove_dir_all(&dir);
	450	+ }
	451	+
	452	+ #[test]
	453	+ fn detect_format_empty_errors() {
	454	+ let dir = std::env::temp_dir().join("justapi_test_empty");
	455	+ let _ = std::fs::create_dir_all(&dir);
	456	+ assert!(detect_format(&dir).is_err());
	457	+ let _ = std::fs::remove_dir_all(&dir);
	458	+ }
	459	+
	460	+ #[test]
	461	+ fn raw_config_parses_and_maps() {
	462	+ let json = r#"{
	463	+ "hidden_size": 768,
	464	+ "intermediate_size": 2048,
	465	+ "num_attention_heads": 12,
	466	+ "num_hidden_layers": 12,
	467	+ "num_key_value_heads": 12,
	468	+ "vocab_size": 50257,
	469	+ "rms_norm_eps": 1e-6,
	470	+ "rope_theta": 10000.0,
	471	+ "max_position_embeddings": 1024
	472	+ }"#;
	473	+ let raw: RawLlamaConfig = serde_json::from_str(json).unwrap();
	474	+ let cfg = raw.to_candle_config();
	475	+ assert_eq!(cfg.hidden_size, 768);
	476	+ assert_eq!(cfg.num_hidden_layers, 12);
	477	+ assert_eq!(cfg.vocab_size, 50257);
	478	+ assert_eq!(cfg.max_position_embeddings, 1024);
	479	+ assert!(cfg.rope_scaling.is_none());
	480	+ assert!(!cfg.use_flash_attn);
	481	+ }
	482	+
	483	+ #[test]
	484	+ fn raw_config_defaults_for_optional_fields() {
	485	+ let json = r#"{
	486	+ "hidden_size": 4096,
	487	+ "intermediate_size": 11008,
	488	+ "num_attention_heads": 32,
	489	+ "num_hidden_layers": 32,
	490	+ "num_key_value_heads": 32,
	491	+ "vocab_size": 32000
	492	+ }"#;
	493	+ let raw: RawLlamaConfig = serde_json::from_str(json).unwrap();
	494	+ let cfg = raw.to_candle_config();
	495	+ assert_eq!(cfg.rms_norm_eps, 1e-5);
	496	+ assert_eq!(cfg.rope_theta, 10000.0);
	497	+ assert_eq!(cfg.max_position_embeddings, 2048);
	498	+ assert!(!cfg.tie_word_embeddings);
	499	+ }
	500	+}
0	501

▼crates/justapi-inference/src/real/quant.rs

+274-0

@@ -0,0 +1,274 @@
	1	+//! Quantization configuration and method selection.
	2	+//!
	3	+//! Pure data structures — no weight loading, no Candle tensors. These types are
	4	+//! always compiled so the scheduler can make quant-aware decisions (e.g. which
	5	+//! adapter to route to) without pulling in heavy dependencies.
	6	+
	7	+use std::collections::HashMap;
	8	+
	9	+/// Supported quantization backends.
	10	+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
	11	+pub enum QuantMethod {
	12	+ /// No quantization (f16/bf16/f32).
	13	+ None,
	14	+ /// GPTQ (Marlin kernel on CUDA, CPU fallback via unpacking).
	15	+ Gptq,
	16	+ /// AWQ (activation-aware quantization).
	17	+ Awq,
	18	+ /// GGUF k-quants (Q2_K through Q8_0, IQ variants).
	19	+ Gguf,
	20	+ /// EXL2 (ExLlamaV2 kernel).
	21	+ Exl2,
	22	+ /// BitsAndBytes (8-bit / 4-bit, widely used in HuggingFace).
	23	+ BitsAndBytes,
	24	+}
	25	+
	26	+impl QuantMethod {
	27	+ /// Human-readable label for logs and metrics.
	28	+ pub fn as_str(&self) -> &'static str {
	29	+ match self {
	30	+ Self::None => "none",
	31	+ Self::Gptq => "gptq",
	32	+ Self::Awq => "awq",
	33	+ Self::Gguf => "gguf",
	34	+ Self::Exl2 => "exl2",
	35	+ Self::BitsAndBytes => "bnb",
	36	+ }
	37	+ }
	38	+
	39	+ /// Approximate bits-per-weight for display / heuristics.
	40	+ pub fn approx_bits(&self) -> f32 {
	41	+ match self {
	42	+ Self::None => 16.0,
	43	+ Self::Gptq => 4.0,
	44	+ Self::Awq => 4.0,
	45	+ Self::Gguf => 4.5,
	46	+ Self::Exl2 => 3.5,
	47	+ Self::BitsAndBytes => 8.0,
	48	+ }
	49	+ }
	50	+}
	51	+
	52	+/// Per-layer quantization metadata (for mixed-precision models).
	53	+#[derive(Debug, Clone)]
	54	+pub struct LayerQuantConfig {
	55	+ pub method: QuantMethod,
	56	+ /// Group size used by the quantization kernel (0 = per-tensor).
	57	+ pub group_size: u32,
	58	+ /// Symmetric quantization flag.
	59	+ pub symmetric: bool,
	60	+}
	61	+
	62	+/// Full quantization descriptor for a model or a specific adapter.
	63	+#[derive(Debug, Clone)]
	64	+pub struct QuantConfig {
	65	+ /// Global quantization method.
	66	+ pub method: QuantMethod,
	67	+ /// Per-layer overrides. Keys are layer name prefixes (e.g. "model.layers.0").
	68	+ pub layers: HashMap<String, LayerQuantConfig>,
	69	+ /// Bits of precision (informational; actual kernels depend on `method`).
	70	+ pub bits: f32,
	71	+}
	72	+
	73	+impl QuantConfig {
	74	+ /// Build a uniform quant config (same method for all layers).
	75	+ pub fn uniform(method: QuantMethod, bits: f32) -> Self {
	76	+ Self {
	77	+ method,
	78	+ layers: HashMap::new(),
	79	+ bits,
	80	+ }
	81	+ }
	82	+
	83	+ /// Effective method for a specific layer (falls back to global).
	84	+ pub fn method_for_layer(&self, layer_name: &str) -> QuantMethod {
	85	+ // Try exact match first, then prefix match (longest prefix wins).
	86	+ let mut best: Option<&LayerQuantConfig> = None;
	87	+ let mut best_len = 0usize;
	88	+ for (prefix, cfg) in &self.layers {
	89	+ if layer_name.starts_with(prefix.as_str()) && prefix.len() > best_len {
	90	+ best = Some(cfg);
	91	+ best_len = prefix.len();
	92	+ }
	93	+ }
	94	+ best.map(\|c\| c.method).unwrap_or(self.method)
	95	+ }
	96	+
	97	+ /// Rough memory estimate in bytes for a model with `param_count` parameters.
	98	+ pub fn memory_bytes(&self, param_count: u64) -> u64 {
	99	+ let bits = self.bits.clamp(1.0, 32.0);
	100	+ param_count * bits as u64 / 8
	101	+ }
	102	+}
	103	+
	104	+/// GGUF k-quant type mapping (used when loading .gguf files).
	105	+pub mod gguf {
	106	+ use super::QuantMethod;
	107	+
	108	+ /// GGUF quantization type IDs (subset used by llama.cpp / candle).
	109	+ #[derive(Debug, Clone, Copy, PartialEq, Eq)]
	110	+ #[repr(u32)]
	111	+ #[allow(non_camel_case_types)]
	112	+ pub enum GgufType {
	113	+ F32 = 0,
	114	+ F16 = 1,
	115	+ Q4_0 = 2,
	116	+ Q4_1 = 3,
	117	+ Q5_0 = 6,
	118	+ Q5_1 = 7,
	119	+ Q8_0 = 8,
	120	+ Q8_1 = 9,
	121	+ Q2_K = 10,
	122	+ Q3_K_S = 11,
	123	+ Q3_K_M = 12,
	124	+ Q3_K_L = 13,
	125	+ Q4_K_S = 14,
	126	+ Q4_K_M = 15,
	127	+ Q5_K_S = 16,
	128	+ Q5_K_M = 17,
	129	+ Q6_K = 18,
	130	+ IQ2_XXS = 19,
	131	+ IQ2_XS = 20,
	132	+ IQ3_XXS = 21,
	133	+ IQ1_S = 22,
	134	+ IQ4_NL = 23,
	135	+ IQ3_S = 24,
	136	+ IQ2_S = 25,
	137	+ IQ4_XS = 26,
	138	+ }
	139	+
	140	+ /// Approximate bits-per-weight for a GGUF type.
	141	+ pub fn bits_for_type(t: GgufType) -> f32 {
	142	+ match t {
	143	+ GgufType::F32 => 32.0,
	144	+ GgufType::F16 => 16.0,
	145	+ GgufType::Q4_0 \| GgufType::Q4_1 => 4.5,
	146	+ GgufType::Q5_0 \| GgufType::Q5_1 => 5.5,
	147	+ GgufType::Q8_0 \| GgufType::Q8_1 => 8.5,
	148	+ GgufType::Q2_K => 2.5,
	149	+ GgufType::Q3_K_S \| GgufType::Q3_K_M \| GgufType::Q3_K_L => 3.5,
	150	+ GgufType::Q4_K_S \| GgufType::Q4_K_M => 4.5,
	151	+ GgufType::Q5_K_S \| GgufType::Q5_K_M => 5.5,
	152	+ GgufType::Q6_K => 6.5,
	153	+ GgufType::IQ2_XXS \| GgufType::IQ2_XS \| GgufType::IQ2_S => 2.0,
	154	+ GgufType::IQ3_XXS \| GgufType::IQ3_S => 3.0,
	155	+ GgufType::IQ1_S => 1.5,
	156	+ GgufType::IQ4_NL \| GgufType::IQ4_XS => 4.0,
	157	+ }
	158	+ }
	159	+
	160	+ /// All supported quant type IDs (for validation).
	161	+ pub fn supported_ids() -> &'static [u32] {
	162	+ &[
	163	+ 0, 1, 2, 3, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,
	164	+ 26,
	165	+ ]
	166	+ }
	167	+
	168	+ /// Whether a GGUF type ID is supported.
	169	+ pub fn is_supported(id: u32) -> bool {
	170	+ supported_ids().contains(&id)
	171	+ }
	172	+
	173	+ /// Map a GGUF type to our QuantMethod.
	174	+ pub fn to_quant_method(t: GgufType) -> QuantMethod {
	175	+ match t {
	176	+ GgufType::F32 \| GgufType::F16 => QuantMethod::None,
	177	+ _ => QuantMethod::Gguf,
	178	+ }
	179	+ }
	180	+}
	181	+
	182	+// ---------------------------------------------------------------------------
	183	+// Tests
	184	+// ---------------------------------------------------------------------------
	185	+
	186	+#[cfg(test)]
	187	+mod tests {
	188	+ use super::*;
	189	+
	190	+ #[test]
	191	+ fn quant_method_as_str_and_bits() {
	192	+ assert_eq!(QuantMethod::None.as_str(), "none");
	193	+ assert_eq!(QuantMethod::None.approx_bits(), 16.0);
	194	+ assert_eq!(QuantMethod::Gptq.approx_bits(), 4.0);
	195	+ }
	196	+
	197	+ #[test]
	198	+ fn quant_config_uniform() {
	199	+ let cfg = QuantConfig::uniform(QuantMethod::Gptq, 4.0);
	200	+ assert_eq!(cfg.method, QuantMethod::Gptq);
	201	+ assert!(cfg.layers.is_empty());
	202	+ assert_eq!(cfg.method_for_layer("anything"), QuantMethod::Gptq);
	203	+ }
	204	+
	205	+ #[test]
	206	+ fn quant_config_layer_override() {
	207	+ let mut cfg = QuantConfig::uniform(QuantMethod::Gptq, 4.0);
	208	+ cfg.layers.insert(
	209	+ "model.layers.0".to_string(),
	210	+ LayerQuantConfig {
	211	+ method: QuantMethod::None,
	212	+ group_size: 128,
	213	+ symmetric: true,
	214	+ },
	215	+ );
	216	+ assert_eq!(cfg.method_for_layer("model.layers.0"), QuantMethod::None);
	217	+ assert_eq!(cfg.method_for_layer("model.layers.1"), QuantMethod::Gptq);
	218	+ }
	219	+
	220	+ #[test]
	221	+ fn quant_config_longest_prefix_wins() {
	222	+ let mut cfg = QuantConfig::uniform(QuantMethod::Awq, 4.0);
	223	+ cfg.layers.insert(
	224	+ "model.layers".to_string(),
	225	+ LayerQuantConfig {
	226	+ method: QuantMethod::Gptq,
	227	+ group_size: 0,
	228	+ symmetric: false,
	229	+ },
	230	+ );
	231	+ cfg.layers.insert(
	232	+ "model.layers.0.self_attn".to_string(),
	233	+ LayerQuantConfig {
	234	+ method: QuantMethod::None,
	235	+ group_size: 128,
	236	+ symmetric: true,
	237	+ },
	238	+ );
	239	+ assert_eq!(
	240	+ cfg.method_for_layer("model.layers.0.self_attn.q_proj"),
	241	+ QuantMethod::None
	242	+ );
	243	+ assert_eq!(
	244	+ cfg.method_for_layer("model.layers.1.ffn"),
	245	+ QuantMethod::Gptq
	246	+ );
	247	+ }
	248	+
	249	+ #[test]
	250	+ fn memory_estimate() {
	251	+ let cfg = QuantConfig::uniform(QuantMethod::Gptq, 4.0);
	252	+ let bytes = cfg.memory_bytes(7_000_000_000); // 7B params
	253	+ assert_eq!(bytes, 3_500_000_000); // 4 bits = 0.5 bytes
	254	+ }
	255	+
	256	+ #[test]
	257	+ fn gguf_bits_and_support() {
	258	+ assert_eq!(gguf::bits_for_type(gguf::GgufType::F32), 32.0);
	259	+ assert!(gguf::is_supported(gguf::GgufType::Q4_K_M as u32));
	260	+ assert!(!gguf::is_supported(999));
	261	+ }
	262	+
	263	+ #[test]
	264	+ fn gguf_to_quant_method() {
	265	+ assert_eq!(
	266	+ gguf::to_quant_method(gguf::GgufType::F16),
	267	+ QuantMethod::None
	268	+ );
	269	+ assert_eq!(
	270	+ gguf::to_quant_method(gguf::GgufType::Q4_K_M),
	271	+ QuantMethod::Gguf
	272	+ );
	273	+ }
	274	+}
0	275

▼crates/justapi-inference/src/real/tokenizer.rs

+144-0

@@ -0,0 +1,144 @@
	1	+//! Thin wrapper around the `tokenizers` crate for decoding token IDs.
	2	+//!
	3	+//! The real model forward pass produces `Vec<u32>` token ids; this module
	4	+//! converts them to text. Only compiled under the `real` feature.
	5	+
	6	+use std::collections::HashMap;
	7	+
	8	+use candle_core::quantized::gguf_file::Value;
	9	+
	10	+use crate::model::ModelError;
	11	+
	12	+/// Wrapper around a HuggingFace `Tokenizer` loaded from a file or built from
	13	+/// GGUF metadata.
	14	+pub struct Tokenizer {
	15	+ inner: tokenizers::Tokenizer,
	16	+}
	17	+
	18	+/// Extract an array of strings from GGUF metadata by key.
	19	+fn extract_string_array(metadata: &HashMap<String, Value>, key: &str) -> Vec<String> {
	20	+ metadata
	21	+ .get(key)
	22	+ .and_then(\|v\| match v {
	23	+ Value::Array(arr) => Some(
	24	+ arr.iter()
	25	+ .filter_map(\|x\| x.to_string().ok())
	26	+ .cloned()
	27	+ .collect(),
	28	+ ),
	29	+ _ => None,
	30	+ })
	31	+ .unwrap_or_default()
	32	+}
	33	+
	34	+impl Tokenizer {
	35	+ /// Load a tokenizer from a `tokenizer.json` file on disk.
	36	+ pub fn from_file(path: impl AsRef<std::path::Path>) -> Result<Self, ModelError> {
	37	+ let path = path.as_ref();
	38	+ let inner = tokenizers::Tokenizer::from_file(path)
	39	+ .map_err(\|e\| ModelError::Generation(format!("tokenizer load failed: {e}")))?;
	40	+ Ok(Self { inner })
	41	+ }
	42	+
	43	+ /// Build a tokenizer from GGUF file metadata.
	44	+ ///
	45	+ /// GGUF files embed tokenizer data (vocab, merges, scores) in their
	46	+ /// metadata section under keys like `tokenizer.ggml.model`, `.tokens`,
	47	+ /// `.merges`, `.scores`. This function converts that metadata into a
	48	+ /// HuggingFace-compatible `tokenizer.json` and loads it in memory.
	49	+ pub fn from_gguf_metadata(metadata: &HashMap<String, Value>) -> Result<Self, ModelError> {
	50	+ // GGUF "gpt2" / "llama" tokenizers are GPT-2-style BPE.
	51	+ let tokens: Vec<String> = extract_string_array(metadata, "tokenizer.ggml.tokens");
	52	+ let merges: Vec<String> = extract_string_array(metadata, "tokenizer.ggml.merges");
	53	+
	54	+ // Build the HuggingFace tokenizer JSON structure.
	55	+ let mut vocab = serde_json::Map::new();
	56	+ for (i, token) in tokens.iter().enumerate() {
	57	+ vocab.insert(token.clone(), serde_json::Value::Number(i.into()));
	58	+ }
	59	+
	60	+ // GPT-2 / Llama BPE uses ByteLevel pre-tokenizer + decoder and the
	61	+ // `Ġ` sentinel for spaces in the vocab.
	62	+ let model_json = serde_json::json!({
	63	+ "type": "BPE",
	64	+ "dropout": null,
	65	+ "unk_token": null,
	66	+ "continuing_subword_prefix": null,
	67	+ "end_of_word_suffix": "",
	68	+ "fuse_unk": false,
	69	+ "byte_fallback": true,
	70	+ "vocab": vocab,
	71	+ "merges": merges,
	72	+ });
	73	+
	74	+ let pre_tokenizer = serde_json::json!({
	75	+ "type": "ByteLevel",
	76	+ "add_prefix_space": false,
	77	+ "trim_offsets": false,
	78	+ });
	79	+
	80	+ let decoder = serde_json::json!({
	81	+ "type": "ByteLevel",
	82	+ "add_prefix_space": false,
	83	+ "trim_offsets": false,
	84	+ });
	85	+
	86	+ let _bos = metadata
	87	+ .get("tokenizer.ggml.bos_token_id")
	88	+ .and_then(\|v\| v.to_u32().ok());
	89	+ let _eos = metadata
	90	+ .get("tokenizer.ggml.eos_token_id")
	91	+ .and_then(\|v\| v.to_u32().ok());
	92	+
	93	+ // BOS/EOS injection is handled by the model's generate path, not the
	94	+ // HuggingFace tokenizer post-processor, so we omit it here to keep the
	95	+ // GGUF-derived tokenizer minimal and loadable.
	96	+ let post_processor = serde_json::Value::Null;
	97	+
	98	+ let tokenizer_json = serde_json::json!({
	99	+ "version": "1.0",
	100	+ "truncation": null,
	101	+ "padding": null,
	102	+ "added_tokens": [],
	103	+ "normalizer": null,
	104	+ "pre_tokenizer": pre_tokenizer,
	105	+ "post_processor": post_processor,
	106	+ "decoder": decoder,
	107	+ "model": model_json,
	108	+ });
	109	+
	110	+ let bytes = serde_json::to_vec(&tokenizer_json)
	111	+ .map_err(\|e\| ModelError::Generation(format!("tokenizer json serialize: {e}")))?;
	112	+
	113	+ let inner = tokenizers::Tokenizer::from_bytes(&bytes)
	114	+ .map_err(\|e\| ModelError::Generation(format!("tokenizer from gguf metadata: {e}")))?;
	115	+
	116	+ Ok(Self { inner })
	117	+ }
	118	+
	119	+ /// Encode text into token IDs (truncation/padding disabled).
	120	+ pub fn encode(&self, text: &str) -> Result<Vec<u32>, ModelError> {
	121	+ let encoding = self
	122	+ .inner
	123	+ .encode(text, true)
	124	+ .map_err(\|e\| ModelError::Generation(format!("tokenizer encode failed: {e}")))?;
	125	+ Ok(encoding.get_ids().to_vec())
	126	+ }
	127	+
	128	+ /// Decode a sequence of token IDs to text.
	129	+ pub fn decode(&self, ids: &[u32]) -> Result<String, ModelError> {
	130	+ self.inner
	131	+ .decode(ids, true)
	132	+ .map_err(\|e\| ModelError::Generation(format!("tokenizer decode failed: {e}")))
	133	+ }
	134	+
	135	+ /// Decode a single token ID to text.
	136	+ pub fn decode_token(&self, id: u32) -> Result<String, ModelError> {
	137	+ self.decode(&[id])
	138	+ }
	139	+
	140	+ /// Vocabulary size.
	141	+ pub fn vocab_size(&self) -> usize {
	142	+ self.inner.get_vocab_size(true)
	143	+ }
	144	+}
0	145

▼crates/justapi-inference/src/router.rs

+416-0

@@ -0,0 +1,416 @@
	1	+//! LLM request router — LoRA-aware and KV-pressure-aware replica selection.
	2	+//!
	3	+//! Pure decision logic (no network, no GPU). Built on top of the [`ControlPlane`]
	4	+//! registry: a request names a model (+ optional version + LoRA routing key), the
	5	+//! control plane resolves it to a concrete [`ResolvedModel`], and the router picks
	6	+//! the best replica currently able to serve it.
	7	+//!
	8	+//! - LoRA-aware: if a routing key resolved to an adapter, only replicas that
	9	+//! already have that adapter resident are considered (avoids a cold LoRA load).
	10	+//! - KV-aware: among healthy replicas with spare capacity, selection honors
	11	+//! the configured [`RoutingStrategy`] (least-loaded, lowest KV pressure, or
	12	+//! round-robin) so hot replicas are not overloaded.
	13	+
	14	+use std::sync::atomic::{AtomicUsize, Ordering};
	15	+use std::sync::Mutex;
	16	+
	17	+use crate::control_plane::{ControlPlane, ResolvedModel};
	18	+use crate::engine::EngineDevice;
	19	+
	20	+/// A single serving replica (one model version on one device/process).
	21	+#[derive(Debug, Clone)]
	22	+pub struct Replica {
	23	+ /// Unique replica id (e.g. "llama-7b-v2-gpu0").
	24	+ pub id: String,
	25	+ /// Model name this replica serves.
	26	+ pub model: String,
	27	+ /// Model version this replica serves.
	28	+ pub version: String,
	29	+ /// Device the replica is pinned to.
	30	+ pub device: EngineDevice,
	31	+ /// Max concurrent sequences this replica admits.
	32	+ pub max_concurrency: usize,
	33	+ /// Currently active sequences (drives load-aware selection).
	34	+ pub active_sequences: usize,
	35	+ /// KV-cache pressure as a percentage (0–100), drives KV-aware selection.
	36	+ pub kv_pressure_pct: f32,
	37	+ /// Whether the replica is healthy and accepting traffic.
	38	+ pub healthy: bool,
	39	+ /// Adapter names currently resident on this replica.
	40	+ pub loaded_adapters: Vec<String>,
	41	+}
	42	+
	43	+impl Replica {
	44	+ /// Spare capacity before the replica hits `max_concurrency`.
	45	+ pub fn available_capacity(&self) -> usize {
	46	+ self.max_concurrency.saturating_sub(self.active_sequences)
	47	+ }
	48	+
	49	+ /// A replica is overloaded when it has no spare concurrency or its KV cache
	50	+ /// is near saturation.
	51	+ pub fn is_overloaded(&self) -> bool {
	52	+ self.active_sequences >= self.max_concurrency \|\| self.kv_pressure_pct >= 95.0
	53	+ }
	54	+}
	55	+
	56	+/// An incoming routing request (what the client asked for).
	57	+#[derive(Debug, Clone, Default)]
	58	+pub struct RouteRequest {
	59	+ pub model: String,
	60	+ pub version: Option<String>,
	61	+ pub routing_key: Option<String>,
	62	+}
	63	+
	64	+/// The outcome of routing: which replica + the fully resolved model target.
	65	+#[derive(Debug, Clone)]
	66	+pub struct RouteDecision {
	67	+ pub replica: String,
	68	+ pub resolved: ResolvedModel,
	69	+}
	70	+
	71	+/// How the router breaks ties among candidate replicas.
	72	+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	73	+pub enum RoutingStrategy {
	74	+ /// Fewest active sequences first.
	75	+ LeastLoaded,
	76	+ /// Lowest KV-cache pressure first.
	77	+ LowestKvPressure,
	78	+ /// Round-robin across healthy candidates (spreads load evenly).
	79	+ RoundRobin,
	80	+}
	81	+
	82	+/// The router: owns its view of live replicas and selects among them.
	83	+pub struct Router {
	84	+ strategy: RoutingStrategy,
	85	+ replicas: Mutex<Vec<Replica>>,
	86	+ rr: AtomicUsize,
	87	+}
	88	+
	89	+impl Router {
	90	+ /// Create a router with the given strategy.
	91	+ pub fn new(strategy: RoutingStrategy) -> Self {
	92	+ Self {
	93	+ strategy,
	94	+ replicas: Mutex::new(Vec::new()),
	95	+ rr: AtomicUsize::new(0),
	96	+ }
	97	+ }
	98	+
	99	+ /// Register or replace a replica by id.
	100	+ pub fn register_replica(&self, replica: Replica) {
	101	+ let mut guard = self.replicas.lock().unwrap();
	102	+ if let Some(slot) = guard.iter_mut().find(\|r\| r.id == replica.id) {
	103	+ *slot = replica;
	104	+ } else {
	105	+ guard.push(replica);
	106	+ }
	107	+ }
	108	+
	109	+ /// Update a replica's live load metrics (no-op if the id is unknown).
	110	+ pub fn update_load(
	111	+ &self,
	112	+ id: &str,
	113	+ active_sequences: usize,
	114	+ kv_pressure_pct: f32,
	115	+ healthy: bool,
	116	+ ) {
	117	+ let mut guard = self.replicas.lock().unwrap();
	118	+ if let Some(r) = guard.iter_mut().find(\|r\| r.id == id) {
	119	+ r.active_sequences = active_sequences;
	120	+ r.kv_pressure_pct = kv_pressure_pct;
	121	+ r.healthy = healthy;
	122	+ }
	123	+ }
	124	+
	125	+ /// Report which adapters are resident on a replica (no-op if unknown).
	126	+ pub fn set_loaded_adapters(&self, id: &str, adapters: Vec<String>) {
	127	+ let mut guard = self.replicas.lock().unwrap();
	128	+ if let Some(r) = guard.iter_mut().find(\|r\| r.id == id) {
	129	+ r.loaded_adapters = adapters;
	130	+ }
	131	+ }
	132	+
	133	+ /// Remove a replica by id.
	134	+ pub fn remove_replica(&self, id: &str) -> Option<Replica> {
	135	+ let mut guard = self.replicas.lock().unwrap();
	136	+ let idx = guard.iter().position(\|r\| r.id == id)?;
	137	+ Some(guard.remove(idx))
	138	+ }
	139	+
	140	+ /// All currently registered replicas (snapshot).
	141	+ pub fn replicas(&self) -> Vec<Replica> {
	142	+ self.replicas.lock().unwrap().clone()
	143	+ }
	144	+
	145	+ /// Replicas eligible to serve `resolved`: matching model + version, healthy,
	146	+ /// not overloaded, and — if an adapter is required — already holding it.
	147	+ fn candidates(&self, resolved: &ResolvedModel) -> Vec<Replica> {
	148	+ self.replicas
	149	+ .lock()
	150	+ .unwrap()
	151	+ .iter()
	152	+ .filter(\|r\| {
	153	+ r.model == resolved.model_name
	154	+ && r.version == resolved.version
	155	+ && r.healthy
	156	+ && !r.is_overloaded()
	157	+ && resolved
	158	+ .adapter
	159	+ .as_ref()
	160	+ .map(\|a\| r.loaded_adapters.iter().any(\|la\| la == a))
	161	+ .unwrap_or(true)
	162	+ })
	163	+ .cloned()
	164	+ .collect()
	165	+ }
	166	+
	167	+ /// Route a request: resolve it via the control plane, then select the best
	168	+ /// replica. Returns `None` if the model is unknown or no replica can serve it.
	169	+ pub fn route(&self, cp: &ControlPlane, request: &RouteRequest) -> Option<RouteDecision> {
	170	+ let resolved = cp.resolve(
	171	+ &request.model,
	172	+ request.version.as_deref(),
	173	+ request.routing_key.as_deref(),
	174	+ )?;
	175	+ let mut candidates = self.candidates(&resolved);
	176	+ if candidates.is_empty() {
	177	+ return None;
	178	+ }
	179	+ let chosen = match self.strategy {
	180	+ RoutingStrategy::LeastLoaded => candidates
	181	+ .into_iter()
	182	+ .min_by_key(\|r\| r.active_sequences)
	183	+ .unwrap(),
	184	+ RoutingStrategy::LowestKvPressure => candidates
	185	+ .into_iter()
	186	+ .min_by(\|a, b\| {
	187	+ a.kv_pressure_pct
	188	+ .partial_cmp(&b.kv_pressure_pct)
	189	+ .unwrap_or(std::cmp::Ordering::Equal)
	190	+ })
	191	+ .unwrap(),
	192	+ RoutingStrategy::RoundRobin => {
	193	+ let n = candidates.len();
	194	+ let idx = self.rr.fetch_add(1, Ordering::Relaxed) % n;
	195	+ candidates.swap_remove(idx)
	196	+ }
	197	+ };
	198	+ Some(RouteDecision {
	199	+ replica: chosen.id,
	200	+ resolved,
	201	+ })
	202	+ }
	203	+}
	204	+
	205	+// ---------------------------------------------------------------------------
	206	+// Tests
	207	+// ---------------------------------------------------------------------------
	208	+
	209	+#[cfg(test)]
	210	+mod tests {
	211	+ use super::*;
	212	+ use crate::control_plane::{ControlPlane, WeightLocation};
	213	+ use std::path::PathBuf;
	214	+
	215	+ fn cp_with_model() -> ControlPlane {
	216	+ let cp = ControlPlane::new();
	217	+ cp.register_version(
	218	+ "llama-7b",
	219	+ crate::control_plane::ModelVersion::new(
	220	+ "v1",
	221	+ WeightLocation::Local(PathBuf::from("/models/v1")),
	222	+ ),
	223	+ );
	224	+ let mut record = cp.get("llama-7b").unwrap();
	225	+ record.route_adapter("domain:math", "math-adapter");
	226	+ record.route_adapter("domain:code", "code-adapter");
	227	+ cp.register(record);
	228	+ cp
	229	+ }
	230	+
	231	+ fn replica(id: &str, active: usize, kv: f32, adapters: &[&str]) -> Replica {
	232	+ Replica {
	233	+ id: id.to_string(),
	234	+ model: "llama-7b".to_string(),
	235	+ version: "v1".to_string(),
	236	+ device: EngineDevice::Cuda(0),
	237	+ max_concurrency: 4,
	238	+ active_sequences: active,
	239	+ kv_pressure_pct: kv,
	240	+ healthy: true,
	241	+ loaded_adapters: adapters.iter().map(\|s\| s.to_string()).collect(),
	242	+ }
	243	+ }
	244	+
	245	+ #[test]
	246	+ fn routes_to_matching_healthy_replica() {
	247	+ let cp = cp_with_model();
	248	+ let router = Router::new(RoutingStrategy::LeastLoaded);
	249	+ router.register_replica(replica("r1", 1, 10.0, &[]));
	250	+ let d = router
	251	+ .route(
	252	+ &cp,
	253	+ &RouteRequest {
	254	+ model: "llama-7b".into(),
	255	+ version: None,
	256	+ routing_key: None,
	257	+ },
	258	+ )
	259	+ .unwrap();
	260	+ assert_eq!(d.replica, "r1");
	261	+ assert_eq!(d.resolved.version, "v1");
	262	+ }
	263	+
	264	+ #[test]
	265	+ fn least_loaded_picks_fewest_active() {
	266	+ let cp = cp_with_model();
	267	+ let router = Router::new(RoutingStrategy::LeastLoaded);
	268	+ router.register_replica(replica("busy", 3, 10.0, &[]));
	269	+ router.register_replica(replica("idle", 0, 80.0, &[]));
	270	+ let d = router
	271	+ .route(
	272	+ &cp,
	273	+ &RouteRequest {
	274	+ model: "llama-7b".into(),
	275	+ version: None,
	276	+ routing_key: None,
	277	+ },
	278	+ )
	279	+ .unwrap();
	280	+ assert_eq!(d.replica, "idle"); // fewer active sequences wins
	281	+ }
	282	+
	283	+ #[test]
	284	+ fn lowest_kv_pressure_picks_emptier_cache() {
	285	+ let cp = cp_with_model();
	286	+ let router = Router::new(RoutingStrategy::LowestKvPressure);
	287	+ router.register_replica(replica("hot", 1, 90.0, &[]));
	288	+ router.register_replica(replica("cool", 1, 20.0, &[]));
	289	+ let d = router
	290	+ .route(
	291	+ &cp,
	292	+ &RouteRequest {
	293	+ model: "llama-7b".into(),
	294	+ version: None,
	295	+ routing_key: None,
	296	+ },
	297	+ )
	298	+ .unwrap();
	299	+ assert_eq!(d.replica, "cool");
	300	+ }
	301	+
	302	+ #[test]
	303	+ fn lora_aware_only_considers_replicas_with_adapter() {
	304	+ let cp = cp_with_model();
	305	+ let router = Router::new(RoutingStrategy::LeastLoaded);
	306	+ // Two replicas: only r2 has the math-adapter resident.
	307	+ router.register_replica(replica("r1", 0, 10.0, &[]));
	308	+ router.register_replica(replica("r2", 0, 10.0, &["math-adapter"]));
	309	+
	310	+ // Request with the math routing key → must land on r2.
	311	+ let d = router
	312	+ .route(
	313	+ &cp,
	314	+ &RouteRequest {
	315	+ model: "llama-7b".into(),
	316	+ version: None,
	317	+ routing_key: Some("domain:math".into()),
	318	+ },
	319	+ )
	320	+ .unwrap();
	321	+ assert_eq!(d.replica, "r2");
	322	+ assert_eq!(d.resolved.adapter.as_deref(), Some("math-adapter"));
	323	+
	324	+ // Request with a routing key whose adapter nobody has → no capacity.
	325	+ assert!(router
	326	+ .route(
	327	+ &cp,
	328	+ &RouteRequest {
	329	+ model: "llama-7b".into(),
	330	+ version: None,
	331	+ routing_key: Some("domain:code".into()),
	332	+ }
	333	+ )
	334	+ .is_none());
	335	+ }
	336	+
	337	+ #[test]
	338	+ fn no_capacity_returns_none() {
	339	+ let cp = cp_with_model();
	340	+ let router = Router::new(RoutingStrategy::LeastLoaded);
	341	+ // Replica exists but is overloaded.
	342	+ router.register_replica(replica("r1", 4, 99.0, &[]));
	343	+ assert!(router
	344	+ .route(
	345	+ &cp,
	346	+ &RouteRequest {
	347	+ model: "llama-7b".into(),
	348	+ version: None,
	349	+ routing_key: None,
	350	+ }
	351	+ )
	352	+ .is_none());
	353	+ }
	354	+
	355	+ #[test]
	356	+ fn unknown_model_returns_none() {
	357	+ let cp = cp_with_model();
	358	+ let router = Router::new(RoutingStrategy::LeastLoaded);
	359	+ assert!(router
	360	+ .route(
	361	+ &cp,
	362	+ &RouteRequest {
	363	+ model: "nope".into(),
	364	+ version: None,
	365	+ routing_key: None,
	366	+ }
	367	+ )
	368	+ .is_none());
	369	+ }
	370	+
	371	+ #[test]
	372	+ fn round_robin_cycles_through_candidates() {
	373	+ let cp = cp_with_model();
	374	+ let router = Router::new(RoutingStrategy::RoundRobin);
	375	+ router.register_replica(replica("a", 0, 10.0, &[]));
	376	+ router.register_replica(replica("b", 0, 10.0, &[]));
	377	+ router.register_replica(replica("c", 0, 10.0, &[]));
	378	+
	379	+ let picks: Vec<String> = (0..6)
	380	+ .map(\|_\| {
	381	+ router
	382	+ .route(
	383	+ &cp,
	384	+ &RouteRequest {
	385	+ model: "llama-7b".into(),
	386	+ version: None,
	387	+ routing_key: None,
	388	+ },
	389	+ )
	390	+ .unwrap()
	391	+ .replica
	392	+ })
	393	+ .collect();
	394	+ // Round-robin over 3 replicas should cycle a, b, c, a, b, c.
	395	+ assert_eq!(picks, vec!["a", "b", "c", "a", "b", "c"]);
	396	+ }
	397	+
	398	+ #[test]
	399	+ fn unhealthy_replica_excluded() {
	400	+ let cp = cp_with_model();
	401	+ let router = Router::new(RoutingStrategy::LeastLoaded);
	402	+ router.register_replica(replica("down", 0, 0.0, &[]));
	403	+ // Mark r1 unhealthy via update_load.
	404	+ router.update_load("down", 0, 0.0, false);
	405	+ assert!(router
	406	+ .route(
	407	+ &cp,
	408	+ &RouteRequest {
	409	+ model: "llama-7b".into(),
	410	+ version: None,
	411	+ routing_key: None,
	412	+ }
	413	+ )
	414	+ .is_none());
	415	+ }
	416	+}
0	417

▼crates/justapi-inference/src/scheduler.rs

+1144-0

@@ -0,0 +1,1144 @@
	1	+//! Continuous-batching scheduler — interleaves prefill and decode across
	2	+//! concurrent sequences to maximise GPU utilisation.
	3	+//!
	4	+//! ## Model
	5	+//!
	6	+//! The [`Scheduler`] owns the block pool and the set of running sequences.
	7	+//! On each step the caller calls [`schedule`](Scheduler::schedule) which
	8	+//! returns a [`Schedule`] describing the work for this forward pass:
	9	+//!
	10	+//! - Prefill sequences consume their prompt (or one chunk of it) and
	11	+//! produce the first token. These are the compute-heavy steps.
	12	+//! - Decode sequences generate one token each. These are memory-bandwidth
	13	+//! bound and benefit greatly from being batched together.
	14	+//!
	15	+//! Finished sequences are removed on the next [`schedule`] call; their blocks
	16	+//! are released to the prefix cache.
	17	+
	18	+use std::collections::{HashMap, VecDeque};
	19	+
	20	+use crate::kv_cache::{BlockId, KvBlockPool, Sequence, BLOCK_SIZE};
	21	+use crate::model::SamplingParams;
	22	+use crate::radix_cache::{RadixPrefixCache, RadixPrefixCacheStats};
	23	+
	24	+// ---------------------------------------------------------------------------
	25	+// Configuration
	26	+// ---------------------------------------------------------------------------
	27	+
	28	+/// Tunable knobs for the scheduler.
	29	+#[derive(Debug, Clone)]
	30	+pub struct SchedulerConfig {
	31	+ /// Maximum number of sequences in a single batch.
	32	+ pub max_num_seqs: usize,
	33	+ /// Maximum number of tokens processed in a single forward pass
	34	+ /// (prefill tokens + decode sequences count toward this budget).
	35	+ pub max_num_batched_tokens: usize,
	36	+ /// Maximum number of tokens a single sequence may generate.
	37	+ pub max_seq_len: usize,
	38	+ /// When `Some(n)`, prompts longer than `n` tokens are split into chunks
	39	+ /// of (at most) `n` tokens so that no single prefill starves decode.
	40	+ /// `None` disables chunking (the entire prompt is prefilled at once).
	41	+ pub chunked_prefill_size: Option<usize>,
	42	+}
	43	+
	44	+impl Default for SchedulerConfig {
	45	+ fn default() -> Self {
	46	+ Self {
	47	+ max_num_seqs: 256,
	48	+ max_num_batched_tokens: 4096,
	49	+ max_seq_len: 4096,
	50	+ chunked_prefill_size: Some(512),
	51	+ }
	52	+ }
	53	+}
	54	+
	55	+// ---------------------------------------------------------------------------
	56	+// Inbound request
	57	+// ---------------------------------------------------------------------------
	58	+
	59	+/// A new sequence arriving from the HTTP / control-plane layer.
	60	+#[derive(Debug, Clone)]
	61	+pub struct NewRequest {
	62	+ /// Opaque caller-provided identifier (echoed back on the response stream).
	63	+ pub id: u64,
	64	+ /// Tokenised prompt (already tokenised by the caller's tokenizer).
	65	+ pub prompt: Vec<u32>,
	66	+ /// Generation parameters.
	67	+ pub sampling_params: SamplingParams,
	68	+ /// Prefix-cache lookup result: the number of tokens that exist in the
	69	+ /// cache so the scheduler can skip computing them.
	70	+ pub prefix_cached_tokens: usize,
	71	+ /// KV-block IDs restored from the prefix cache (length matches
	72	+ /// `prefix_cached_tokens / BLOCK_SIZE`).
	73	+ pub cached_blocks: Vec<BlockId>,
	74	+}
	75	+
	76	+// ---------------------------------------------------------------------------
	77	+// Schedule output
	78	+// ---------------------------------------------------------------------------
	79	+
	80	+/// A single prefill sequence for this forward step.
	81	+#[derive(Debug)]
	82	+pub struct PrefillStep {
	83	+ pub seq_id: u64,
	84	+ /// The tokens to process in this step (one chunk or the full remainder).
	85	+ pub tokens: Vec<u32>,
	86	+ /// How many tokens of this prompt have already been computed (from
	87	+ /// prefix-cache or earlier chunks).
	88	+ pub computed_tokens: usize,
	89	+}
	90	+
	91	+/// The set of work the engine should execute in the next forward pass.
	92	+#[derive(Debug)]
	93	+pub struct Schedule {
	94	+ /// Sequences that need a prefill forward pass.
	95	+ pub prefill: Vec<PrefillStep>,
	96	+ /// Sequence ids that need a single-token decode forward pass.
	97	+ pub decode: Vec<u64>,
	98	+ /// If `true` the scheduler could not admit a waiting request because
	99	+ /// either `max_num_seqs` or `max_num_batched_tokens` would be exceeded.
	100	+ /// The caller can use this to produce a 503 / Retry-After.
	101	+ pub back_pressure: bool,
	102	+}
	103	+
	104	+// ---------------------------------------------------------------------------
	105	+// Scheduler statistics
	106	+// ---------------------------------------------------------------------------
	107	+
	108	+#[derive(Debug, Clone)]
	109	+pub struct SchedulerStats {
	110	+ pub num_waiting: usize,
	111	+ pub num_running: usize,
	112	+ /// How many of the running sequences are still prefilling.
	113	+ pub num_prefilling: usize,
	114	+ /// GPU token budget utilisation for the last schedule.
	115	+ pub budget_used_pct: f32,
	116	+ /// Whether the last schedule hit a resource limit.
	117	+ pub back_pressure: bool,
	118	+ /// Prefix-cache (RadixAttention) reuse effectiveness — the operator-visible
	119	+ /// signal that shared prompts are skipping recomputation.
	120	+ pub prefix: RadixPrefixCacheStats,
	121	+ /// Number of KV blocks currently resident as cached prefix entries (pinned
	122	+ /// in the pool and tracked by the radix tree).
	123	+ pub cached_kv_blocks: usize,
	124	+}
	125	+
	126	+// ---------------------------------------------------------------------------
	127	+// Internal per-sequence state
	128	+// ---------------------------------------------------------------------------
	129	+
	130	+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
	131	+enum SeqState {
	132	+ /// Still processing the prompt (one or more prefill steps remain).
	133	+ Prefilling,
	134	+ /// Prompt fully consumed; generating tokens one at a time.
	135	+ Decoding,
	136	+}
	137	+
	138	+struct RunningSeq {
	139	+ seq: Sequence,
	140	+ /// Full original prompt (kept so we can feed remaining chunks).
	141	+ prompt: Vec<u32>,
	142	+ params: SamplingParams,
	143	+ /// Number of prompt tokens that have been consumed across all prefill
	144	+ /// steps so far (including any prefix-cache restoration).
	145	+ num_prefilled: usize,
	146	+ /// Number of decode tokens produced so far.
	147	+ num_generated: usize,
	148	+ state: SeqState,
	149	+}
	150	+
	151	+// ---------------------------------------------------------------------------
	152	+// Transferable sequence (for disaggregated prefill/decode)
	153	+// ---------------------------------------------------------------------------
	154	+
	155	+/// A sequence that has finished prefill on one scheduler (pool) and is ready
	156	+/// to be handed off to another (e.g. a separate decode scheduler/pool).
	157	+///
	158	+/// Mirrors the KV-cache transfer in disaggregated serving (DistServe /
	159	+/// NVIDIA Dynamo): the prefill GPU computes the prompt's KV cache and ships
	160	+/// it to the decode GPU, which then runs token-by-token generation. Here the
	161	+/// "KV cache" is the logical `Sequence` (block ids + token count); a real
	162	+/// engine would copy the actual tensors and record `num_prefilled` tokens as
	163	+/// the transfer volume.
	164	+#[derive(Debug, Clone)]
	165	+pub struct TransferableSequence {
	166	+ /// The internal sequence id assigned by the source (prefill) scheduler.
	167	+ pub src_seq_id: u64,
	168	+ /// The caller-provided request id (echoed from [`NewRequest::id`]).
	169	+ pub req_id: u64,
	170	+ /// The full tokenised prompt.
	171	+ pub prompt: Vec<u32>,
	172	+ /// Generation parameters.
	173	+ pub params: SamplingParams,
	174	+ /// Number of prompt tokens that were computed during prefill.
	175	+ pub num_prefilled: usize,
	176	+ /// Tokens already produced (the first generated token is emitted by the
	177	+ /// final prefill step, so this is typically 1).
	178	+ pub num_generated: usize,
	179	+}
	180	+
	181	+// ---------------------------------------------------------------------------
	182	+// Scheduler
	183	+// ---------------------------------------------------------------------------
	184	+
	185	+/// Owns the block pool, the wait queue, and the running batch of sequences.
	186	+///
	187	+/// ## Prefix caching
	188	+///
	189	+/// The scheduler owns a [`RadixPrefixCache`] (RadixAttention-style). On
	190	+/// admission it looks up the incoming prompt; any matched prefix reuses the
	191	+/// cached KV blocks (the prompt is not recomputed) and only the unique suffix
	192	+/// is prefilled. When a sequence finishes, its blocks are inserted into the
	193	+/// tree and pinned in the pool so the clock-sweep evictor can never recycle
	194	+/// them out from under the cache. Under memory pressure the scheduler drives
	195	+/// eviction itself via [`RadixPrefixCache::evict_filter`], freeing the LRU
	196	+/// leaf's blocks back to the pool — this keeps a single eviction authority and
	197	+/// avoids stale block ids.
	198	+///
	199	+/// Usage cycle (called from the engine's generation thread):
	200	+///
	201	+/// ```ignore
	202	+/// loop {
	203	+/// let schedule = scheduler.schedule();
	204	+/// if schedule.is_empty() { break; }
	205	+/// // engine.forward(schedule) -> Vec<SequenceEvent>
	206	+/// scheduler.on_step_complete(events);
	207	+/// }
	208	+/// ```
	209	+pub struct Scheduler {
	210	+ config: SchedulerConfig,
	211	+ pool: KvBlockPool,
	212	+ /// RadixAttention-style prefix cache shared across requests.
	213	+ prefix: RadixPrefixCache,
	214	+ /// Number of live (running) sequences currently referencing each
	215	+ /// prefix-cached block. A block with `prefix_refs == 0` is owned purely by
	216	+ /// the tree (pinned in the pool); `== 1` means exactly one live sequence
	217	+ /// re-used it. Admission only reuses a block when its live-ref is 0, which
	218	+ /// keeps a cached block from being claimed by two concurrent sequences.
	219	+ prefix_refs: HashMap<BlockId, usize>,
	220	+ running: Vec<RunningSeq>,
	221	+ wait_queue: VecDeque<NewRequest>,
	222	+ next_seq_id: u64,
	223	+ last_back_pressure: bool,
	224	+}
	225	+
	226	+impl Scheduler {
	227	+ /// Create a new scheduler.
	228	+ pub fn new(config: SchedulerConfig, pool: KvBlockPool) -> Self {
	229	+ Self {
	230	+ config,
	231	+ pool,
	232	+ prefix: RadixPrefixCache::new(),
	233	+ prefix_refs: HashMap::new(),
	234	+ running: Vec::new(),
	235	+ wait_queue: VecDeque::new(),
	236	+ next_seq_id: 1,
	237	+ last_back_pressure: false,
	238	+ }
	239	+ }
	240	+
	241	+ /// Borrow the scheduler's prefix cache (for stats / observability).
	242	+ pub fn prefix_cache(&self) -> &RadixPrefixCache {
	243	+ &self.prefix
	244	+ }
	245	+
	246	+ /// Prefix-cache (RadixAttention) effectiveness snapshot.
	247	+ pub fn prefix_cache_stats(&self) -> RadixPrefixCacheStats {
	248	+ self.prefix.stats()
	249	+ }
	250	+
	251	+ // -- accessors ----------------------------------------------------------
	252	+
	253	+ pub fn config(&self) -> &SchedulerConfig {
	254	+ &self.config
	255	+ }
	256	+
	257	+ pub fn pool(&self) -> &KvBlockPool {
	258	+ &self.pool
	259	+ }
	260	+
	261	+ pub fn pool_mut(&mut self) -> &mut KvBlockPool {
	262	+ &mut self.pool
	263	+ }
	264	+
	265	+ pub fn num_waiting(&self) -> usize {
	266	+ self.wait_queue.len()
	267	+ }
	268	+
	269	+ pub fn num_running(&self) -> usize {
	270	+ self.running.len()
	271	+ }
	272	+
	273	+ /// Snapshot of scheduler state.
	274	+ pub fn stats(&self) -> SchedulerStats {
	275	+ let num_prefilling = self
	276	+ .running
	277	+ .iter()
	278	+ .filter(\|s\| s.state == SeqState::Prefilling)
	279	+ .count();
	280	+ SchedulerStats {
	281	+ num_waiting: self.wait_queue.len(),
	282	+ num_running: self.running.len(),
	283	+ num_prefilling,
	284	+ budget_used_pct: 0.0,
	285	+ back_pressure: self.last_back_pressure,
	286	+ prefix: self.prefix.stats(),
	287	+ cached_kv_blocks: self.pool.cached(),
	288	+ }
	289	+ }
	290	+
	291	+ // -- request admission --------------------------------------------------
	292	+
	293	+ /// Enqueue a new generation request.
	294	+ ///
	295	+ /// Assigns the sequence id the scheduler will use for this request and
	296	+ /// returns it so callers can correlate the request with scheduler events
	297	+ /// (e.g. map a scheduler seq id → a token stream). The request may be
	298	+ /// scheduled immediately if capacity is available, or placed in the wait
	299	+ /// queue.
	300	+ pub fn add_request(&mut self, mut req: NewRequest) -> u64 {
	301	+ let seq_id = self.next_seq_id;
	302	+ self.next_seq_id += 1;
	303	+ req.id = seq_id;
	304	+ self.wait_queue.push_back(req);
	305	+ seq_id
	306	+ }
	307	+
	308	+ /// The scheduler seq ids currently running (for completion detection by an
	309	+ /// external driver).
	310	+ pub fn running_ids(&self) -> Vec<u64> {
	311	+ self.running.iter().map(\|s\| s.seq.id()).collect()
	312	+ }
	313	+
	314	+ /// How many tokens a running sequence has generated so far, if it is
	315	+ /// currently running.
	316	+ pub fn running_seq_generated(&self, seq_id: u64) -> Option<usize> {
	317	+ self.running
	318	+ .iter()
	319	+ .find(\|s\| s.seq.id() == seq_id)
	320	+ .map(\|s\| s.num_generated)
	321	+ }
	322	+
	323	+ // -- scheduling ---------------------------------------------------------
	324	+
	325	+ /// Produce a [`Schedule`] for the next forward pass.
	326	+ ///
	327	+ /// 1. Cache and remove completed sequences (their KV blocks become reusable
	328	+ /// prefix-cache entries for future requests).
	329	+ /// 2. Admit as many waiting requests as the budget allows, reusing any
	330	+ /// cached prefix blocks via the radix tree.
	331	+ /// 3. Assign prefill / decode roles.
	332	+ pub fn schedule(&mut self) -> Schedule {
	333	+ // ---- 1. cache + remove finished / overflow sequences ----
	334	+ // Collect the finished sequences' prompts + blocks before dropping them
	335	+ // so their KV blocks can be promoted into the prefix cache (this also
	336	+ // fixes the block leak the naive `retain` drop would otherwise cause).
	337	+ let mut finished = Vec::new();
	338	+ self.running.retain(\|s\| {
	339	+ let done = s.num_generated >= s.params.max_tokens
	340	+ \|\| s.seq.num_tokens() >= self.config.max_seq_len;
	341	+ if done {
	342	+ finished.push((s.prompt.clone(), s.seq.blocks().to_vec()));
	343	+ false
	344	+ } else {
	345	+ true
	346	+ }
	347	+ });
	348	+ for (prompt, blocks) in &finished {
	349	+ self.cache_completed(prompt, blocks);
	350	+ }
	351	+
	352	+ // ---- 2. admit waiting requests ----
	353	+ let mut blocked = false;
	354	+ while self.running.len() < self.config.max_num_seqs {
	355	+ let Some(req_ref) = self.wait_queue.front() else {
	356	+ break;
	357	+ };
	358	+
	359	+ // Resolve the prefix-cache hit: prefer an explicit caller-provided
	360	+ // hint (`cached_blocks`), otherwise consult the radix tree directly.
	361	+ let (cached_blocks, cached_tokens) = if !req_ref.cached_blocks.is_empty() {
	362	+ (req_ref.cached_blocks.clone(), req_ref.prefix_cached_tokens)
	363	+ } else {
	364	+ match self.prefix.lookup(&req_ref.prompt) {
	365	+ Some((matched, blocks)) => (blocks, matched),
	366	+ None => (Vec::new(), 0),
	367	+ }
	368	+ };
	369	+
	370	+ let total_blocks = req_ref.prompt.len().div_ceil(BLOCK_SIZE);
	371	+ let cached_n = cached_blocks.len();
	372	+ let uncached_blocks = total_blocks.saturating_sub(cached_n);
	373	+ // Admission back-pressure uses a token-count estimate (excluding the
	374	+ // reusable cached prefix) so the heuristic matches prefill cost.
	375	+ let uncached_tokens = req_ref.prompt.len().saturating_sub(cached_tokens);
	376	+ let est_total = uncached_tokens + req_ref.sampling_params.max_tokens;
	377	+ if est_total > self.config.max_num_batched_tokens * 2 {
	378	+ blocked = true;
	379	+ break;
	380	+ }
	381	+ let req = self.wait_queue.pop_front().unwrap();
	382	+
	383	+ // The seq id was assigned in `add_request` and is echoed on
	384	+ // `req.id` so external drivers can correlate scheduler events.
	385	+ let mut seq = Sequence::new(req.id);
	386	+
	387	+ // Reuse cached prefix blocks (skip recomputing them).
	388	+ if !cached_blocks.is_empty() {
	389	+ seq.restore_from_cache(&cached_blocks, cached_tokens, &mut self.pool);
	390	+ for &b in &cached_blocks {
	391	+ // Block is now owned by a live sequence; unpin and record a
	392	+ // live reference so eviction won't free it under us.
	393	+ self.pool.unpin(b);
	394	+ *self.prefix_refs.entry(b).or_insert(0) += 1;
	395	+ }
	396	+ }
	397	+
	398	+ // Allocate the uncached tail, triggering radix eviction if the pool
	399	+ // is exhausted so blocks become available on the retry.
	400	+ let mut alloc_failed = false;
	401	+ for _ in 0..uncached_blocks {
	402	+ if seq.grow(&mut self.pool).is_none() {
	403	+ let freed = self.prefix.evict_filter(uncached_blocks, \|bs\| {
	404	+ bs.iter()
	405	+ .all(\|b\| self.prefix_refs.get(b).copied().unwrap_or(0) == 0)
	406	+ });
	407	+ for b in freed {
	408	+ self.pool.unpin(b);
	409	+ self.pool.free_cached(b);
	410	+ self.prefix_refs.remove(&b);
	411	+ }
	412	+ if seq.grow(&mut self.pool).is_none() {
	413	+ alloc_failed = true;
	414	+ break;
	415	+ }
	416	+ }
	417	+ }
	418	+ if alloc_failed {
	419	+ blocked = true;
	420	+ self.wait_queue.push_front(req);
	421	+ break;
	422	+ }
	423	+
	424	+ let num_prefilled = cached_n * BLOCK_SIZE;
	425	+ let state = if num_prefilled >= req.prompt.len() {
	426	+ SeqState::Decoding
	427	+ } else {
	428	+ SeqState::Prefilling
	429	+ };
	430	+
	431	+ self.running.push(RunningSeq {
	432	+ seq,
	433	+ prompt: req.prompt,
	434	+ params: req.sampling_params,
	435	+ num_prefilled,
	436	+ num_generated: 0,
	437	+ state,
	438	+ });
	439	+ }
	440	+
	441	+ // Back-pressure: true when there are waiting requests that couldn't
	442	+ // be admitted (either blocked by estimate or by capacity).
	443	+ let back_pressure = blocked \|\| !self.wait_queue.is_empty();
	444	+
	445	+ // ---- 3. build schedule ----
	446	+ let mut prefill = Vec::new();
	447	+ let mut decode = Vec::new();
	448	+ let mut token_budget = self.config.max_num_batched_tokens;
	449	+
	450	+ // Decode capacity: at most one token per running sequence that is
	451	+ // in decode mode.
	452	+ let decode_count = self
	453	+ .running
	454	+ .iter()
	455	+ .filter(\|s\| s.state == SeqState::Decoding)
	456	+ .count();
	457	+ let decode_tokens = decode_count.min(token_budget);
	458	+ token_budget = token_budget.saturating_sub(decode_tokens);
	459	+
	460	+ // Prefill capacity: fill remaining budget with prefill steps.
	461	+ for s in &mut self.running {
	462	+ if s.state != SeqState::Prefilling {
	463	+ continue;
	464	+ }
	465	+
	466	+ let remaining = s.prompt.len().saturating_sub(s.num_prefilled);
	467	+ if remaining == 0 {
	468	+ s.state = SeqState::Decoding;
	469	+ continue;
	470	+ }
	471	+
	472	+ let chunk = match self.config.chunked_prefill_size {
	473	+ Some(chunk_size) => remaining.min(chunk_size).min(token_budget),
	474	+ None => remaining.min(token_budget),
	475	+ };
	476	+ if chunk == 0 {
	477	+ break; // budget exhausted
	478	+ }
	479	+
	480	+ let end = s.num_prefilled + chunk;
	481	+ let tokens = s.prompt[s.num_prefilled..end].to_vec();
	482	+ let computed_tokens = s.num_prefilled;
	483	+
	484	+ token_budget = token_budget.saturating_sub(chunk);
	485	+ s.num_prefilled = end;
	486	+ if s.num_prefilled >= s.prompt.len() {
	487	+ // This chunk finished the prompt — next step will be decode.
	488	+ s.state = SeqState::Decoding;
	489	+ }
	490	+
	491	+ prefill.push(PrefillStep {
	492	+ seq_id: s.seq.id(),
	493	+ tokens,
	494	+ computed_tokens,
	495	+ });
	496	+ }
	497	+
	498	+ // Fill decode list.
	499	+ for s in &self.running {
	500	+ if s.state == SeqState::Decoding && decode.len() < decode_tokens {
	501	+ decode.push(s.seq.id());
	502	+ }
	503	+ }
	504	+
	505	+ self.last_back_pressure = back_pressure;
	506	+
	507	+ Schedule {
	508	+ prefill,
	509	+ decode,
	510	+ back_pressure,
	511	+ }
	512	+ }
	513	+
	514	+ // -- step completion ----------------------------------------------------
	515	+
	516	+ /// Promote a finished sequence's KV blocks into the prefix cache.
	517	+ ///
	518	+ /// Inserts the block-aligned prompt prefix into the radix tree (idempotent
	519	+ /// for shared prefixes) and pins the blocks in the pool so the clock-sweep
	520	+ /// evictor cannot recycle them out from under the cache. Blocks still
	521	+ /// referenced by another live sequence are left in-use.
	522	+ fn cache_completed(&mut self, prompt: &[u32], blocks: &[BlockId]) {
	523	+ // Cache only the block-aligned prompt prefix — generated tokens vary
	524	+ // per request and an unaligned tail can't form a fixed shared prefix.
	525	+ let cache_blocks = (prompt.len() / BLOCK_SIZE).min(blocks.len());
	526	+ if cache_blocks == 0 {
	527	+ return;
	528	+ }
	529	+ let cache_tokens = cache_blocks * BLOCK_SIZE;
	530	+ self.prefix
	531	+ .insert(&prompt[..cache_tokens], &blocks[..cache_blocks]);
	532	+ for &b in &blocks[..cache_blocks] {
	533	+ let entry = self.prefix_refs.entry(b).or_insert(0);
	534	+ if *entry == 0 {
	535	+ // Fresh or pure-radix block: this sequence is its sole live
	536	+ // owner. Release to the cached state and pin for the tree.
	537	+ self.pool.release(b, 0);
	538	+ self.pool.pin(b);
	539	+ } else {
	540	+ // Block was reused by this sequence (live-ref == 1). Drop our
	541	+ // hold; if no other live sequence references it, pin it back.
	542	+ *entry -= 1;
	543	+ if *entry == 0 {
	544	+ self.pool.release(b, 0);
	545	+ self.pool.pin(b);
	546	+ }
	547	+ // else: another live sequence still holds it — leave in-use.
	548	+ }
	549	+ }
	550	+ }
	551	+
	552	+ /// Inform the scheduler that a forward step completed.
	553	+ ///
	554	+ /// `seq_id` is the sequence that was processed; `num_new_tokens` is how
	555	+ /// many tokens were generated (1 for decode, >1 for prefill that
	556	+ /// produced multiple tokens via speculation — typically 1 per prefill
	557	+ /// step as well since the first generated token is sampled from logits).
	558	+ pub fn on_step_complete(&mut self, seq_id: u64, num_new_tokens: usize) {
	559	+ if let Some(s) = self.running.iter_mut().find(\|s\| s.seq.id() == seq_id) {
	560	+ s.num_generated += num_new_tokens;
	561	+ s.seq.advance(num_new_tokens);
	562	+
	563	+ // After a prefill step that finished the prompt, the next call
	564	+ // to schedule will treat this sequence as decode.
	565	+ }
	566	+ }
	567	+
	568	+ /// Immediately remove a sequence (client disconnected, error, etc.).
	569	+ ///
	570	+ /// Its blocks are promoted into the prefix cache (reusable by later
	571	+ /// requests) rather than leaked.
	572	+ pub fn cancel(&mut self, seq_id: u64) {
	573	+ if let Some(pos) = self.running.iter().position(\|s\| s.seq.id() == seq_id) {
	574	+ let s = self.running.remove(pos);
	575	+ self.cache_completed(&s.prompt, s.seq.blocks());
	576	+ }
	577	+ }
	578	+
	579	+ /// Return true if no sequences are running or waiting.
	580	+ pub fn is_idle(&self) -> bool {
	581	+ self.running.is_empty() && self.wait_queue.is_empty()
	582	+ }
	583	+
	584	+ // -- disaggregated prefill/decode (KV transfer) -------------------------
	585	+
	586	+ /// Remove every sequence that has finished prefill (transitioned to decode
	587	+ /// with the full prompt consumed) from this scheduler, releasing its blocks
	588	+ /// back to this pool, and return them as [`TransferableSequence`]s ready to
	589	+ /// be admitted to another scheduler.
	590	+ ///
	591	+ /// Sequences still mid-prefill are retained. This is the prefill-side half
	592	+ /// of disaggregated P/D; pair with [`Scheduler::admit_transferred`].
	593	+ pub fn take_completed_prefill(&mut self) -> Vec<TransferableSequence> {
	594	+ let mut out = Vec::new();
	595	+ let mut keep = Vec::new();
	596	+ for s in self.running.drain(..) {
	597	+ let ready = s.state == SeqState::Decoding && s.num_prefilled >= s.prompt.len();
	598	+ if ready {
	599	+ for &b in s.seq.blocks() {
	600	+ // KV is handed to another pool; clear this pool's live-ref
	601	+ // bookkeeping and release the blocks here.
	602	+ self.prefix_refs.remove(&b);
	603	+ self.pool.release(b, 0);
	604	+ }
	605	+ out.push(TransferableSequence {
	606	+ src_seq_id: s.seq.id(),
	607	+ // The caller id is not stored on RunningSeq; echo the
	608	+ // internal seq id so callers can correlate if needed.
	609	+ req_id: s.seq.id(),
	610	+ prompt: s.prompt,
	611	+ params: s.params,
	612	+ num_prefilled: s.num_prefilled,
	613	+ num_generated: s.num_generated,
	614	+ });
	615	+ } else {
	616	+ keep.push(s);
	617	+ }
	618	+ }
	619	+ self.running = keep;
	620	+ out
	621	+ }
	622	+
	623	+ /// Admit a sequence that arrived pre-prefilled from another scheduler
	624	+ /// (see [`Scheduler::take_completed_prefill`]). Allocates fresh blocks in
	625	+ /// this pool for the already-computed token count and resumes decode.
	626	+ pub fn admit_transferred(&mut self, t: TransferableSequence) {
	627	+ let mut seq = Sequence::new(self.next_seq_id);
	628	+ self.next_seq_id += 1;
	629	+ let total_tokens = t.num_prefilled + t.num_generated;
	630	+ let needed_blocks = total_tokens.div_ceil(BLOCK_SIZE);
	631	+ for _ in 0..needed_blocks {
	632	+ if seq.grow(&mut self.pool).is_none() {
	633	+ // Pool exhausted — drop the transfer. Callers must size the
	634	+ // decode pool to accept the expected concurrency.
	635	+ return;
	636	+ }
	637	+ }
	638	+ seq.advance(total_tokens);
	639	+ self.running.push(RunningSeq {
	640	+ seq,
	641	+ prompt: t.prompt,
	642	+ params: t.params,
	643	+ num_prefilled: t.num_prefilled,
	644	+ num_generated: t.num_generated,
	645	+ state: SeqState::Decoding,
	646	+ });
	647	+ }
	648	+}
	649	+
	650	+// ---------------------------------------------------------------------------
	651	+// Tests
	652	+// ---------------------------------------------------------------------------
	653	+
	654	+#[cfg(test)]
	655	+mod tests {
	656	+ use super::*;
	657	+ use crate::model::SamplingParams;
	658	+
	659	+ fn basic_pool() -> KvBlockPool {
	660	+ KvBlockPool::new(1024)
	661	+ }
	662	+
	663	+ fn default_config() -> SchedulerConfig {
	664	+ SchedulerConfig {
	665	+ max_num_seqs: 4,
	666	+ max_num_batched_tokens: 64,
	667	+ max_seq_len: 128,
	668	+ chunked_prefill_size: Some(16),
	669	+ }
	670	+ }
	671	+
	672	+ fn short_prompt(id: u64) -> NewRequest {
	673	+ NewRequest {
	674	+ id,
	675	+ prompt: (0..8).collect(),
	676	+ sampling_params: SamplingParams {
	677	+ max_tokens: 16,
	678	+ ..Default::default()
	679	+ },
	680	+ prefix_cached_tokens: 0,
	681	+ cached_blocks: Vec::new(),
	682	+ }
	683	+ }
	684	+
	685	+ #[test]
	686	+ fn scheduler_starts_idle() {
	687	+ let pool = basic_pool();
	688	+ let sched = Scheduler::new(default_config(), pool);
	689	+ assert!(sched.is_idle());
	690	+ assert_eq!(sched.num_running(), 0);
	691	+ assert_eq!(sched.num_waiting(), 0);
	692	+ }
	693	+
	694	+ #[test]
	695	+ fn schedule_single_prefill() {
	696	+ let mut sched = Scheduler::new(default_config(), basic_pool());
	697	+ sched.add_request(short_prompt(1));
	698	+
	699	+ let s = sched.schedule();
	700	+ assert_eq!(s.prefill.len(), 1);
	701	+ assert!(s.decode.is_empty());
	702	+ assert!(!s.back_pressure);
	703	+
	704	+ let step = &s.prefill[0];
	705	+ assert_eq!(step.seq_id, 1);
	706	+ // Chunked prefill caps at 16, prompt is 8 — fits in one chunk.
	707	+ assert_eq!(step.tokens.len(), 8);
	708	+ assert_eq!(step.computed_tokens, 0);
	709	+ }
	710	+
	711	+ #[test]
	712	+ fn schedule_chunked_prefill() {
	713	+ let cfg = SchedulerConfig {
	714	+ max_num_seqs: 4,
	715	+ max_num_batched_tokens: 64,
	716	+ max_seq_len: 128,
	717	+ chunked_prefill_size: Some(4),
	718	+ };
	719	+ let mut sched = Scheduler::new(cfg, basic_pool());
	720	+ sched.add_request(NewRequest {
	721	+ id: 1,
	722	+ prompt: (0..10).collect(), // 10 tokens
	723	+ sampling_params: SamplingParams {
	724	+ max_tokens: 8,
	725	+ ..Default::default()
	726	+ },
	727	+ prefix_cached_tokens: 0,
	728	+ cached_blocks: Vec::new(),
	729	+ });
	730	+
	731	+ // First schedule: chunk of 4 tokens.
	732	+ let s = sched.schedule();
	733	+ assert_eq!(s.prefill.len(), 1);
	734	+ assert_eq!(s.prefill[0].tokens.len(), 4);
	735	+ assert_eq!(s.prefill[0].computed_tokens, 0);
	736	+
	737	+ // After step complete: 0 generated (prefill just computed KV).
	738	+ sched.on_step_complete(1, 1);
	739	+
	740	+ // Second schedule: next chunk of 4.
	741	+ let s = sched.schedule();
	742	+ assert_eq!(s.prefill.len(), 1);
	743	+ assert_eq!(s.prefill[0].tokens.len(), 4);
	744	+ assert_eq!(s.prefill[0].computed_tokens, 4);
	745	+
	746	+ sched.on_step_complete(1, 1);
	747	+
	748	+ // Third schedule: remaining 2 tokens, then transitions to decode.
	749	+ let s = sched.schedule();
	750	+ assert_eq!(s.prefill.len(), 1);
	751	+ assert_eq!(s.prefill[0].tokens.len(), 2);
	752	+ assert_eq!(s.prefill[0].computed_tokens, 8);
	753	+
	754	+ sched.on_step_complete(1, 1);
	755	+
	756	+ // Now the sequence is in decode mode.
	757	+ let s = sched.schedule();
	758	+ assert!(s.prefill.is_empty());
	759	+ assert_eq!(s.decode.len(), 1);
	760	+ assert_eq!(s.decode[0], 1);
	761	+ }
	762	+
	763	+ #[test]
	764	+ fn schedule_multiple_sequences_interleaved() {
	765	+ let cfg = SchedulerConfig {
	766	+ max_num_seqs: 4,
	767	+ max_num_batched_tokens: 32,
	768	+ max_seq_len: 128,
	769	+ chunked_prefill_size: Some(8),
	770	+ };
	771	+ let mut sched = Scheduler::new(cfg, basic_pool());
	772	+
	773	+ // Add two requests with different-length prompts.
	774	+ // Internal seq_ids are assigned sequentially (1, 2, 3, …).
	775	+ sched.add_request(short_prompt(10)); // internal seq_id = 1
	776	+ sched.add_request(NewRequest {
	777	+ id: 20,
	778	+ prompt: (0..16).collect(),
	779	+ sampling_params: SamplingParams {
	780	+ max_tokens: 8,
	781	+ ..Default::default()
	782	+ },
	783	+ prefix_cached_tokens: 0,
	784	+ cached_blocks: Vec::new(),
	785	+ }); // internal seq_id = 2
	786	+
	787	+ // Batch has capacity for both in one step.
	788	+ let s = sched.schedule();
	789	+ assert_eq!(s.prefill.len(), 2);
	790	+ assert!(s.decode.is_empty());
	791	+
	792	+ // Complete first prefill for both.
	793	+ for step in &s.prefill {
	794	+ sched.on_step_complete(step.seq_id, 1);
	795	+ }
	796	+
	797	+ // Second step: seq 1 finishes prefill (prompt=8, all consumed),
	798	+ // seq 2 gets second chunk (prompt=16, 8 consumed, 8 remaining).
	799	+ let s = sched.schedule();
	800	+ assert_eq!(s.prefill.len(), 1);
	801	+ assert_eq!(s.prefill[0].seq_id, 2); // internal seq_id
	802	+ assert_eq!(s.decode.len(), 1);
	803	+ assert_eq!(s.decode[0], 1); // internal seq_id
	804	+ }
	805	+
	806	+ #[test]
	807	+ fn schedule_back_pressure_when_full() {
	808	+ let mut sched = Scheduler::new(
	809	+ SchedulerConfig {
	810	+ max_num_seqs: 1,
	811	+ max_num_batched_tokens: 64,
	812	+ max_seq_len: 128,
	813	+ chunked_prefill_size: None,
	814	+ },
	815	+ basic_pool(),
	816	+ );
	817	+
	818	+ sched.add_request(short_prompt(1));
	819	+ let _ = sched.schedule();
	820	+ assert!(!sched.stats().back_pressure);
	821	+
	822	+ // Second request must wait.
	823	+ sched.add_request(short_prompt(2));
	824	+ let s = sched.schedule();
	825	+ // First is still running (just prefilled, not finished).
	826	+ assert!(s.back_pressure);
	827	+ assert_eq!(sched.num_waiting(), 1);
	828	+ }
	829	+
	830	+ #[test]
	831	+ fn schedule_finished_sequence_is_removed() {
	832	+ let mut sched = Scheduler::new(
	833	+ SchedulerConfig {
	834	+ max_num_seqs: 4,
	835	+ max_num_batched_tokens: 64,
	836	+ max_seq_len: 128,
	837	+ chunked_prefill_size: None,
	838	+ },
	839	+ basic_pool(),
	840	+ );
	841	+
	842	+ sched.add_request(NewRequest {
	843	+ id: 1,
	844	+ prompt: (0..4).collect(),
	845	+ sampling_params: SamplingParams {
	846	+ max_tokens: 3,
	847	+ ..Default::default()
	848	+ },
	849	+ prefix_cached_tokens: 0,
	850	+ cached_blocks: Vec::new(),
	851	+ });
	852	+
	853	+ // Prefill step (generates 1 token).
	854	+ let s = sched.schedule();
	855	+ assert_eq!(s.prefill.len(), 1);
	856	+ sched.on_step_complete(1, 1);
	857	+
	858	+ // Decode 1 (2nd token).
	859	+ let s = sched.schedule();
	860	+ assert_eq!(s.decode.len(), 1);
	861	+ sched.on_step_complete(1, 1);
	862	+
	863	+ // Decode 2 (3rd token = max_tokens = 3).
	864	+ let s = sched.schedule();
	865	+ assert_eq!(s.decode.len(), 1);
	866	+ sched.on_step_complete(1, 1);
	867	+
	868	+ // Next schedule should remove it (num_generated >= max_tokens).
	869	+ let s = sched.schedule();
	870	+ assert!(s.prefill.is_empty());
	871	+ assert!(s.decode.is_empty());
	872	+ assert_eq!(sched.num_running(), 0);
	873	+ }
	874	+
	875	+ #[test]
	876	+ fn schedule_cancel_sequence() {
	877	+ let mut sched = Scheduler::new(default_config(), basic_pool());
	878	+ sched.add_request(short_prompt(1));
	879	+ let _ = sched.schedule();
	880	+ assert_eq!(sched.num_running(), 1);
	881	+
	882	+ sched.cancel(1);
	883	+ assert_eq!(sched.num_running(), 0);
	884	+ assert!(sched.is_idle());
	885	+ }
	886	+
	887	+ #[test]
	888	+ fn schedule_prefix_cache_restoration() {
	889	+ let mut pool = KvBlockPool::new(64);
	890	+ // Pre-allocate 2 blocks and assign them as "cached".
	891	+ let b0 = pool.alloc().unwrap();
	892	+ let b1 = pool.alloc().unwrap();
	893	+ // Release to simulate them being in the prefix cache.
	894	+ pool.release(b0, 42);
	895	+ pool.release(b1, 43);
	896	+
	897	+ let mut sched = Scheduler::new(
	898	+ SchedulerConfig {
	899	+ max_num_seqs: 4,
	900	+ max_num_batched_tokens: 64,
	901	+ max_seq_len: 128,
	902	+ chunked_prefill_size: Some(16),
	903	+ },
	904	+ pool,
	905	+ );
	906	+
	907	+ sched.add_request(NewRequest {
	908	+ id: 1,
	909	+ prompt: (0..48).collect(),
	910	+ sampling_params: SamplingParams {
	911	+ max_tokens: 8,
	912	+ ..Default::default()
	913	+ },
	914	+ prefix_cached_tokens: 32, // 2 blocks × 16 tokens
	915	+ cached_blocks: vec![b0, b1],
	916	+ });
	917	+
	918	+ let s = sched.schedule();
	919	+ assert_eq!(s.prefill.len(), 1);
	920	+ // Only the uncached remainder (48 - 32 = 16) is prefilled.
	921	+ assert_eq!(s.prefill[0].tokens.len(), 16);
	922	+ assert_eq!(s.prefill[0].computed_tokens, 32);
	923	+ }
	924	+
	925	+ #[test]
	926	+ fn schedule_budget_limits_batch_size() {
	927	+ let cfg = SchedulerConfig {
	928	+ max_num_seqs: 10,
	929	+ max_num_batched_tokens: 12,
	930	+ max_seq_len: 128,
	931	+ chunked_prefill_size: Some(8),
	932	+ };
	933	+ let mut sched = Scheduler::new(cfg, basic_pool());
	934	+
	935	+ // Add many requests with 8-token prompts.
	936	+ for i in 0..5 {
	937	+ sched.add_request(short_prompt(i + 1));
	938	+ }
	939	+
	940	+ let s = sched.schedule();
	941	+ // Budget = 12. First prefill consumes 8, leaving 4 budget.
	942	+ // Second prefill gets a 4-token chunk (partial).
	943	+ assert_eq!(s.prefill.len(), 2);
	944	+ assert_eq!(s.prefill[0].tokens.len(), 8);
	945	+ assert_eq!(s.prefill[1].tokens.len(), 4);
	946	+ // Wait queue is empty — all 5 were admitted into running.
	947	+ // The token budget only limits how many prefill in one step, not admission.
	948	+ assert!(sched.num_waiting() == 0);
	949	+ }
	950	+
	951	+ // ---- radix prefix-cache integration ----------------------------------
	952	+
	953	+ /// Config that prefills the whole prompt in one step (no chunking) so a
	954	+ /// request finishes its prefix in a single forward pass.
	955	+ fn full_prefix_config() -> SchedulerConfig {
	956	+ SchedulerConfig {
	957	+ max_num_seqs: 4,
	958	+ max_num_batched_tokens: 256,
	959	+ max_seq_len: 4096,
	960	+ chunked_prefill_size: None,
	961	+ }
	962	+ }
	963	+
	964	+ /// Drive the scheduler to idle, simulating one generated token per step.
	965	+ fn drain(sched: &mut Scheduler) {
	966	+ let mut guard = 0;
	967	+ while !sched.is_idle() && guard < 2000 {
	968	+ let s = sched.schedule();
	969	+ for p in &s.prefill {
	970	+ sched.on_step_complete(p.seq_id, 1);
	971	+ }
	972	+ for &d in &s.decode {
	973	+ sched.on_step_complete(d, 1);
	974	+ }
	975	+ guard += 1;
	976	+ }
	977	+ }
	978	+
	979	+ #[test]
	980	+ fn schedule_radix_reuses_shared_prefix() {
	981	+ // Two requests with an identical 48-token prompt. The first computes
	982	+ // its prefix; the second must reuse it from the radix cache instead of
	983	+ // recomputing (no prefill step, hit recorded).
	984	+ let mut sched = Scheduler::new(full_prefix_config(), KvBlockPool::new(64));
	985	+
	986	+ // Request 1: full prompt, 1 generated token.
	987	+ sched.add_request(NewRequest {
	988	+ id: 1,
	989	+ prompt: (0..48).collect(),
	990	+ sampling_params: SamplingParams {
	991	+ max_tokens: 1,
	992	+ ..Default::default()
	993	+ },
	994	+ prefix_cached_tokens: 0,
	995	+ cached_blocks: Vec::new(),
	996	+ });
	997	+ let s = sched.schedule();
	998	+ assert_eq!(s.prefill.len(), 1);
	999	+ assert_eq!(s.prefill[0].tokens.len(), 48);
	1000	+ assert_eq!(s.prefill[0].computed_tokens, 0);
	1001	+ sched.on_step_complete(1, 1); // generated the 1st token → finished
	1002	+
	1003	+ // Run again so request 1 is dropped and cached.
	1004	+ let _ = sched.schedule();
	1005	+ assert_eq!(sched.prefix_cache().stats().nodes, 1);
	1006	+
	1007	+ // Request 2: same prompt → full prefix cache hit.
	1008	+ sched.add_request(NewRequest {
	1009	+ id: 2,
	1010	+ prompt: (0..48).collect(),
	1011	+ sampling_params: SamplingParams {
	1012	+ max_tokens: 1,
	1013	+ ..Default::default()
	1014	+ },
	1015	+ prefix_cached_tokens: 0,
	1016	+ cached_blocks: Vec::new(),
	1017	+ });
	1018	+ let s = sched.schedule();
	1019	+ // Nothing to prefill; the cached prefix covers the whole prompt.
	1020	+ assert!(
	1021	+ s.prefill.is_empty(),
	1022	+ "shared prefix should be reused, not prefilled"
	1023	+ );
	1024	+ assert_eq!(s.decode.len(), 1);
	1025	+ // The hit should have been recorded and 48 tokens saved.
	1026	+ let stats = sched.prefix_cache().stats();
	1027	+ assert_eq!(stats.hits, 1);
	1028	+ assert_eq!(stats.tokens_saved, 48);
	1029	+ }
	1030	+
	1031	+ #[test]
	1032	+ fn schedule_radix_reuses_partial_prefix() {
	1033	+ // Request 2 extends request 1's prompt by one block. Only the new
	1034	+ // block should be prefilled (computed_tokens == 48).
	1035	+ let mut sched = Scheduler::new(full_prefix_config(), KvBlockPool::new(64));
	1036	+
	1037	+ sched.add_request(NewRequest {
	1038	+ id: 1,
	1039	+ prompt: (0..48).collect(), // 3 blocks
	1040	+ sampling_params: SamplingParams {
	1041	+ max_tokens: 1,
	1042	+ ..Default::default()
	1043	+ },
	1044	+ prefix_cached_tokens: 0,
	1045	+ cached_blocks: Vec::new(),
	1046	+ });
	1047	+ let s = sched.schedule();
	1048	+ sched.on_step_complete(s.prefill[0].seq_id, 1);
	1049	+ let _ = sched.schedule(); // cache request 1
	1050	+
	1051	+ // Request 2: 4 blocks, first 3 identical to request 1.
	1052	+ sched.add_request(NewRequest {
	1053	+ id: 2,
	1054	+ prompt: (0..64).collect(),
	1055	+ sampling_params: SamplingParams {
	1056	+ max_tokens: 1,
	1057	+ ..Default::default()
	1058	+ },
	1059	+ prefix_cached_tokens: 0,
	1060	+ cached_blocks: Vec::new(),
	1061	+ });
	1062	+ let s = sched.schedule();
	1063	+ assert_eq!(s.prefill.len(), 1);
	1064	+ // Only the unique 4th block (16 tokens) is prefilled.
	1065	+ assert_eq!(s.prefill[0].tokens.len(), 16);
	1066	+ assert_eq!(s.prefill[0].computed_tokens, 48);
	1067	+ }
	1068	+
	1069	+ #[test]
	1070	+ fn schedule_radix_evicts_under_pressure() {
	1071	+ // Tiny pool: 6 blocks total. Each request needs 3 blocks and distinct
	1072	+ // prompts (no cache hits), so the scheduler must evict cached prefixes
	1073	+ // to make room. It must not deadlock or panic.
	1074	+ let pool = KvBlockPool::new(6);
	1075	+ let mut sched = Scheduler::new(full_prefix_config(), pool);
	1076	+
	1077	+ for i in 0..10u64 {
	1078	+ // Distinct 48-token prompts: shift the token range per request.
	1079	+ let base = i as u32 * 100;
	1080	+ let prompt: Vec<u32> = (base..base + 48).collect();
	1081	+ sched.add_request(NewRequest {
	1082	+ id: i,
	1083	+ prompt,
	1084	+ sampling_params: SamplingParams {
	1085	+ max_tokens: 1,
	1086	+ ..Default::default()
	1087	+ },
	1088	+ prefix_cached_tokens: 0,
	1089	+ cached_blocks: Vec::new(),
	1090	+ });
	1091	+ }
	1092	+
	1093	+ drain(&mut sched);
	1094	+ // All 10 requests were admitted, computed, and cached without deadlock.
	1095	+ assert!(sched.is_idle());
	1096	+ assert_eq!(sched.num_running(), 0);
	1097	+ assert_eq!(sched.num_waiting(), 0);
	1098	+ }
	1099	+
	1100	+ #[test]
	1101	+ fn schedule_prefix_stats_reflect_reuse() {
	1102	+ // The SchedulerStats surface must expose prefix reuse so a metrics
	1103	+ // layer can observe it without reaching into the cache directly.
	1104	+ let mut sched = Scheduler::new(full_prefix_config(), KvBlockPool::new(64));
	1105	+
	1106	+ sched.add_request(NewRequest {
	1107	+ id: 1,
	1108	+ prompt: (0..48).collect(),
	1109	+ sampling_params: SamplingParams {
	1110	+ max_tokens: 1,
	1111	+ ..Default::default()
	1112	+ },
	1113	+ prefix_cached_tokens: 0,
	1114	+ cached_blocks: Vec::new(),
	1115	+ });
	1116	+ let s = sched.schedule();
	1117	+ sched.on_step_complete(s.prefill[0].seq_id, 1);
	1118	+ let _ = sched.schedule(); // cache request 1
	1119	+
	1120	+ // After caching, resident prefix KV blocks > 0 and stats are visible.
	1121	+ let cached = sched.stats();
	1122	+ assert!(cached.cached_kv_blocks > 0, "prefix KV should be resident");
	1123	+ assert_eq!(cached.prefix.nodes, 1);
	1124	+ assert_eq!(cached.prefix.hits, 0);
	1125	+ assert_eq!(cached.prefix.hash_collisions, 0);
	1126	+
	1127	+ sched.add_request(NewRequest {
	1128	+ id: 2,
	1129	+ prompt: (0..48).collect(),
	1130	+ sampling_params: SamplingParams {
	1131	+ max_tokens: 1,
	1132	+ ..Default::default()
	1133	+ },
	1134	+ prefix_cached_tokens: 0,
	1135	+ cached_blocks: Vec::new(),
	1136	+ });
	1137	+ let _ = sched.schedule(); // reuses the cached prefix
	1138	+
	1139	+ let reused = sched.stats();
	1140	+ assert_eq!(reused.prefix.hits, 1);
	1141	+ assert_eq!(reused.prefix.tokens_saved, 48);
	1142	+ assert_eq!(reused.prefix.hash_collisions, 0);
	1143	+ }
	1144	+}
0	1145

▼crates/justapi-inference/src/scheduler_engine.rs

+693-0

@@ -0,0 +1,693 @@
	1	+//! Bridges the [`Scheduler`] (continuous batching + prefix cache) with the
	2	+//! [`Engine`] (model registry) for scheduler-backed token generation.
	3	+//!
	4	+//! Unlike the naive `Engine::generate()` loop (which runs one request to
	5	+//! completion inside [`Model::generate`]), the [`SchedulerEngine`] runs a
	6	+//! single persistent scheduler thread that interleaves all in-flight
	7	+//! requests. This is true continuous batching: concurrent requests share the
	8	+//! same `schedule()` → `forward_logits()` → `on_step_complete()` loop, exactly
	9	+//! like vLLM's `SchedulerLoop`.
	10	+//!
	11	+//! # Architecture
	12	+//!
	13	+//! [`SchedulerEngine::new`] spawns the scheduler loop thread. Each call to
	14	+//! [`generate`](SchedulerEngine::generate):
	15	+//!
	16	+//! 1. Resolves the model handle and assigns the request a scheduler seq id
	17	+//! via [`Scheduler::add_request`].
	18	+//! 2. Registers a per-request `mpsc` channel keyed by seq id.
	19	+//! 3. Returns the receiving end; the loop streams tokens into it.
	20	+//!
	21	+//! The loop thread:
	22	+//! - Calls [`Scheduler::schedule`] to get prefill/decode work.
	23	+//! - Feeds each step's context through [`Model::forward_logits`] and samples
	24	+//! one token (greedy argmax).
	25	+//! - Routes tokens to the right per-request channel.
	26	+//! - Calls [`Scheduler::on_step_complete`] and, when a sequence leaves the
	27	+//! running set, sends a final `finish_reason` token and closes its channel.
	28	+
	29	+use std::collections::HashMap;
	30	+use std::sync::atomic::{AtomicBool, Ordering};
	31	+use std::sync::{Arc, Mutex};
	32	+use std::time::Duration;
	33	+
	34	+use tokio::sync::mpsc;
	35	+
	36	+use crate::engine::Engine;
	37	+use crate::model::{FinishReason, GeneratedToken, Model, ModelError, SamplingParams};
	38	+use crate::scheduler::{NewRequest, Scheduler, SchedulerStats};
	39	+use crate::MockModel;
	40	+
	41	+/// Wraps an [`Engine`] and a [`Scheduler`] to produce a scheduler-backed
	42	+/// token stream with continuous batching across concurrent requests.
	43	+pub struct SchedulerEngine {
	44	+ engine: Arc<Engine>,
	45	+ scheduler: Arc<Mutex<Scheduler>>,
	46	+ /// Per-seq-id token sinks (one per in-flight request).
	47	+ senders: Arc<Mutex<HashMap<u64, mpsc::UnboundedSender<GeneratedToken>>>>,
	48	+ /// Per-seq-id model handle (a request may target any registered model).
	49	+ seq_model: Arc<Mutex<HashMap<u64, Arc<dyn Model>>>>,
	50	+ /// Per-seq-id token context for forward passes.
	51	+ seq_context: Arc<Mutex<HashMap<u64, Vec<u32>>>>,
	52	+ /// Per-seq-id `max_tokens` so we can emit a finish reason.
	53	+ seq_max_tokens: Arc<Mutex<HashMap<u64, usize>>>,
	54	+ /// Per-seq-id sampling params (temperature / top_p / top_k) for sampling.
	55	+ seq_params: Arc<Mutex<HashMap<u64, SamplingParams>>>,
	56	+ /// Liveness flag for the scheduler loop thread.
	57	+ alive: Arc<AtomicBool>,
	58	+}
	59	+
	60	+impl SchedulerEngine {
	61	+ /// Create a new scheduler-backed engine and start its scheduler loop.
	62	+ pub fn new(engine: Arc<Engine>, scheduler: Arc<Mutex<Scheduler>>) -> Self {
	63	+ let senders = Arc::new(Mutex::new(HashMap::new()));
	64	+ let seq_model = Arc::new(Mutex::new(HashMap::new()));
	65	+ let seq_context = Arc::new(Mutex::new(HashMap::new()));
	66	+ let seq_max_tokens = Arc::new(Mutex::new(HashMap::new()));
	67	+ let seq_params = Arc::new(Mutex::new(HashMap::new()));
	68	+ let alive = Arc::new(AtomicBool::new(true));
	69	+
	70	+ let se = Self {
	71	+ engine,
	72	+ scheduler,
	73	+ senders,
	74	+ seq_model,
	75	+ seq_context,
	76	+ seq_max_tokens,
	77	+ seq_params,
	78	+ alive,
	79	+ };
	80	+ se.spawn_loop();
	81	+ se
	82	+ }
	83	+
	84	+ /// Borrow the engine.
	85	+ pub fn engine(&self) -> &Arc<Engine> {
	86	+ &self.engine
	87	+ }
	88	+
	89	+ /// Borrow the scheduler.
	90	+ pub fn scheduler(&self) -> &Arc<Mutex<Scheduler>> {
	91	+ &self.scheduler
	92	+ }
	93	+
	94	+ /// Snapshot of scheduler statistics (for metrics / observability).
	95	+ pub fn stats(&self) -> SchedulerStats {
	96	+ self.scheduler.lock().unwrap().stats()
	97	+ }
	98	+
	99	+ /// Generate tokens for a prompt, streaming them over an unbounded channel.
	100	+ pub fn generate(
	101	+ &self,
	102	+ model_name: &str,
	103	+ prompt: &[u32],
	104	+ params: SamplingParams,
	105	+ ) -> Result<mpsc::UnboundedReceiver<GeneratedToken>, ModelError> {
	106	+ let model = self
	107	+ .engine
	108	+ .get(model_name)
	109	+ .ok_or_else(\|\| ModelError::NotFound(model_name.to_string()))?;
	110	+
	111	+ let (tx, rx) = mpsc::unbounded_channel();
	112	+
	113	+ // Assign a scheduler seq id up front so we can route its tokens.
	114	+ let seq_id = {
	115	+ let mut sched = self.scheduler.lock().unwrap();
	116	+ sched.add_request(NewRequest {
	117	+ id: 0,
	118	+ prompt: prompt.to_vec(),
	119	+ sampling_params: params.clone(),
	120	+ prefix_cached_tokens: 0,
	121	+ cached_blocks: Vec::new(),
	122	+ })
	123	+ };
	124	+
	125	+ self.senders.lock().unwrap().insert(seq_id, tx);
	126	+ self.seq_model.lock().unwrap().insert(seq_id, model.clone());
	127	+ self.seq_max_tokens
	128	+ .lock()
	129	+ .unwrap()
	130	+ .insert(seq_id, params.max_tokens);
	131	+ self.seq_params
	132	+ .lock()
	133	+ .unwrap()
	134	+ .insert(seq_id, params.clone());
	135	+
	136	+ Ok(rx)
	137	+ }
	138	+
	139	+ /// Start the persistent scheduler loop thread.
	140	+ fn spawn_loop(&self) {
	141	+ let scheduler = self.scheduler.clone();
	142	+ let senders = self.senders.clone();
	143	+ let seq_model = self.seq_model.clone();
	144	+ let seq_context = self.seq_context.clone();
	145	+ let seq_max_tokens = self.seq_max_tokens.clone();
	146	+ let seq_params = self.seq_params.clone();
	147	+ let alive = self.alive.clone();
	148	+
	149	+ std::thread::spawn(move \|\| {
	150	+ while alive.load(Ordering::Relaxed) {
	151	+ let before: Vec<u64> = {
	152	+ let sched = scheduler.lock().unwrap();
	153	+ sched.running_ids()
	154	+ };
	155	+
	156	+ let schedule = {
	157	+ let mut sched = scheduler.lock().unwrap();
	158	+ sched.schedule()
	159	+ };
	160	+
	161	+ let empty = schedule.prefill.is_empty() && schedule.decode.is_empty();
	162	+
	163	+ if !empty {
	164	+ // --- Prefill steps ---
	165	+ for step in &schedule.prefill {
	166	+ let params = get_seq_params(&seq_params, step.seq_id);
	167	+ let token = match forward_for_seq(
	168	+ step.seq_id,
	169	+ &step.tokens,
	170	+ &seq_model,
	171	+ &seq_context,
	172	+ &params,
	173	+ ) {
	174	+ Some(t) => t,
	175	+ None => continue,
	176	+ };
	177	+ let text = MockModel::detokenize(&[token]);
	178	+
	179	+ let done = is_seq_done(step.seq_id, &scheduler, &seq_max_tokens);
	180	+ deliver(
	181	+ &senders,
	182	+ step.seq_id,
	183	+ GeneratedToken {
	184	+ id: token,
	185	+ text,
	186	+ logprob: 0.0,
	187	+ finish_reason: if done {
	188	+ Some(FinishReason::Length)
	189	+ } else {
	190	+ None
	191	+ },
	192	+ },
	193	+ );
	194	+ }
	195	+
	196	+ // --- Decode steps ---
	197	+ for &seq_id in &schedule.decode {
	198	+ let params = get_seq_params(&seq_params, seq_id);
	199	+ let token = match forward_decode(seq_id, &seq_model, &seq_context, &params)
	200	+ {
	201	+ Some(t) => t,
	202	+ None => continue,
	203	+ };
	204	+ let text = MockModel::detokenize(&[token]);
	205	+
	206	+ let done = is_seq_done(seq_id, &scheduler, &seq_max_tokens);
	207	+ deliver(
	208	+ &senders,
	209	+ seq_id,
	210	+ GeneratedToken {
	211	+ id: token,
	212	+ text,
	213	+ logprob: 0.0,
	214	+ finish_reason: if done {
	215	+ Some(FinishReason::Length)
	216	+ } else {
	217	+ None
	218	+ },
	219	+ },
	220	+ );
	221	+ }
	222	+
	223	+ // --- Notify scheduler of completed steps ---
	224	+ {
	225	+ let mut sched = scheduler.lock().unwrap();
	226	+ for step in &schedule.prefill {
	227	+ sched.on_step_complete(step.seq_id, 1);
	228	+ }
	229	+ for &seq_id in &schedule.decode {
	230	+ sched.on_step_complete(seq_id, 1);
	231	+ }
	232	+ }
	233	+ }
	234	+
	235	+ // --- Detect completions (seq left the running set) ---
	236	+ // Runs on EVERY iteration, including when `schedule()` returned
	237	+ // no work (that is exactly when a finished sequence was removed
	238	+ // and must receive its final finish token).
	239	+ let after: Vec<u64> = {
	240	+ let sched = scheduler.lock().unwrap();
	241	+ sched.running_ids()
	242	+ };
	243	+ for id in &before {
	244	+ if !after.contains(id) {
	245	+ deliver(
	246	+ &senders,
	247	+ *id,
	248	+ GeneratedToken {
	249	+ id: 0,
	250	+ text: String::new(),
	251	+ logprob: 0.0,
	252	+ finish_reason: Some(FinishReason::Length),
	253	+ },
	254	+ );
	255	+ senders.lock().unwrap().remove(id);
	256	+ seq_model.lock().unwrap().remove(id);
	257	+ seq_context.lock().unwrap().remove(id);
	258	+ seq_max_tokens.lock().unwrap().remove(id);
	259	+ seq_params.lock().unwrap().remove(id);
	260	+ }
	261	+ }
	262	+
	263	+ if empty {
	264	+ // Nothing to do this cycle; idle briefly to avoid a hot spin.
	265	+ std::thread::sleep(Duration::from_millis(2));
	266	+ }
	267	+ }
	268	+ });
	269	+ }
	270	+}
	271	+
	272	+/// Fetch a sequence's sampling params (defaults to greedy if unknown).
	273	+fn get_seq_params(
	274	+ seq_params: &Arc<Mutex<HashMap<u64, SamplingParams>>>,
	275	+ seq_id: u64,
	276	+) -> SamplingParams {
	277	+ seq_params
	278	+ .lock()
	279	+ .unwrap()
	280	+ .get(&seq_id)
	281	+ .cloned()
	282	+ .unwrap_or_default()
	283	+}
	284	+
	285	+/// Sample a token id from logits according to `params`.
	286	+///
	287	+/// - `temperature == 0` (or effectively 0): greedy argmax.
	288	+/// - Otherwise: scale logits by `1/temperature`, optionally keep only the top-`k`
	289	+/// logits, apply top-`p` (nucleus) truncation, then draw from the categorical
	290	+/// distribution. `rand` provides the stochastic draw.
	291	+fn sample_token_with_params(logits: &[f32], params: &SamplingParams) -> u32 {
	292	+ if params.temperature <= f32::EPSILON {
	293	+ return argmax(logits);
	294	+ }
	295	+
	296	+ let inv_temp = 1.0 / params.temperature;
	297	+ let mut scaled: Vec<f32> = logits.iter().map(\|&l\| l * inv_temp).collect();
	298	+
	299	+ // Top-k: zero out everything outside the k largest logits.
	300	+ if params.top_k > 0 && params.top_k < scaled.len() {
	301	+ let mut idx: Vec<usize> = (0..scaled.len()).collect();
	302	+ let k = params.top_k;
	303	+ idx.select_nth_unstable_by(k - 1, \|&a, &b\| {
	304	+ scaled[b]
	305	+ .partial_cmp(&scaled[a])
	306	+ .unwrap_or(std::cmp::Ordering::Equal)
	307	+ });
	308	+ for &i in &idx[k..] {
	309	+ scaled[i] = f32::NEG_INFINITY;
	310	+ }
	311	+ }
	312	+
	313	+ // Numerically stable softmax.
	314	+ let max = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
	315	+ let mut probs: Vec<f32> = scaled.iter().map(\|&s\| (s - max).exp()).collect();
	316	+ let mut sum = probs.iter().sum::<f32>();
	317	+ if sum <= 0.0 {
	318	+ return argmax(logits);
	319	+ }
	320	+ for p in probs.iter_mut() {
	321	+ *p /= sum;
	322	+ }
	323	+
	324	+ // Top-p (nucleus): keep the smallest prefix of the sorted distribution that
	325	+ // reaches cumulative mass `top_p`, renormalise.
	326	+ if params.top_p < 1.0 {
	327	+ let mut order: Vec<usize> = (0..probs.len()).collect();
	328	+ order.sort_by(\|&a, &b\| {
	329	+ probs[b]
	330	+ .partial_cmp(&probs[a])
	331	+ .unwrap_or(std::cmp::Ordering::Equal)
	332	+ });
	333	+ let mut cum = 0.0_f32;
	334	+ for (rank, &i) in order.iter().enumerate() {
	335	+ cum += probs[i];
	336	+ if cum >= params.top_p {
	337	+ // Zero everything beyond this rank.
	338	+ for &j in &order[rank + 1..] {
	339	+ probs[j] = 0.0;
	340	+ }
	341	+ break;
	342	+ }
	343	+ }
	344	+ sum = probs.iter().sum::<f32>();
	345	+ if sum > 0.0 {
	346	+ for p in probs.iter_mut() {
	347	+ *p /= sum;
	348	+ }
	349	+ }
	350	+ }
	351	+
	352	+ // Draw from the categorical distribution.
	353	+ let mut rng = rand::thread_rng();
	354	+ let r = rand::Rng::gen::<f32>(&mut rng);
	355	+ let mut cum = 0.0_f32;
	356	+ for (i, &p) in probs.iter().enumerate() {
	357	+ cum += p;
	358	+ if r < cum {
	359	+ return i as u32;
	360	+ }
	361	+ }
	362	+ argmax(logits)
	363	+}
	364	+
	365	+/// Greedy argmax: index of the maximum logit value.
	366	+fn argmax(logits: &[f32]) -> u32 {
	367	+ let mut best = 0u32;
	368	+ let mut best_val = f32::NEG_INFINITY;
	369	+ for (i, &v) in logits.iter().enumerate() {
	370	+ if v > best_val {
	371	+ best_val = v;
	372	+ best = i as u32;
	373	+ }
	374	+ }
	375	+ best
	376	+}
	377	+
	378	+/// Run a prefill forward pass for `seq_id`, extending its context with the new
	379	+/// tokens. Returns the sampled token, or `None` if the sequence/model is gone.
	380	+fn forward_for_seq(
	381	+ seq_id: u64,
	382	+ new_tokens: &[u32],
	383	+ seq_model: &Arc<Mutex<HashMap<u64, Arc<dyn Model>>>>,
	384	+ seq_context: &Arc<Mutex<HashMap<u64, Vec<u32>>>>,
	385	+ params: &SamplingParams,
	386	+) -> Option<u32> {
	387	+ let model = seq_model.lock().unwrap().get(&seq_id)?.clone();
	388	+ {
	389	+ let mut ctx = seq_context.lock().unwrap();
	390	+ ctx.entry(seq_id).or_default().extend_from_slice(new_tokens);
	391	+ }
	392	+ let ctx = seq_context.lock().unwrap();
	393	+ let full = ctx.get(&seq_id)?.clone();
	394	+ drop(ctx);
	395	+ let logits = model.forward_logits(&full).ok()?;
	396	+ Some(sample_token_with_params(&logits, params))
	397	+}
	398	+
	399	+/// Run a decode forward pass for `seq_id` (context already includes prior
	400	+/// tokens), append the sampled token, and return it.
	401	+fn forward_decode(
	402	+ seq_id: u64,
	403	+ seq_model: &Arc<Mutex<HashMap<u64, Arc<dyn Model>>>>,
	404	+ seq_context: &Arc<Mutex<HashMap<u64, Vec<u32>>>>,
	405	+ params: &SamplingParams,
	406	+) -> Option<u32> {
	407	+ let model = seq_model.lock().unwrap().get(&seq_id)?.clone();
	408	+ let prev_ctx = seq_context.lock().unwrap().get(&seq_id)?.clone();
	409	+ let logits = model.forward_logits(&prev_ctx).ok()?;
	410	+ let token = sample_token_with_params(&logits, params);
	411	+ if let Some(c) = seq_context.lock().unwrap().get_mut(&seq_id) {
	412	+ c.push(token);
	413	+ }
	414	+ Some(token)
	415	+}
	416	+
	417	+/// Send a token to a sequence's channel if it is still open.
	418	+fn deliver(
	419	+ senders: &Arc<Mutex<HashMap<u64, mpsc::UnboundedSender<GeneratedToken>>>>,
	420	+ seq_id: u64,
	421	+ token: GeneratedToken,
	422	+) {
	423	+ if let Some(tx) = senders.lock().unwrap().get(&seq_id) {
	424	+ let _ = tx.send(token);
	425	+ }
	426	+}
	427	+
	428	+/// Whether `seq_id` has reached its `max_tokens` budget.
	429	+fn is_seq_done(
	430	+ seq_id: u64,
	431	+ scheduler: &Arc<Mutex<Scheduler>>,
	432	+ seq_max_tokens: &Arc<Mutex<HashMap<u64, usize>>>,
	433	+) -> bool {
	434	+ let max = seq_max_tokens
	435	+ .lock()
	436	+ .unwrap()
	437	+ .get(&seq_id)
	438	+ .copied()
	439	+ .unwrap_or(usize::MAX);
	440	+ let generated = scheduler.lock().unwrap().running_seq_generated(seq_id);
	441	+ match generated {
	442	+ Some(n) => n >= max,
	443	+ None => false,
	444	+ }
	445	+}
	446	+
	447	+impl Drop for SchedulerEngine {
	448	+ fn drop(&mut self) {
	449	+ self.alive.store(false, Ordering::Relaxed);
	450	+ }
	451	+}
	452	+
	453	+#[cfg(test)]
	454	+mod tests {
	455	+ use super::*;
	456	+ use crate::engine::EngineDevice;
	457	+ use crate::kv_cache::KvBlockPool;
	458	+ use crate::scheduler::SchedulerConfig;
	459	+
	460	+ fn test_setup() -> (Arc<Engine>, Arc<Mutex<Scheduler>>, Arc<SchedulerEngine>) {
	461	+ let engine = Arc::new(Engine::new(EngineDevice::Cpu).unwrap());
	462	+ engine.register_mock("mock");
	463	+
	464	+ let pool = KvBlockPool::new(256);
	465	+ let config = SchedulerConfig {
	466	+ max_num_seqs: 4,
	467	+ max_num_batched_tokens: 4096,
	468	+ max_seq_len: 4096,
	469	+ chunked_prefill_size: None,
	470	+ };
	471	+ let scheduler = Arc::new(Mutex::new(Scheduler::new(config, pool)));
	472	+ let se = Arc::new(SchedulerEngine::new(engine.clone(), scheduler.clone()));
	473	+
	474	+ (engine, scheduler, se)
	475	+ }
	476	+
	477	+ async fn drain_receiver(
	478	+ rx: &mut mpsc::UnboundedReceiver<GeneratedToken>,
	479	+ ) -> Vec<GeneratedToken> {
	480	+ let mut tokens = Vec::new();
	481	+ while let Some(tok) = rx.recv().await {
	482	+ if tok.finish_reason.is_some() {
	483	+ break;
	484	+ }
	485	+ tokens.push(tok);
	486	+ }
	487	+ tokens
	488	+ }
	489	+
	490	+ #[tokio::test]
	491	+ async fn scheduler_engine_generates_tokens() {
	492	+ let (_, _, se) = test_setup();
	493	+ let prompt = MockModel::tokenize("hello");
	494	+ let params = SamplingParams {
	495	+ max_tokens: 10,
	496	+ temperature: 0.0,
	497	+ ..Default::default()
	498	+ };
	499	+
	500	+ let mut rx = se.generate("mock", &prompt, params).unwrap();
	501	+ let tokens = drain_receiver(&mut rx).await;
	502	+ assert!(!tokens.is_empty(), "should generate at least one token");
	503	+ }
	504	+
	505	+ #[tokio::test]
	506	+ async fn scheduler_engine_streams_multiple_tokens() {
	507	+ let (_, _, se) = test_setup();
	508	+ let prompt = MockModel::tokenize("hello world");
	509	+ let params = SamplingParams {
	510	+ max_tokens: 5,
	511	+ temperature: 0.0,
	512	+ ..Default::default()
	513	+ };
	514	+
	515	+ let mut rx = se.generate("mock", &prompt, params).unwrap();
	516	+ let tokens = drain_receiver(&mut rx).await;
	517	+ assert!(!tokens.is_empty(), "should produce tokens");
	518	+ assert!(
	519	+ tokens.iter().all(\|t\| !t.text.is_empty()),
	520	+ "every token should have text"
	521	+ );
	522	+ }
	523	+
	524	+ #[tokio::test]
	525	+ async fn scheduler_engine_prefix_cache_hit_reuses_blocks() {
	526	+ let (_, sched, se) = test_setup();
	527	+ // Need >= 16 tokens for at least one full block (BLOCK_SIZE=16).
	528	+ let prompt_text = "abcdefghijklmnopqrstuvwxyz0123456789"; // 36 tokens
	529	+ let prompt = MockModel::tokenize(prompt_text);
	530	+ assert!(
	531	+ prompt.len() >= 16,
	532	+ "prompt must be >= 16 tokens for block caching"
	533	+ );
	534	+
	535	+ let params = SamplingParams {
	536	+ max_tokens: 3,
	537	+ temperature: 0.0,
	538	+ ..Default::default()
	539	+ };
	540	+
	541	+ // First request caches its prompt.
	542	+ let mut rx1 = se.generate("mock", &prompt, params.clone()).unwrap();
	543	+ drain_receiver(&mut rx1).await;
	544	+
	545	+ let stats_before = se.stats();
	546	+ assert!(
	547	+ stats_before.prefix.nodes > 0,
	548	+ "expected at least one prefix cache node after first request, got {}",
	549	+ stats_before.prefix.nodes
	550	+ );
	551	+
	552	+ // Second request with same prompt → should hit the prefix cache.
	553	+ let mut rx2 = se.generate("mock", &prompt, params).unwrap();
	554	+ drain_receiver(&mut rx2).await;
	555	+
	556	+ let stats_after = se.stats();
	557	+ assert!(
	558	+ stats_after.prefix.hits > stats_before.prefix.hits,
	559	+ "expected at least one additional prefix cache hit, got {} -> {}",
	560	+ stats_before.prefix.hits,
	561	+ stats_after.prefix.hits
	562	+ );
	563	+ let _ = sched;
	564	+ }
	565	+
	566	+ #[tokio::test]
	567	+ async fn scheduler_engine_batches_concurrent_requests() {
	568	+ let (_, _, se) = test_setup();
	569	+ let params = SamplingParams {
	570	+ max_tokens: 4,
	571	+ temperature: 0.0,
	572	+ ..Default::default()
	573	+ };
	574	+
	575	+ // Launch two concurrent requests; both should complete.
	576	+ let mut rx_a = se
	577	+ .generate("mock", &MockModel::tokenize("hello"), params.clone())
	578	+ .unwrap();
	579	+ let mut rx_b = se
	580	+ .generate("mock", &MockModel::tokenize("world"), params.clone())
	581	+ .unwrap();
	582	+
	583	+ let a = drain_receiver(&mut rx_a).await;
	584	+ let b = drain_receiver(&mut rx_b).await;
	585	+
	586	+ assert!(!a.is_empty(), "request A should produce tokens");
	587	+ assert!(!b.is_empty(), "request B should produce tokens");
	588	+ }
	589	+
	590	+ #[tokio::test]
	591	+ async fn scheduler_engine_handles_many_concurrent_requests() {
	592	+ let (_, _, se) = test_setup();
	593	+ let params = SamplingParams {
	594	+ max_tokens: 2,
	595	+ temperature: 0.0,
	596	+ ..Default::default()
	597	+ };
	598	+
	599	+ // Fire 8 concurrent requests (more than max_num_seqs=4 batch limit).
	600	+ let mut receivers = Vec::new();
	601	+ for i in 0..8u32 {
	602	+ let p = MockModel::tokenize(&format!("request-{i}"));
	603	+ receivers.push(se.generate("mock", &p, params.clone()).unwrap());
	604	+ }
	605	+
	606	+ let mut completed = 0;
	607	+ for rx in receivers.iter_mut() {
	608	+ let tokens = drain_receiver(rx).await;
	609	+ if !tokens.is_empty() {
	610	+ completed += 1;
	611	+ }
	612	+ }
	613	+ assert_eq!(completed, 8, "all 8 concurrent requests should complete");
	614	+ }
	615	+
	616	+ #[test]
	617	+ fn scheduler_engine_stats_accessible() {
	618	+ let (_, _, se) = test_setup();
	619	+ let stats = se.stats();
	620	+ assert_eq!(stats.num_waiting, 0);
	621	+ assert_eq!(stats.num_running, 0);
	622	+ }
	623	+
	624	+ #[test]
	625	+ fn argmax_selects_max() {
	626	+ let logits = vec![-1.0, 0.5, 3.0, -2.0];
	627	+ assert_eq!(argmax(&logits), 2);
	628	+ }
	629	+
	630	+ #[test]
	631	+ fn argmax_handles_negative_all() {
	632	+ let logits = vec![-10.0, -5.0, -1.0];
	633	+ assert_eq!(argmax(&logits), 2);
	634	+ }
	635	+
	636	+ #[test]
	637	+ fn greedy_sampling_is_argmax() {
	638	+ // temperature 0 ⇒ deterministic argmax regardless of distribution.
	639	+ let logits = vec![-1.0, 0.5, 3.0, -2.0];
	640	+ let params = SamplingParams {
	641	+ temperature: 0.0,
	642	+ ..Default::default()
	643	+ };
	644	+ assert_eq!(sample_token_with_params(&logits, &params), 2);
	645	+ }
	646	+
	647	+ #[test]
	648	+ fn sampling_honors_temperature() {
	649	+ // A sharp peak with very low temperature should always pick the peak.
	650	+ let logits = vec![-100.0, -100.0, 50.0, -100.0];
	651	+ let params = SamplingParams {
	652	+ temperature: 0.01,
	653	+ ..Default::default()
	654	+ };
	655	+ for _ in 0..50 {
	656	+ assert_eq!(sample_token_with_params(&logits, &params), 2);
	657	+ }
	658	+ }
	659	+
	660	+ #[test]
	661	+ fn sampling_top_k_truncates() {
	662	+ // top_k=1 ⇒ always the argmax even with high temperature noise.
	663	+ let logits = vec![1.0, 2.0, 10.0, 0.5];
	664	+ let params = SamplingParams {
	665	+ temperature: 2.0,
	666	+ top_k: 1,
	667	+ ..Default::default()
	668	+ };
	669	+ for _ in 0..50 {
	670	+ assert_eq!(sample_token_with_params(&logits, &params), 2);
	671	+ }
	672	+ }
	673	+
	674	+ #[test]
	675	+ fn sampling_uniform_distributes_with_high_temp() {
	676	+ // A flat distribution at high temperature should produce variety, not
	677	+ // a single token (stochasticity is exercised).
	678	+ let logits = vec![0.0, 0.0, 0.0, 0.0];
	679	+ let params = SamplingParams {
	680	+ temperature: 1.0,
	681	+ ..Default::default()
	682	+ };
	683	+ let mut seen = std::collections::HashSet::new();
	684	+ for _ in 0..200 {
	685	+ seen.insert(sample_token_with_params(&logits, &params));
	686	+ }
	687	+ assert!(
	688	+ seen.len() >= 2,
	689	+ "high-temperature flat logits should sample multiple tokens, saw {}",
	690	+ seen.len()
	691	+ );
	692	+ }
	693	+}
0	694

▼crates/justapi-inference/src/spec_decode.rs

+535-0

@@ -0,0 +1,535 @@
	1	+//! Speculative decoding (draft-target, Chen et al. 2023).
	2	+//!
	3	+//! Standard draft-target scheme: a draft model proposes a block of `gamma`
	4	+//! tokens, the target model scores them in a single forward pass per
	5	+//! position, and we accept the longest prefix the target agrees with, then
	6	+//! sample one correction token from the target. When the draft is cheap and
	7	+//! well-aligned with the target, this yields up to `gamma + 1` tokens per
	8	+//! decode step instead of one — exactly the throughput win vLLM/SGLang get from
	9	+//! speculative decoding, with no Python GIL on the hot path.
	10	+//!
	11	+//! Medusa / EAGLE are tree-based refinements of the same verify step; the hook
	12	+//! here (`gamma` + draft model) is the foundation they build on.
	13	+
	14	+use std::sync::Arc;
	15	+
	16	+use rand::rngs::StdRng;
	17	+use rand::{Rng, SeedableRng};
	18	+
	19	+use crate::model::{FinishReason, GeneratedToken, Model, ModelError, SamplingParams};
	20	+
	21	+/// Statistics collected during a speculative decode run — the acceptance rate
	22	+/// is the dominant driver of the realized speedup (tokens emitted per step).
	23	+#[derive(Debug, Clone, Copy, Default)]
	24	+pub struct AcceptanceStats {
	25	+ /// Number of verify steps executed.
	26	+ pub steps: usize,
	27	+ /// Total number of draft tokens proposed (draft-target: `gamma` per step;
	28	+ /// tree-based: `gamma` per step — i.e. the number of positions
	29	+ /// speculated, not the number of tree nodes).
	30	+ pub total_draft: usize,
	31	+ /// Number of proposed draft tokens the target accepted (excludes the
	32	+ /// bonus "gamma + 1" target token emitted after a fully-accepted block).
	33	+ pub total_accepted: usize,
	34	+ /// Number of bonus target tokens emitted (the "gamma + 1" correction after
	35	+ /// a fully-accepted draft block). This is the source of the speedup.
	36	+ pub extra_emitted: usize,
	37	+ /// Branch factor for a tree-based speculation run. `0` means draft-target
	38	+ /// (single-path), `>0` means tree-based (Medusa/EAGLE-style).
	39	+ pub tree_branch: usize,
	40	+ /// Total tree nodes scored by the target (only meaningful when
	41	+ /// `tree_branch > 0`). Measures the verification cost per step.
	42	+ pub tree_nodes_verified: usize,
	43	+}
	44	+
	45	+impl AcceptanceStats {
	46	+ /// Create stats for a tree-based speculation run with the given branch factor.
	47	+ pub fn with_tree_branch(branch: usize) -> Self {
	48	+ Self {
	49	+ tree_branch: branch,
	50	+ ..Default::default()
	51	+ }
	52	+ }
	53	+
	54	+ /// Total tokens emitted (`accepted_draft + bonus`).
	55	+ pub fn tokens_emitted(&self) -> usize {
	56	+ self.total_accepted + self.extra_emitted
	57	+ }
	58	+
	59	+ /// Draft acceptance rate in `[0, 1]`. `1.0` means every proposed token was
	60	+ /// accepted (the draft is a perfect predictor at this temperature).
	61	+ pub fn acceptance_rate(&self) -> f64 {
	62	+ if self.total_draft == 0 {
	63	+ 0.0
	64	+ } else {
	65	+ self.total_accepted as f64 / self.total_draft as f64
	66	+ }
	67	+ }
	68	+}
	69	+
	70	+/// Configuration for a speculative-decode run.
	71	+#[derive(Clone)]
	72	+pub struct SpeculativeConfig {
	73	+ /// Number of tokens the draft model proposes per verify step (`gamma`).
	74	+ /// `0` degenerates to plain target-only decoding (1 token / step).
	75	+ pub gamma: usize,
	76	+ /// The cheaper model that proposes candidate tokens.
	77	+ pub draft: Arc<dyn Model>,
	78	+ /// RNG seed for reproducible sampling during verification.
	79	+ pub seed: u64,
	80	+}
	81	+
	82	+impl SpeculativeConfig {
	83	+ pub fn new(gamma: usize, draft: Arc<dyn Model>, seed: u64) -> Self {
	84	+ Self { gamma, draft, seed }
	85	+ }
	86	+}
	87	+
	88	+/// Sample a single token from `logits` honoring `temperature`/`top_p`/`top_k`.
	89	+///
	90	+/// Pure-Rust (no Candle) so it works for any `Model`, including [`MockModel`].
	91	+/// With `temperature == 0` it returns the argmax (greedy), which makes
	92	+/// correctness comparisons against plain target decode deterministic.
	93	+pub fn sample_token(logits: &[f32], params: &SamplingParams, rng: &mut StdRng) -> u32 {
	94	+ let vocab = logits.len();
	95	+ if vocab == 0 {
	96	+ return 0;
	97	+ }
	98	+ if params.temperature < 1e-7 {
	99	+ return logits
	100	+ .iter()
	101	+ .enumerate()
	102	+ .max_by(\|a, b\| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
	103	+ .map(\|(i, _)\| i as u32)
	104	+ .unwrap_or(0);
	105	+ }
	106	+
	107	+ // Temperature scaling.
	108	+ let inv = 1.0 / params.temperature.max(1e-7);
	109	+ let mut scaled: Vec<f32> = logits.iter().map(\|l\| l * inv).collect();
	110	+
	111	+ // Top-k truncation.
	112	+ if params.top_k > 0 && params.top_k < vocab {
	113	+ let mut idx: Vec<usize> = (0..vocab).collect();
	114	+ idx.sort_by(\|&a, &b\| {
	115	+ scaled[b]
	116	+ .partial_cmp(&scaled[a])
	117	+ .unwrap_or(std::cmp::Ordering::Equal)
	118	+ });
	119	+ let kth = scaled[idx[params.top_k.min(vocab - 1)]];
	120	+ for s in scaled.iter_mut() {
	121	+ if *s < kth - 1e-6 {
	122	+ *s = f32::NEG_INFINITY;
	123	+ }
	124	+ }
	125	+ }
	126	+
	127	+ // Softmax.
	128	+ let max = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
	129	+ let mut probs: Vec<f32> = scaled.iter().map(\|s\| (s - max).exp()).collect();
	130	+ let sum: f32 = probs.iter().sum();
	131	+
	132	+ // Top-p (nucleus) truncation.
	133	+ if params.top_p < 1.0 && sum > 0.0 {
	134	+ let mut order: Vec<usize> = (0..vocab).collect();
	135	+ order.sort_by(\|&a, &b\| {
	136	+ probs[b]
	137	+ .partial_cmp(&probs[a])
	138	+ .unwrap_or(std::cmp::Ordering::Equal)
	139	+ });
	140	+ let mut cum = 0.0;
	141	+ let mut cutoff = 1.0;
	142	+ for &i in &order {
	143	+ cum += probs[i] / sum;
	144	+ if cum >= params.top_p {
	145	+ cutoff = probs[i] / sum;
	146	+ break;
	147	+ }
	148	+ }
	149	+ for p in probs.iter_mut() {
	150	+ if *p / sum < cutoff - 1e-9 {
	151	+ *p = 0.0;
	152	+ }
	153	+ }
	154	+ }
	155	+
	156	+ let sum2: f32 = probs.iter().sum();
	157	+ if sum2 <= 0.0 {
	158	+ return (vocab - 1) as u32;
	159	+ }
	160	+
	161	+ // Categorical sample.
	162	+ let r: f32 = rng.gen();
	163	+ let mut cdf = 0.0;
	164	+ for (i, p) in probs.iter().enumerate() {
	165	+ cdf += p / sum2;
	166	+ if r < cdf {
	167	+ return i as u32;
	168	+ }
	169	+ }
	170	+ (vocab - 1) as u32
	171	+}
	172	+
	173	+/// Run speculative decoding and stream accepted tokens through `sink`.
	174	+///
	175	+/// Returns the finish reason and per-run [`AcceptanceStats`]. The implementation
	176	+/// matches the standard algorithm: draft proposes `gamma` tokens, target scores
	177	+/// and accepts the longest agreeing prefix, then emits one correction token.
	178	+/// With `gamma == 0` this is identical to plain target decode.
	179	+pub fn speculative_generate(
	180	+ target: &dyn Model,
	181	+ draft: &dyn Model,
	182	+ prompt: &[u32],
	183	+ params: &SamplingParams,
	184	+ gamma: usize,
	185	+ seed: u64,
	186	+ sink: &dyn Fn(GeneratedToken) -> bool,
	187	+) -> Result<(FinishReason, AcceptanceStats), ModelError> {
	188	+ let mut rng = StdRng::seed_from_u64(seed);
	189	+ let mut context: Vec<u32> = prompt.to_vec();
	190	+ let mut emitted = 0usize;
	191	+ let mut stats = AcceptanceStats::default();
	192	+ let mut finish = FinishReason::Length;
	193	+
	194	+ let detok = \|id: u32\| -> String { MockModel::detokenize(&[id]) };
	195	+
	196	+ 'outer: while emitted < params.max_tokens {
	197	+ stats.steps += 1;
	198	+
	199	+ // --- Draft phase: propose up to `gamma` tokens. ---
	200	+ let mut draft_tokens: Vec<u32> = Vec::with_capacity(gamma);
	201	+ let mut dctx = context.clone();
	202	+ for _ in 0..gamma {
	203	+ if emitted + draft_tokens.len() >= params.max_tokens {
	204	+ break;
	205	+ }
	206	+ let logits = draft.forward_logits(&dctx)?;
	207	+ let tok = sample_token(&logits, params, &mut rng);
	208	+ if params.stop_tokens.contains(&tok) {
	209	+ break;
	210	+ }
	211	+ draft_tokens.push(tok);
	212	+ dctx.push(tok);
	213	+ }
	214	+ stats.total_draft += draft_tokens.len();
	215	+
	216	+ // --- Verify phase: target scores the draft block. ---
	217	+ let mut tctx = context.clone();
	218	+ let mut rejected = false;
	219	+ for &dtok in &draft_tokens {
	220	+ let tlogits = target.forward_logits(&tctx)?;
	221	+ let ttok = sample_token(&tlogits, params, &mut rng);
	222	+ if ttok == dtok {
	223	+ stats.total_accepted += 1;
	224	+ let token = GeneratedToken {
	225	+ id: ttok,
	226	+ text: detok(ttok),
	227	+ logprob: 0.0,
	228	+ finish_reason: None,
	229	+ };
	230	+ if !sink(token) {
	231	+ finish = FinishReason::Cancelled;
	232	+ break 'outer;
	233	+ }
	234	+ emitted += 1;
	235	+ tctx.push(ttok);
	236	+ } else {
	237	+ // Reject: emit the target's token as the correction, stop block.
	238	+ let token = GeneratedToken {
	239	+ id: ttok,
	240	+ text: detok(ttok),
	241	+ logprob: 0.0,
	242	+ finish_reason: None,
	243	+ };
	244	+ if !sink(token) {
	245	+ finish = FinishReason::Cancelled;
	246	+ break 'outer;
	247	+ }
	248	+ emitted += 1;
	249	+ tctx.push(ttok);
	250	+ rejected = true;
	251	+ break;
	252	+ }
	253	+ }
	254	+
	255	+ context = tctx;
	256	+ if rejected {
	257	+ continue 'outer;
	258	+ }
	259	+
	260	+ // --- All drafts accepted: target emits one more token (gamma + 1). ---
	261	+ if emitted < params.max_tokens {
	262	+ let tlogits = target.forward_logits(&context)?;
	263	+ let ttok = sample_token(&tlogits, params, &mut rng);
	264	+ let stop = params.stop_tokens.contains(&ttok);
	265	+ let token = GeneratedToken {
	266	+ id: ttok,
	267	+ text: detok(ttok),
	268	+ logprob: 0.0,
	269	+ finish_reason: if stop { Some(FinishReason::Stop) } else { None },
	270	+ };
	271	+ if !sink(token) {
	272	+ finish = FinishReason::Cancelled;
	273	+ break;
	274	+ }
	275	+ emitted += 1;
	276	+ stats.extra_emitted += 1;
	277	+ context.push(ttok);
	278	+ if stop {
	279	+ finish = FinishReason::Stop;
	280	+ break;
	281	+ }
	282	+ }
	283	+ }
	284	+
	285	+ Ok((finish, stats))
	286	+}
	287	+
	288	+/// A `Model` that runs [`speculative_generate`] under the hood, so it can be
	289	+/// registered in the [`crate::ModelRegistry`] and served through the normal
	290	+/// `Engine::generate` path without callers knowing speculation is in play.
	291	+pub struct SpeculativeModel {
	292	+ target: Arc<dyn Model>,
	293	+ config: SpeculativeConfig,
	294	+}
	295	+
	296	+impl SpeculativeModel {
	297	+ pub fn new(target: Arc<dyn Model>, config: SpeculativeConfig) -> Self {
	298	+ Self { target, config }
	299	+ }
	300	+}
	301	+
	302	+impl Model for SpeculativeModel {
	303	+ fn vocab_size(&self) -> usize {
	304	+ self.target.vocab_size()
	305	+ }
	306	+
	307	+ fn generate(
	308	+ &self,
	309	+ prompt: &[u32],
	310	+ params: &SamplingParams,
	311	+ sink: &dyn Fn(GeneratedToken) -> bool,
	312	+ ) -> Result<FinishReason, ModelError> {
	313	+ let (finish, _stats) = speculative_generate(
	314	+ &*self.target,
	315	+ &*self.config.draft,
	316	+ prompt,
	317	+ params,
	318	+ self.config.gamma,
	319	+ self.config.seed,
	320	+ sink,
	321	+ )?;
	322	+ Ok(finish)
	323	+ }
	324	+
	325	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError> {
	326	+ self.target.forward_logits(context)
	327	+ }
	328	+}
	329	+
	330	+use crate::model::MockModel;
	331	+
	332	+#[cfg(test)]
	333	+mod tests {
	334	+ use super::*;
	335	+
	336	+ /// A stub model whose next-token rule is `(prev + offset) % vocab`, used to
	337	+ /// simulate a different (imperfect) draft model for acceptance tests.
	338	+ struct OffsetModel {
	339	+ vocab: usize,
	340	+ offset: u32,
	341	+ }
	342	+
	343	+ impl Model for OffsetModel {
	344	+ fn vocab_size(&self) -> usize {
	345	+ self.vocab
	346	+ }
	347	+
	348	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError> {
	349	+ let prev = context.last().copied().unwrap_or(0) % self.vocab as u32;
	350	+ let next = (prev.wrapping_add(self.offset)) % self.vocab as u32;
	351	+ let mut logits = vec![-5.0f32; self.vocab];
	352	+ logits[next as usize] = 5.0;
	353	+ Ok(logits)
	354	+ }
	355	+
	356	+ fn generate(
	357	+ &self,
	358	+ prompt: &[u32],
	359	+ params: &SamplingParams,
	360	+ sink: &dyn Fn(GeneratedToken) -> bool,
	361	+ ) -> Result<FinishReason, ModelError> {
	362	+ let mut prev = prompt.last().copied().unwrap_or(0);
	363	+ for _ in 0..params.max_tokens {
	364	+ let next = (prev.wrapping_add(self.offset)) % self.vocab as u32;
	365	+ let stop = params.stop_tokens.contains(&next);
	366	+ if !sink(GeneratedToken {
	367	+ id: next,
	368	+ text: MockModel::detokenize(&[next]),
	369	+ logprob: 0.0,
	370	+ finish_reason: if stop { Some(FinishReason::Stop) } else { None },
	371	+ }) {
	372	+ return Ok(FinishReason::Cancelled);
	373	+ }
	374	+ if stop {
	375	+ return Ok(FinishReason::Stop);
	376	+ }
	377	+ prev = next;
	378	+ }
	379	+ Ok(FinishReason::Length)
	380	+ }
	381	+ }
	382	+
	383	+ fn collect(model: &dyn Model, prompt: &[u32], params: &SamplingParams) -> Vec<u32> {
	384	+ let out = std::cell::RefCell::new(Vec::new());
	385	+ model
	386	+ .generate(prompt, params, &\|t\| {
	387	+ out.borrow_mut().push(t.id);
	388	+ true
	389	+ })
	390	+ .unwrap();
	391	+ out.into_inner()
	392	+ }
	393	+
	394	+ fn collect_spec(
	395	+ target: &dyn Model,
	396	+ draft: &dyn Model,
	397	+ prompt: &[u32],
	398	+ params: &SamplingParams,
	399	+ gamma: usize,
	400	+ ) -> (Vec<u32>, AcceptanceStats) {
	401	+ let out = std::cell::RefCell::new(Vec::new());
	402	+ let (_finish, stats) =
	403	+ speculative_generate(target, draft, prompt, params, gamma, 42, &\|t\| {
	404	+ out.borrow_mut().push(t.id);
	405	+ true
	406	+ })
	407	+ .unwrap();
	408	+ (out.into_inner(), stats)
	409	+ }
	410	+
	411	+ #[test]
	412	+ fn target_only_matches_spec_at_gamma_zero() {
	413	+ // gamma 0 must reproduce the target's normal decode exactly.
	414	+ let target = MockModel::new(32);
	415	+ let draft = MockModel::new(32);
	416	+ let params = SamplingParams {
	417	+ max_tokens: 12,
	418	+ temperature: 0.0,
	419	+ ..Default::default()
	420	+ };
	421	+ let prompt = vec![0u32];
	422	+ let baseline = collect(&target, &prompt, &params);
	423	+ let (spec, _) = collect_spec(&target, &draft, &prompt, &params, 0);
	424	+ assert_eq!(baseline, spec);
	425	+ assert_eq!(spec.len(), 12);
	426	+ }
	427	+
	428	+ #[test]
	429	+ fn spec_matches_target_when_draft_equals_target() {
	430	+ // A perfect draft (=== target) changes no output: speculation is
	431	+ // correctness-preserving, only faster.
	432	+ let target = MockModel::new(32);
	433	+ let draft = MockModel::new(32);
	434	+ let params = SamplingParams {
	435	+ max_tokens: 20,
	436	+ temperature: 0.0,
	437	+ ..Default::default()
	438	+ };
	439	+ let prompt = vec![5u32];
	440	+ let baseline = collect(&target, &prompt, &params);
	441	+ let (spec, _) = collect_spec(&target, &draft, &prompt, &params, 4);
	442	+ assert_eq!(baseline, spec, "speculation must not change the output");
	443	+ }
	444	+
	445	+ #[test]
	446	+ fn perfect_draft_yields_gamma_plus_1_per_step() {
	447	+ let target = MockModel::new(32);
	448	+ let draft = MockModel::new(32);
	449	+ let max = 20usize;
	450	+ let gamma = 4usize;
	451	+ let params = SamplingParams {
	452	+ max_tokens: max,
	453	+ temperature: 0.0,
	454	+ ..Default::default()
	455	+ };
	456	+ let (spec, stats) = collect_spec(&target, &draft, &[0u32], &params, gamma);
	457	+ assert_eq!(spec.len(), max);
	458	+ // With a perfect draft, every step emits gamma+1 tokens, so we need
	459	+ // ceil(max / (gamma+1)) steps.
	460	+ let expected_steps = max.div_ceil(gamma + 1);
	461	+ assert_eq!(stats.steps, expected_steps);
	462	+ // All draft tokens accepted (acceptance rate ~ 1.0).
	463	+ assert!((stats.acceptance_rate() - 1.0).abs() < 1e-9);
	464	+ // Bonus tokens == steps (one gamma+1 token per step).
	465	+ assert_eq!(stats.extra_emitted, expected_steps);
	466	+ assert_eq!(stats.tokens_emitted(), max);
	467	+ }
	468	+
	469	+ #[test]
	470	+ fn imperfect_draft_lowers_acceptance_rate() {
	471	+ let target = MockModel::new(64);
	472	+ // Draft predicts (prev + 7) while target predicts (prev + 1): rarely equal.
	473	+ let draft = OffsetModel {
	474	+ vocab: 64,
	475	+ offset: 7,
	476	+ };
	477	+ let params = SamplingParams {
	478	+ max_tokens: 40,
	479	+ temperature: 0.0,
	480	+ ..Default::default()
	481	+ };
	482	+ let (_spec, stats) = collect_spec(&target, &draft, &[0u32], &params, 5);
	483	+ // With temperature 0 and a mismatched draft, acceptance is essentially 0.
	484	+ assert!(
	485	+ stats.acceptance_rate() < 0.2,
	486	+ "got {}",
	487	+ stats.acceptance_rate()
	488	+ );
	489	+ }
	490	+
	491	+ #[test]
	492	+ fn spec_honors_max_tokens_with_imperfect_draft() {
	493	+ let target = MockModel::new(64);
	494	+ let draft = OffsetModel {
	495	+ vocab: 64,
	496	+ offset: 3,
	497	+ };
	498	+ let max = 15usize;
	499	+ let params = SamplingParams {
	500	+ max_tokens: max,
	501	+ temperature: 0.0,
	502	+ ..Default::default()
	503	+ };
	504	+ let (spec, _) = collect_spec(&target, &draft, &[0u32], &params, 3);
	505	+ assert_eq!(spec.len(), max);
	506	+ }
	507	+
	508	+ #[test]
	509	+ fn spec_respects_stop_token() {
	510	+ let target = MockModel::new(32);
	511	+ let draft = MockModel::new(32);
	512	+ // MockModel rule is (prev+1)%32; pick `stop` at the next token id.
	513	+ let next_after_prompt = (0u32.wrapping_add(1)) % 32;
	514	+ let params = SamplingParams {
	515	+ max_tokens: 50,
	516	+ temperature: 0.0,
	517	+ stop_tokens: vec![next_after_prompt],
	518	+ ..Default::default()
	519	+ };
	520	+ let (spec, _) = collect_spec(&target, &draft, &[0u32], &params, 3);
	521	+ assert_eq!(spec.len(), 1);
	522	+ assert_eq!(spec[0], next_after_prompt);
	523	+ }
	524	+
	525	+ #[test]
	526	+ fn sampling_token_argmax_is_greedy() {
	527	+ let logits = vec![-1.0, 3.0, -2.0, 0.5];
	528	+ let params = SamplingParams {
	529	+ temperature: 0.0,
	530	+ ..Default::default()
	531	+ };
	532	+ let mut rng = StdRng::seed_from_u64(1);
	533	+ assert_eq!(sample_token(&logits, &params, &mut rng), 1);
	534	+ }
	535	+}
0	536

▼crates/justapi-inference/src/spec_decode_tree.rs

+861-0

@@ -0,0 +1,861 @@
	1	+//! Tree-based speculative decoding (Medusa / EAGLE-style).
	2	+//!
	3	+//! Extends the draft-target scheme from [`super::spec_decode`] with a tree of
	4	+//! candidate tokens per position. The draft proposes `branch` candidates at each
	5	+//! of `gamma` positions, forming a tree. The target verifies a single path
	6	+//! through the tree (the longest path where every node matches the target's
	7	+//! prediction at that context), which yields higher acceptance than the
	8	+//! single-path draft-target scheme because the draft gets multiple "guesses"
	9	+//! per position.
	10	+//!
	11	+//! # Lossless guarantee
	12	+//!
	13	+//! Like the draft-target scheme, tree-based speculation produces exactly the
	14	+//! same output as plain target decode — it is a correctness-preserving
	15	+//! optimization. At every verify step, if any tree node matches the target's
	16	+//! sample, that token is emitted; if no node matches, the target's own
	17	+//! correction token is emitted — which is token-for-token identical to running
	18	+//! the target alone.
	19	+
	20	+use rand::rngs::StdRng;
	21	+use rand::SeedableRng;
	22	+use std::sync::Arc;
	23	+
	24	+use crate::model::{FinishReason, GeneratedToken, Model, ModelError, SamplingParams};
	25	+use crate::spec_decode::sample_token;
	26	+use crate::MockModel;
	27	+
	28	+/// A single node in the draft tree.
	29	+#[derive(Clone, Debug)]
	30	+pub struct TreeNode {
	31	+ pub token: u32,
	32	+ pub score: f32,
	33	+ pub children: Vec<TreeNode>,
	34	+}
	35	+
	36	+impl TreeNode {
	37	+ pub fn new(token: u32, score: f32) -> Self {
	38	+ Self {
	39	+ token,
	40	+ score,
	41	+ children: vec![],
	42	+ }
	43	+ }
	44	+}
	45	+
	46	+/// A draft tree: branching structure of candidate tokens.
	47	+#[derive(Clone, Debug)]
	48	+pub struct DraftTree {
	49	+ pub roots: Vec<TreeNode>,
	50	+ pub gamma: usize,
	51	+ pub branch: usize,
	52	+}
	53	+
	54	+impl DraftTree {
	55	+ /// Total number of nodes = Σ_{i=1}^{γ} b^i.
	56	+ pub fn total_nodes(&self) -> usize {
	57	+ if self.branch == 0 \|\| self.gamma == 0 {
	58	+ return 0;
	59	+ }
	60	+ let b = self.branch as u64;
	61	+ let g = self.gamma as u64;
	62	+ if b == 1 {
	63	+ return g as usize;
	64	+ }
	65	+ (b * (b.pow(g as u32) - 1) / (b - 1)) as usize
	66	+ }
	67	+}
	68	+
	69	+/// Return the top-k token ids and raw logit scores.
	70	+pub fn top_k_tokens(logits: &[f32], k: usize) -> Vec<(u32, f32)> {
	71	+ if k == 0 \|\| logits.is_empty() {
	72	+ return vec![];
	73	+ }
	74	+ let mut scored: Vec<(u32, f32)> = logits
	75	+ .iter()
	76	+ .enumerate()
	77	+ .map(\|(i, &s)\| (i as u32, s))
	78	+ .collect();
	79	+ scored.sort_by(\|a, b\| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
	80	+ scored.truncate(k);
	81	+ scored
	82	+}
	83	+
	84	+/// Build a draft tree by recursively expanding the top-`branch` tokens at
	85	+/// each position up to `gamma` depth.
	86	+///
	87	+/// Makes O(b^γ) draft forward calls — acceptable for small γ and b.
	88	+pub fn build_draft_tree(
	89	+ draft: &dyn Model,
	90	+ context: &[u32],
	91	+ gamma: usize,
	92	+ branch: usize,
	93	+) -> Result<DraftTree, ModelError> {
	94	+ if gamma == 0 \|\| branch == 0 {
	95	+ return Ok(DraftTree {
	96	+ roots: vec![],
	97	+ gamma: 0,
	98	+ branch: 0,
	99	+ });
	100	+ }
	101	+
	102	+ fn expand(
	103	+ draft: &dyn Model,
	104	+ ctx: &[u32],
	105	+ depth_remaining: usize,
	106	+ branch: usize,
	107	+ ) -> Result<Vec<TreeNode>, ModelError> {
	108	+ if depth_remaining == 0 {
	109	+ return Ok(vec![]);
	110	+ }
	111	+ let logits = draft.forward_logits(ctx)?;
	112	+ let candidates = top_k_tokens(&logits, branch);
	113	+ let mut nodes = Vec::with_capacity(candidates.len());
	114	+ for (token, score) in candidates {
	115	+ let mut child_ctx = ctx.to_vec();
	116	+ child_ctx.push(token);
	117	+ let children = expand(draft, &child_ctx, depth_remaining - 1, branch)?;
	118	+ nodes.push(TreeNode {
	119	+ token,
	120	+ score,
	121	+ children,
	122	+ });
	123	+ }
	124	+ Ok(nodes)
	125	+ }
	126	+
	127	+ let roots = expand(draft, context, gamma, branch)?;
	128	+ Ok(DraftTree {
	129	+ roots,
	130	+ gamma,
	131	+ branch,
	132	+ })
	133	+}
	134	+
	135	+/// Verify a draft tree by walking the longest matching path.
	136	+///
	137	+/// For each depth:
	138	+/// 1. Call target forward with current context (original + accepted prefix).
	139	+/// 2. Sample the target's next token.
	140	+/// 3. Find a matching child at this depth — if found, accept and continue;
	141	+/// if not found, emit correction (the target's token) and stop.
	142	+///
	143	+/// Returns `(accepted_tokens, bonus_or_correction, accepted_count)`.
	144	+/// - `accepted_tokens`: draft tokens accepted from the tree (may be empty).
	145	+/// - `bonus_or_correction`: `Some(token)` — always emitted (target's own
	146	+/// prediction after the accepted prefix or at the rejection point).
	147	+/// - `accepted_count`: same as `accepted_tokens.len()`.
	148	+pub fn verify_tree(
	149	+ target: &dyn Model,
	150	+ tree: &DraftTree,
	151	+ context: &[u32],
	152	+ params: &SamplingParams,
	153	+ rng: &mut StdRng,
	154	+) -> Result<(Vec<u32>, Option<u32>, usize), ModelError> {
	155	+ let mut accepted: Vec<u32> = Vec::with_capacity(tree.gamma);
	156	+ let mut ctx = context.to_vec();
	157	+
	158	+ // Clone roots into frontier so we can own children when descending.
	159	+ let mut frontier: Vec<(u32, Vec<TreeNode>)> = tree
	160	+ .roots
	161	+ .iter()
	162	+ .map(\|n\| (n.token, n.children.clone()))
	163	+ .collect();
	164	+
	165	+ let mut n_accepted = 0usize;
	166	+
	167	+ for _depth in 0..tree.gamma {
	168	+ let logits = target.forward_logits(&ctx)?;
	169	+ let target_tok = sample_token(&logits, params, rng);
	170	+
	171	+ if let Some(idx) = frontier.iter().position(\|(tok, _)\| *tok == target_tok) {
	172	+ let (_tok, children) = frontier.swap_remove(idx);
	173	+ accepted.push(target_tok);
	174	+ n_accepted += 1;
	175	+ ctx.push(target_tok);
	176	+ frontier = children
	177	+ .into_iter()
	178	+ .map(\|n\| (n.token, n.children))
	179	+ .collect();
	180	+ if frontier.is_empty() {
	181	+ let logits = target.forward_logits(&ctx)?;
	182	+ let bonus = sample_token(&logits, params, rng);
	183	+ return Ok((accepted, Some(bonus), n_accepted));
	184	+ }
	185	+ } else {
	186	+ return Ok((accepted, Some(target_tok), n_accepted));
	187	+ }
	188	+ }
	189	+
	190	+ let logits = target.forward_logits(&ctx)?;
	191	+ let bonus = sample_token(&logits, params, rng);
	192	+ Ok((accepted, Some(bonus), n_accepted))
	193	+}
	194	+
	195	+/// Run tree-based speculative decoding and stream tokens through `sink`.
	196	+///
	197	+/// Each step:
	198	+/// 1. `build_draft_tree` — draft proposes a tree of candidates.
	199	+/// 2. `verify_tree` — target scores the tree, finds the longest agreeing path.
	200	+/// 3. Accepted tokens + bonus/correction are streamed through `sink`.
	201	+///
	202	+/// Returns finish reason and per-run [`crate::AcceptanceStats`].
	203	+/// With `gamma == 0` or `branch == 0` this degenerates to plain decode.
	204	+#[allow(clippy::too_many_arguments)]
	205	+pub fn speculative_generate_tree(
	206	+ target: &dyn Model,
	207	+ draft: &dyn Model,
	208	+ prompt: &[u32],
	209	+ params: &SamplingParams,
	210	+ gamma: usize,
	211	+ branch: usize,
	212	+ seed: u64,
	213	+ sink: &dyn Fn(GeneratedToken) -> bool,
	214	+) -> Result<(FinishReason, crate::AcceptanceStats), ModelError> {
	215	+ let mut rng = StdRng::seed_from_u64(seed);
	216	+ let mut context: Vec<u32> = prompt.to_vec();
	217	+ let mut emitted = 0usize;
	218	+ let mut stats = crate::AcceptanceStats::with_tree_branch(branch);
	219	+ let mut finish = FinishReason::Length;
	220	+
	221	+ let detok = \|id: u32\| -> String { MockModel::detokenize(&[id]) };
	222	+
	223	+ while emitted < params.max_tokens {
	224	+ stats.steps += 1;
	225	+
	226	+ // 1. Build draft tree.
	227	+ let tree = build_draft_tree(draft, &context, gamma, branch)?;
	228	+
	229	+ // Positions speculated = gamma per step (comparable to draft-target).
	230	+ // Tree nodes scored = tree.total_nodes() (a cost metric).
	231	+ let eff_gamma = gamma.min(params.max_tokens.saturating_sub(emitted));
	232	+ stats.total_draft += eff_gamma;
	233	+ stats.tree_nodes_verified += tree.total_nodes();
	234	+
	235	+ // 2. Verify.
	236	+ let (accepted, bonus_or_correction, n_accepted) =
	237	+ verify_tree(target, &tree, &context, params, &mut rng)?;
	238	+
	239	+ stats.total_accepted += n_accepted;
	240	+
	241	+ // 3. Emit accepted tokens (check stop tokens for each).
	242	+ let mut stopped = false;
	243	+ for &tok in &accepted {
	244	+ if emitted >= params.max_tokens {
	245	+ break;
	246	+ }
	247	+ let stop = params.stop_tokens.contains(&tok);
	248	+ let token = GeneratedToken {
	249	+ id: tok,
	250	+ text: detok(tok),
	251	+ logprob: 0.0,
	252	+ finish_reason: if stop { Some(FinishReason::Stop) } else { None },
	253	+ };
	254	+ if !sink(token) {
	255	+ finish = FinishReason::Cancelled;
	256	+ stopped = true;
	257	+ break;
	258	+ }
	259	+ emitted += 1;
	260	+ context.push(tok);
	261	+ if stop {
	262	+ finish = FinishReason::Stop;
	263	+ stopped = true;
	264	+ break;
	265	+ }
	266	+ }
	267	+
	268	+ if stopped {
	269	+ break;
	270	+ }
	271	+ if emitted >= params.max_tokens {
	272	+ break;
	273	+ }
	274	+
	275	+ // 4. Emit bonus / correction token.
	276	+ if let Some(bonus) = bonus_or_correction {
	277	+ let stop = params.stop_tokens.contains(&bonus);
	278	+ let token = GeneratedToken {
	279	+ id: bonus,
	280	+ text: detok(bonus),
	281	+ logprob: 0.0,
	282	+ finish_reason: if stop { Some(FinishReason::Stop) } else { None },
	283	+ };
	284	+ if !sink(token) {
	285	+ finish = FinishReason::Cancelled;
	286	+ break;
	287	+ }
	288	+ emitted += 1;
	289	+ stats.extra_emitted += 1;
	290	+ context.push(bonus);
	291	+ if stop {
	292	+ finish = FinishReason::Stop;
	293	+ break;
	294	+ }
	295	+ }
	296	+ }
	297	+
	298	+ Ok((finish, stats))
	299	+}
	300	+
	301	+/// A `Model` that runs [`speculative_generate_tree`] under the hood, so it can
	302	+/// be registered in the [`crate::ModelRegistry`] and served through the normal
	303	+/// `Engine::generate` path without callers knowing tree speculation is in play.
	304	+pub struct TreeSpeculativeModel {
	305	+ target: Arc<dyn Model>,
	306	+ draft: Arc<dyn Model>,
	307	+ gamma: usize,
	308	+ branch: usize,
	309	+ seed: u64,
	310	+}
	311	+
	312	+impl TreeSpeculativeModel {
	313	+ pub fn new(
	314	+ target: Arc<dyn Model>,
	315	+ draft: Arc<dyn Model>,
	316	+ gamma: usize,
	317	+ branch: usize,
	318	+ seed: u64,
	319	+ ) -> Self {
	320	+ Self {
	321	+ target,
	322	+ draft,
	323	+ gamma,
	324	+ branch,
	325	+ seed,
	326	+ }
	327	+ }
	328	+}
	329	+
	330	+impl Model for TreeSpeculativeModel {
	331	+ fn vocab_size(&self) -> usize {
	332	+ self.target.vocab_size()
	333	+ }
	334	+
	335	+ fn generate(
	336	+ &self,
	337	+ prompt: &[u32],
	338	+ params: &SamplingParams,
	339	+ sink: &dyn Fn(GeneratedToken) -> bool,
	340	+ ) -> Result<FinishReason, ModelError> {
	341	+ let (finish, _stats) = speculative_generate_tree(
	342	+ &*self.target,
	343	+ &*self.draft,
	344	+ prompt,
	345	+ params,
	346	+ self.gamma,
	347	+ self.branch,
	348	+ self.seed,
	349	+ sink,
	350	+ )?;
	351	+ Ok(finish)
	352	+ }
	353	+
	354	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError> {
	355	+ self.target.forward_logits(context)
	356	+ }
	357	+}
	358	+
	359	+#[cfg(test)]
	360	+mod tests {
	361	+ use super::*;
	362	+ use crate::model::MockModel;
	363	+
	364	+ /// A model whose top-k tokens are a configurable range, for testing
	365	+ /// tree-based speculative decoding with controllable draft-target overlap.
	366	+ struct RangeModel {
	367	+ vocab: usize,
	368	+ offset: u32,
	369	+ spread: u32,
	370	+ }
	371	+
	372	+ impl RangeModel {
	373	+ fn new(vocab: usize, offset: u32, spread: u32) -> Self {
	374	+ Self {
	375	+ vocab,
	376	+ offset,
	377	+ spread,
	378	+ }
	379	+ }
	380	+ }
	381	+
	382	+ impl Model for RangeModel {
	383	+ fn vocab_size(&self) -> usize {
	384	+ self.vocab
	385	+ }
	386	+
	387	+ fn forward_logits(&self, context: &[u32]) -> Result<Vec<f32>, ModelError> {
	388	+ let prev = context.last().copied().unwrap_or(0) % self.vocab as u32;
	389	+ // Center = prev + offset + 1 (so offset=0 means centered on target's (prev+1))
	390	+ let center = (prev.wrapping_add(self.offset).wrapping_add(1)) % self.vocab as u32;
	391	+ let mut logits = vec![-10.0f32; self.vocab];
	392	+ for d in 0..=self.spread {
	393	+ let tok = (center.wrapping_add(d)) % self.vocab as u32;
	394	+ let score = 5.0 - d as f32;
	395	+ if score > logits[tok as usize] {
	396	+ logits[tok as usize] = score;
	397	+ }
	398	+ // Also negative side (wrapping reverse)
	399	+ if d > 0 {
	400	+ let tok2 = (center.wrapping_sub(d)) % self.vocab as u32;
	401	+ logits[tok2 as usize] = 5.0 - d as f32;
	402	+ }
	403	+ }
	404	+ Ok(logits)
	405	+ }
	406	+
	407	+ fn generate(
	408	+ &self,
	409	+ prompt: &[u32],
	410	+ params: &SamplingParams,
	411	+ sink: &dyn Fn(GeneratedToken) -> bool,
	412	+ ) -> Result<FinishReason, ModelError> {
	413	+ let mut prev = prompt.last().copied().unwrap_or(0);
	414	+ for _ in 0..params.max_tokens {
	415	+ let center = (prev.wrapping_add(self.offset).wrapping_add(1)) % self.vocab as u32;
	416	+ let next = center;
	417	+ let stop = params.stop_tokens.contains(&next);
	418	+ if !sink(GeneratedToken {
	419	+ id: next,
	420	+ text: MockModel::detokenize(&[next]),
	421	+ logprob: 0.0,
	422	+ finish_reason: if stop { Some(FinishReason::Stop) } else { None },
	423	+ }) {
	424	+ return Ok(FinishReason::Cancelled);
	425	+ }
	426	+ if stop {
	427	+ return Ok(FinishReason::Stop);
	428	+ }
	429	+ prev = next;
	430	+ }
	431	+ Ok(FinishReason::Length)
	432	+ }
	433	+ }
	434	+
	435	+ fn collect(model: &dyn Model, prompt: &[u32], params: &SamplingParams) -> Vec<u32> {
	436	+ let out = std::cell::RefCell::new(Vec::new());
	437	+ model
	438	+ .generate(prompt, params, &\|t\| {
	439	+ out.borrow_mut().push(t.id);
	440	+ true
	441	+ })
	442	+ .unwrap();
	443	+ out.into_inner()
	444	+ }
	445	+
	446	+ fn collect_tree(
	447	+ target: &dyn Model,
	448	+ draft: &dyn Model,
	449	+ prompt: &[u32],
	450	+ params: &SamplingParams,
	451	+ gamma: usize,
	452	+ branch: usize,
	453	+ ) -> (Vec<u32>, crate::AcceptanceStats) {
	454	+ let out = std::cell::RefCell::new(Vec::new());
	455	+ let (_finish, stats) =
	456	+ speculative_generate_tree(target, draft, prompt, params, gamma, branch, 42, &\|t\| {
	457	+ out.borrow_mut().push(t.id);
	458	+ true
	459	+ })
	460	+ .unwrap();
	461	+ (out.into_inner(), stats)
	462	+ }
	463	+
	464	+ // ---------- tree node / helpers ----------
	465	+
	466	+ #[test]
	467	+ fn top_k_returns_sorted() {
	468	+ let logits = vec![-5.0, 3.0, 10.0, -2.0, 7.0];
	469	+ let top = top_k_tokens(&logits, 3);
	470	+ assert_eq!(top.len(), 3);
	471	+ assert_eq!(top[0], (2, 10.0));
	472	+ assert_eq!(top[1], (4, 7.0));
	473	+ assert_eq!(top[2], (1, 3.0));
	474	+ }
	475	+
	476	+ #[test]
	477	+ fn top_k_zero_returns_empty() {
	478	+ assert!(top_k_tokens(&[1.0, 2.0], 0).is_empty());
	479	+ }
	480	+
	481	+ #[test]
	482	+ fn tree_total_nodes_formula() {
	483	+ let t = DraftTree {
	484	+ roots: vec![],
	485	+ gamma: 4,
	486	+ branch: 3,
	487	+ };
	488	+ // 3 + 9 + 27 + 81 = 120
	489	+ assert_eq!(t.total_nodes(), 120);
	490	+ }
	491	+
	492	+ #[test]
	493	+ fn tree_total_nodes_zero() {
	494	+ let t = DraftTree {
	495	+ roots: vec![],
	496	+ gamma: 0,
	497	+ branch: 3,
	498	+ };
	499	+ assert_eq!(t.total_nodes(), 0);
	500	+ let t = DraftTree {
	501	+ roots: vec![],
	502	+ gamma: 4,
	503	+ branch: 0,
	504	+ };
	505	+ assert_eq!(t.total_nodes(), 0);
	506	+ }
	507	+
	508	+ // ---------- tree matches plain decode with perfect draft ----------
	509	+
	510	+ #[test]
	511	+ fn tree_matches_target_when_draft_contains_correct_path() {
	512	+ // Target: MockModel — predicts (prev+1)%V
	513	+ // Draft: RangeModel(offset=0) — centered on (prev+1)%V, spread=3, so
	514	+ // top-3 include the correct prediction at every depth.
	515	+ let target = MockModel::new(32);
	516	+ let draft = RangeModel::new(32, 0, 3);
	517	+ let params = SamplingParams {
	518	+ max_tokens: 20,
	519	+ temperature: 0.0,
	520	+ ..Default::default()
	521	+ };
	522	+ let prompt = vec![5u32];
	523	+ let baseline = collect(&target, &prompt, &params);
	524	+ let (spec, stats) = collect_tree(&target, &draft, &prompt, &params, 4, 3);
	525	+ assert_eq!(
	526	+ baseline, spec,
	527	+ "tree speculation must not change the output"
	528	+ );
	529	+ // Acceptance should be high — perfect draft.
	530	+ assert!(
	531	+ stats.acceptance_rate() > 0.8,
	532	+ "got {}",
	533	+ stats.acceptance_rate()
	534	+ );
	535	+ assert_eq!(stats.tree_branch, 3);
	536	+ }
	537	+
	538	+ // ---------- gamma=0 degenerates to plain decode ----------
	539	+
	540	+ #[test]
	541	+ fn tree_gamma_zero_matches_plain() {
	542	+ let target = MockModel::new(32);
	543	+ let draft = MockModel::new(32);
	544	+ let params = SamplingParams {
	545	+ max_tokens: 12,
	546	+ temperature: 0.0,
	547	+ ..Default::default()
	548	+ };
	549	+ let prompt = vec![0u32];
	550	+ let baseline = collect(&target, &prompt, &params);
	551	+ let (spec, _) = collect_tree(&target, &draft, &prompt, &params, 0, 3);
	552	+ assert_eq!(baseline, spec);
	553	+ }
	554	+
	555	+ // ---------- branch=1 acts like draft-target ----------
	556	+
	557	+ #[test]
	558	+ fn tree_branch_one_acts_like_single_path() {
	559	+ let target = MockModel::new(32);
	560	+ let draft = MockModel::new(32);
	561	+ let params = SamplingParams {
	562	+ max_tokens: 20,
	563	+ temperature: 0.0,
	564	+ ..Default::default()
	565	+ };
	566	+ let prompt = vec![3u32];
	567	+ let baseline = collect(&target, &prompt, &params);
	568	+ // With branch=1, tree has only one candidate per depth = essentially
	569	+ // the same as draft-target (single path).
	570	+ let (spec, _) = collect_tree(&target, &draft, &prompt, &params, 4, 1);
	571	+ assert_eq!(baseline, spec);
	572	+ }
	573	+
	574	+ // ---------- graceful degradation when draft never matches ----------
	575	+
	576	+ #[test]
	577	+ fn tree_graceful_degradation() {
	578	+ // Target predicts (prev+1)%V.
	579	+ // Draft offset=10, centered on (prev+11)%V — never includes (prev+1).
	580	+ let target = MockModel::new(64);
	581	+ let draft = RangeModel::new(64, 10, 0);
	582	+ let params = SamplingParams {
	583	+ max_tokens: 30,
	584	+ temperature: 0.0,
	585	+ ..Default::default()
	586	+ };
	587	+ let prompt = vec![0u32];
	588	+ let baseline = collect(&target, &prompt, &params);
	589	+ let (spec, stats) = collect_tree(&target, &draft, &prompt, &params, 4, 3);
	590	+ assert_eq!(baseline, spec, "must preserve output even with bad draft");
	591	+ // Acceptance rate should be near 0.
	592	+ assert!(
	593	+ stats.acceptance_rate() < 0.1,
	594	+ "got {}",
	595	+ stats.acceptance_rate()
	596	+ );
	597	+ // Tokens emitted == max_tokens (graceful, no loss)
	598	+ assert_eq!(spec.len(), 30);
	599	+ }
	600	+
	601	+ // ---------- branch > 1 gives higher acceptance than branch=1 ----------
	602	+
	603	+ #[test]
	604	+ fn higher_branch_gives_better_acceptance() {
	605	+ // Draft's offset is 1 (predicts (prev+2) as center), spread=2,
	606	+ // so top-3 at each position includes (prev+1), (prev+2), (prev+3).
	607	+ // Target predicts (prev+1). So acceptance should be high.
	608	+ let target = MockModel::new(64);
	609	+ let draft = RangeModel::new(64, 0, 2);
	610	+ let params = SamplingParams {
	611	+ max_tokens: 40,
	612	+ temperature: 0.0,
	613	+ ..Default::default()
	614	+ };
	615	+ let prompt = vec![0u32];
	616	+
	617	+ let (_spec1, stats1) = collect_tree(&target, &draft, &prompt, &params, 4, 1);
	618	+ let (_spec2, stats2) = collect_tree(&target, &draft, &prompt, &params, 4, 3);
	619	+
	620	+ // branch=3 should have higher or equal acceptance than branch=1
	621	+ // (because it has more candidates per depth).
	622	+ assert!(
	623	+ stats2.acceptance_rate() >= stats1.acceptance_rate() - 1e-9,
	624	+ "branch=3 ({}) < branch=1 ({})",
	625	+ stats2.acceptance_rate(),
	626	+ stats1.acceptance_rate()
	627	+ );
	628	+ }
	629	+
	630	+ // ---------- honors max_tokens ----------
	631	+
	632	+ #[test]
	633	+ fn tree_honors_max_tokens() {
	634	+ let target = MockModel::new(64);
	635	+ let draft = RangeModel::new(64, 0, 2);
	636	+ let params = SamplingParams {
	637	+ max_tokens: 15,
	638	+ temperature: 0.0,
	639	+ ..Default::default()
	640	+ };
	641	+ let (spec, _) = collect_tree(&target, &draft, &[0u32], &params, 4, 3);
	642	+ assert_eq!(spec.len(), 15);
	643	+ }
	644	+
	645	+ // ---------- respects stop tokens ----------
	646	+
	647	+ #[test]
	648	+ fn tree_respects_stop_token() {
	649	+ let target = MockModel::new(32);
	650	+ let draft = MockModel::new(32);
	651	+ // MockModel rule is (prev+1)%32; stop at the next token after prompt.
	652	+ let next_after_prompt = (0u32.wrapping_add(1)) % 32;
	653	+ let params = SamplingParams {
	654	+ max_tokens: 50,
	655	+ temperature: 0.0,
	656	+ stop_tokens: vec![next_after_prompt],
	657	+ ..Default::default()
	658	+ };
	659	+ let (spec, _) = collect_tree(&target, &draft, &[0u32], &params, 3, 2);
	660	+ assert_eq!(spec.len(), 1);
	661	+ assert_eq!(spec[0], next_after_prompt);
	662	+ }
	663	+
	664	+ // ---------- verify_tree cases ----------
	665	+
	666	+ #[test]
	667	+ fn verify_accepts_when_child_matches() {
	668	+ let target = MockModel::new(16);
	669	+ // Context: [0]. Target predicts 1.
	670	+ let tree = DraftTree {
	671	+ roots: vec![TreeNode::new(1, 5.0), TreeNode::new(2, 4.0)],
	672	+ gamma: 1,
	673	+ branch: 2,
	674	+ };
	675	+ let params = SamplingParams {
	676	+ temperature: 0.0,
	677	+ ..Default::default()
	678	+ };
	679	+ let mut rng = StdRng::seed_from_u64(0);
	680	+ let (accepted, bonus, n) = verify_tree(&target, &tree, &[0u32], &params, &mut rng).unwrap();
	681	+ assert_eq!(accepted, vec![1]);
	682	+ assert_eq!(n, 1);
	683	+ assert!(bonus.is_some());
	684	+ }
	685	+
	686	+ #[test]
	687	+ fn verify_rejects_when_no_child_matches() {
	688	+ let target = MockModel::new(16);
	689	+ // Context: [0]. Target predicts 1. Draft proposes [2, 3] — no match.
	690	+ let tree = DraftTree {
	691	+ roots: vec![TreeNode::new(2, 5.0), TreeNode::new(3, 4.0)],
	692	+ gamma: 1,
	693	+ branch: 2,
	694	+ };
	695	+ let params = SamplingParams {
	696	+ temperature: 0.0,
	697	+ ..Default::default()
	698	+ };
	699	+ let mut rng = StdRng::seed_from_u64(0);
	700	+ let (accepted, bonus, n) = verify_tree(&target, &tree, &[0u32], &params, &mut rng).unwrap();
	701	+ assert_eq!(accepted.len(), 0);
	702	+ assert_eq!(n, 0);
	703	+ // Bonus should be 1 (target's prediction).
	704	+ assert_eq!(bonus, Some(1));
	705	+ }
	706	+
	707	+ // ---------- multi-depth tree verified correctly ----------
	708	+
	709	+ #[test]
	710	+ fn verify_multiple_depths() {
	711	+ // Context [0]. Target predicts 1, then 2, then 3 (MockModel chain).
	712	+ // Draft tree has gamma=3, branch=2, and one path matches perfectly.
	713	+ let mut node_a = TreeNode::new(1, 5.0);
	714	+ node_a.children = vec![
	715	+ TreeNode::new(2, 5.0), // matches target
	716	+ TreeNode::new(9, 4.0),
	717	+ ];
	718	+ node_a.children[0].children = vec![
	719	+ TreeNode::new(3, 5.0), // matches target
	720	+ TreeNode::new(8, 4.0),
	721	+ ];
	722	+ let mut node_b = TreeNode::new(7, 3.0);
	723	+ node_b.children = vec![TreeNode::new(8, 3.0)];
	724	+
	725	+ let tree = DraftTree {
	726	+ roots: vec![node_a, node_b],
	727	+ gamma: 3,
	728	+ branch: 2,
	729	+ };
	730	+ let target = MockModel::new(16);
	731	+ let params = SamplingParams {
	732	+ temperature: 0.0,
	733	+ ..Default::default()
	734	+ };
	735	+ let mut rng = StdRng::seed_from_u64(0);
	736	+ let (accepted, bonus, n) = verify_tree(&target, &tree, &[0u32], &params, &mut rng).unwrap();
	737	+ // Accepted: [1, 2, 3] — the matching path.
	738	+ assert_eq!(accepted, vec![1, 2, 3]);
	739	+ assert_eq!(n, 3);
	740	+ // Bonus: target predicts 4 at context [0,1,2,3].
	741	+ assert_eq!(bonus, Some(4));
	742	+ }
	743	+
	744	+ #[test]
	745	+ fn verify_deeply_rejected() {
	746	+ // Context [0]. Some match at depth 0, then reject at depth 1.
	747	+ let mut node_a = TreeNode::new(1, 5.0);
	748	+ node_a.children = vec![
	749	+ TreeNode::new(9, 5.0), // doesn't match (target expects 2)
	750	+ TreeNode::new(10, 4.0),
	751	+ ];
	752	+
	753	+ let tree = DraftTree {
	754	+ roots: vec![node_a],
	755	+ gamma: 2,
	756	+ branch: 2,
	757	+ };
	758	+ let target = MockModel::new(16);
	759	+ let params = SamplingParams {
	760	+ temperature: 0.0,
	761	+ ..Default::default()
	762	+ };
	763	+ let mut rng = StdRng::seed_from_u64(0);
	764	+ let (accepted, bonus, n) = verify_tree(&target, &tree, &[0u32], &params, &mut rng).unwrap();
	765	+ assert_eq!(accepted, vec![1]); // only first token accepted
	766	+ assert_eq!(n, 1);
	767	+ // Correction: target at [0,1] predicts 2.
	768	+ assert_eq!(bonus, Some(2));
	769	+ }
	770	+
	771	+ // ---------- acceptance stats are recorded correctly ----------
	772	+
	773	+ #[test]
	774	+ fn tree_stats_record_branch() {
	775	+ let target = MockModel::new(32);
	776	+ let draft = MockModel::new(32);
	777	+ let params = SamplingParams {
	778	+ max_tokens: 10,
	779	+ temperature: 0.0,
	780	+ ..Default::default()
	781	+ };
	782	+ let (_, stats) = collect_tree(&target, &draft, &[0u32], &params, 3, 3);
	783	+ assert_eq!(stats.tree_branch, 3);
	784	+ assert!(stats.steps > 0);
	785	+ }
	786	+
	787	+ // ---------- tree and draft-target produce same output (both lossless) ----------
	788	+
	789	+ #[test]
	790	+ fn tree_and_draft_target_produce_same_output() {
	791	+ // When branch=1, tree is essentially the same as single-path draft-target
	792	+ // (both produce the same tokens as plain decode).
	793	+ let target = MockModel::new(32);
	794	+ let draft = MockModel::new(32);
	795	+ let params = SamplingParams {
	796	+ max_tokens: 20,
	797	+ temperature: 0.0,
	798	+ ..Default::default()
	799	+ };
	800	+ let prompt = vec![0u32];
	801	+ let baseline = collect(&target, &prompt, &params);
	802	+
	803	+ let (tree_spec, _) = collect_tree(&target, &draft, &prompt, &params, 4, 1);
	804	+ assert_eq!(baseline, tree_spec);
	805	+
	806	+ let (tree_spec3, _) = collect_tree(&target, &draft, &prompt, &params, 4, 3);
	807	+ assert_eq!(baseline, tree_spec3);
	808	+ }
	809	+
	810	+ // ---------- empty prompt ----------
	811	+
	812	+ #[test]
	813	+ fn tree_handles_empty_prompt() {
	814	+ let target = MockModel::new(32);
	815	+ let draft = MockModel::new(32);
	816	+ let params = SamplingParams {
	817	+ max_tokens: 5,
	818	+ temperature: 0.0,
	819	+ ..Default::default()
	820	+ };
	821	+ let (spec, _) = collect_tree(&target, &draft, &[], &params, 2, 2);
	822	+ assert_eq!(spec.len(), 5);
	823	+ }
	824	+
	825	+ // ---------- TreeSpeculativeModel wrapper works as a Model ----------
	826	+
	827	+ #[test]
	828	+ fn tree_speculative_model_wrapper_matches_target() {
	829	+ // Wrap MockModel as a tree-speculative model with a perfect draft
	830	+ // (the draft is the same MockModel, so every path matches) and confirm
	831	+ // the wrapper emits the same token stream as plain target decode — the
	832	+ // property required for it to be servable through Engine::generate.
	833	+ let target = Arc::new(MockModel::new(32));
	834	+ let draft = Arc::new(MockModel::new(32));
	835	+ let spec = TreeSpeculativeModel::new(target.clone(), draft, 4, 3, 0);
	836	+
	837	+ let params = SamplingParams {
	838	+ max_tokens: 20,
	839	+ temperature: 0.0,
	840	+ ..Default::default()
	841	+ };
	842	+
	843	+ let plain = collect(&*target, &[3u32], &params);
	844	+ let wrapped = collect(&spec, &[3u32], &params);
	845	+ assert_eq!(plain, wrapped);
	846	+ assert_eq!(wrapped.len(), 20);
	847	+ }
	848	+
	849	+ #[test]
	850	+ fn tree_speculative_model_vocab_and_logits_delegated() {
	851	+ let target = Arc::new(MockModel::new(64));
	852	+ let draft = Arc::new(MockModel::new(64));
	853	+ let spec = TreeSpeculativeModel::new(target.clone(), draft, 4, 2, 0);
	854	+ assert_eq!(spec.vocab_size(), 64);
	855	+ // forward_logits should be delegated to the target.
	856	+ let ctx = vec![1u32, 2, 3];
	857	+ let target_logits = target.forward_logits(&ctx).unwrap();
	858	+ let spec_logits = spec.forward_logits(&ctx).unwrap();
	859	+ assert_eq!(target_logits, spec_logits);
	860	+ }
	861	+}
0	862

▼crates/justapi-inference/src/supervisor.rs

+355-0

@@ -0,0 +1,355 @@
	1	+//! Multi-replica supervisor — the lifecycle/orchestration layer.
	2	+//!
	3	+//! Ties the [`Autoscaler`] (how many replicas we want) and the [`Router`]
	4	+//! (which replicas can serve a request) to actual replicas. It is pure decision
	5	+//! logic: rather than spawning processes, it emits [`SupervisorAction`] intents
	6	+//! (start/stop a replica) that an external executor fulfills — that executor is
	7	+//! where real process/GPU management lives.
	8	+//!
	9	+//! Lifecycle:
	10	+//! 1. `reconcile(metrics)` asks the autoscaler for a desired replica count and
	11	+//! emits `StartReplica` / `StopReplica` intents to close the gap with the
	12	+//! live set, then syncs the router's replica view.
	13	+//! 2. Newly started replicas register in the router as not-yet-healthy; the
	14	+//! executor calls `on_replica_started` once the replica is actually serving,
	15	+//! flipping it healthy so the router will route to it.
	16	+//! 3. `on_replica_report` feeds live load metrics back in; `on_replica_stopped`
	17	+//! removes a replica the executor terminated.
	18	+
	19	+use std::collections::HashSet;
	20	+
	21	+use crate::engine::EngineDevice;
	22	+use crate::{Autoscaler, AutoscalerConfig, LlmMetrics, Replica, Router, RoutingStrategy};
	23	+
	24	+/// The supervisor's view of a single live replica.
	25	+#[derive(Debug, Clone)]
	26	+pub struct LiveReplica {
	27	+ pub id: String,
	28	+ /// Whether the replica is healthy and accepting traffic.
	29	+ pub healthy: bool,
	30	+ /// Current active sequences (drives router load-aware selection).
	31	+ pub active_sequences: usize,
	32	+ /// KV-cache pressure (0–100).
	33	+ pub kv_pressure_pct: f32,
	34	+}
	35	+
	36	+/// An intent the supervisor emits for the executor to fulfill.
	37	+#[derive(Debug, Clone)]
	38	+pub enum SupervisorAction {
	39	+ /// Start a new replica for `model` @ `version`.
	40	+ StartReplica {
	41	+ id: String,
	42	+ model: String,
	43	+ version: String,
	44	+ },
	45	+ /// Terminate the replica with `id`.
	46	+ StopReplica { id: String },
	47	+}
	48	+
	49	+/// Tuning for the supervisor (mirrors the autoscaler knobs plus replica shape).
	50	+#[derive(Debug, Clone)]
	51	+pub struct SupervisorConfig {
	52	+ /// Model this supervisor manages.
	53	+ pub model: String,
	54	+ /// Version of the model this supervisor manages.
	55	+ pub version: String,
	56	+ /// Device new replicas are pinned to.
	57	+ pub device: EngineDevice,
	58	+ pub min_replicas: usize,
	59	+ pub max_replicas: usize,
	60	+ /// Max concurrent sequences per replica (advertised to the router).
	61	+ pub max_concurrency_per_replica: usize,
	62	+ pub target_ttft_ms: f32,
	63	+ pub target_kv_pressure_pct: f32,
	64	+ pub queue_depth_scale_up: usize,
	65	+ pub queue_depth_scale_down: usize,
	66	+ /// Ticks to wait after a scaling action before another (flap protection).
	67	+ pub cooldown_ticks: usize,
	68	+}
	69	+
	70	+/// Multi-replica supervisor: owns an [`Autoscaler`] + [`Router`] and drives a
	71	+/// live replica set toward the autoscaler's desired count.
	72	+pub struct Supervisor {
	73	+ config: SupervisorConfig,
	74	+ autoscaler: Autoscaler,
	75	+ router: Router,
	76	+ live: Vec<LiveReplica>,
	77	+ next_index: usize,
	78	+}
	79	+
	80	+impl Supervisor {
	81	+ /// Create a supervisor for `config.model`@`config.version`. The embedded
	82	+ /// autoscaler seeds at `min_replicas`.
	83	+ pub fn new(config: SupervisorConfig) -> Self {
	84	+ let autoscaler = Autoscaler::new(AutoscalerConfig {
	85	+ min_replicas: config.min_replicas,
	86	+ max_replicas: config.max_replicas,
	87	+ target_ttft_ms: config.target_ttft_ms,
	88	+ target_kv_pressure_pct: config.target_kv_pressure_pct,
	89	+ queue_depth_scale_up: config.queue_depth_scale_up,
	90	+ queue_depth_scale_down: config.queue_depth_scale_down,
	91	+ cooldown_ticks: config.cooldown_ticks,
	92	+ });
	93	+ Self {
	94	+ config,
	95	+ autoscaler,
	96	+ router: Router::new(RoutingStrategy::LeastLoaded),
	97	+ live: Vec::new(),
	98	+ next_index: 0,
	99	+ }
	100	+ }
	101	+
	102	+ /// Current desired replica count (from the autoscaler).
	103	+ pub fn desired(&self) -> usize {
	104	+ self.autoscaler.current()
	105	+ }
	106	+
	107	+ /// Number of live replicas (healthy or starting).
	108	+ pub fn live_count(&self) -> usize {
	109	+ self.live.len()
	110	+ }
	111	+
	112	+ /// Number of healthy live replicas.
	113	+ pub fn healthy_count(&self) -> usize {
	114	+ self.live.iter().filter(\|r\| r.healthy).count()
	115	+ }
	116	+
	117	+ /// Access the supervisor's router (so requests can be routed against the
	118	+ /// live replica set).
	119	+ pub fn router(&self) -> &Router {
	120	+ &self.router
	121	+ }
	122	+
	123	+ /// Reconcile the live replica set with the desired count derived from
	124	+ /// `metrics`. Emits start/stop intents and syncs the router.
	125	+ pub fn reconcile(&mut self, metrics: &LlmMetrics) -> Vec<SupervisorAction> {
	126	+ self.autoscaler.decide(metrics);
	127	+ let desired = self.autoscaler.current();
	128	+ let mut actions = Vec::new();
	129	+
	130	+ // Scale up: register not-yet-healthy replicas and emit start intents.
	131	+ while self.live.len() < desired && self.live.len() < self.config.max_replicas {
	132	+ let id = format!(
	133	+ "{}-{}-{}",
	134	+ self.config.model, self.config.version, self.next_index
	135	+ );
	136	+ self.next_index += 1;
	137	+ self.live.push(LiveReplica {
	138	+ id: id.clone(),
	139	+ healthy: false,
	140	+ active_sequences: 0,
	141	+ kv_pressure_pct: 0.0,
	142	+ });
	143	+ actions.push(SupervisorAction::StartReplica {
	144	+ id,
	145	+ model: self.config.model.clone(),
	146	+ version: self.config.version.clone(),
	147	+ });
	148	+ }
	149	+
	150	+ // Scale down: prefer removing stuck (unhealthy) replicas first.
	151	+ while self.live.len() > desired {
	152	+ let idx = self
	153	+ .live
	154	+ .iter()
	155	+ .position(\|r\| !r.healthy)
	156	+ .or_else(\|\| Some(self.live.len().saturating_sub(1)));
	157	+ if let Some(idx) = idx {
	158	+ let removed = self.live.remove(idx);
	159	+ actions.push(SupervisorAction::StopReplica { id: removed.id });
	160	+ } else {
	161	+ break;
	162	+ }
	163	+ }
	164	+
	165	+ self.sync_router();
	166	+ actions
	167	+ }
	168	+
	169	+ /// Executor reports a replica has finished starting and is now healthy.
	170	+ pub fn on_replica_started(&mut self, id: &str) {
	171	+ if let Some(r) = self.live.iter_mut().find(\|r\| r.id == id) {
	172	+ r.healthy = true;
	173	+ }
	174	+ self.sync_router();
	175	+ }
	176	+
	177	+ /// Executor reports live load for a replica.
	178	+ pub fn on_replica_report(&mut self, id: &str, active_sequences: usize, kv_pressure_pct: f32) {
	179	+ if let Some(r) = self.live.iter_mut().find(\|r\| r.id == id) {
	180	+ r.active_sequences = active_sequences;
	181	+ r.kv_pressure_pct = kv_pressure_pct;
	182	+ }
	183	+ self.sync_router();
	184	+ }
	185	+
	186	+ /// Executor reports a replica was terminated.
	187	+ pub fn on_replica_stopped(&mut self, id: &str) {
	188	+ self.live.retain(\|r\| r.id != id);
	189	+ self.router.remove_replica(id);
	190	+ }
	191	+
	192	+ /// Push the live replica set into the router (add/update live, drop gone).
	193	+ fn sync_router(&mut self) {
	194	+ let live_ids: HashSet<String> = self.live.iter().map(\|r\| r.id.clone()).collect();
	195	+ for r in self.router.replicas() {
	196	+ if !live_ids.contains(&r.id) {
	197	+ self.router.remove_replica(&r.id);
	198	+ }
	199	+ }
	200	+ for lr in &self.live {
	201	+ self.router.register_replica(Replica {
	202	+ id: lr.id.clone(),
	203	+ model: self.config.model.clone(),
	204	+ version: self.config.version.clone(),
	205	+ device: self.config.device,
	206	+ max_concurrency: self.config.max_concurrency_per_replica,
	207	+ active_sequences: lr.active_sequences,
	208	+ kv_pressure_pct: lr.kv_pressure_pct,
	209	+ healthy: lr.healthy,
	210	+ loaded_adapters: Vec::new(),
	211	+ });
	212	+ }
	213	+ }
	214	+}
	215	+
	216	+// ---------------------------------------------------------------------------
	217	+// Tests
	218	+// ---------------------------------------------------------------------------
	219	+
	220	+#[cfg(test)]
	221	+mod tests {
	222	+ use super::*;
	223	+
	224	+ fn cfg() -> SupervisorConfig {
	225	+ SupervisorConfig {
	226	+ model: "llama-7b".into(),
	227	+ version: "v1".into(),
	228	+ device: EngineDevice::Cuda(0),
	229	+ min_replicas: 1,
	230	+ max_replicas: 4,
	231	+ max_concurrency_per_replica: 8,
	232	+ target_ttft_ms: 500.0,
	233	+ target_kv_pressure_pct: 80.0,
	234	+ queue_depth_scale_up: 16,
	235	+ queue_depth_scale_down: 2,
	236	+ cooldown_ticks: 2,
	237	+ }
	238	+ }
	239	+
	240	+ fn metrics(kv: f32, ttft: f32, q: usize) -> LlmMetrics {
	241	+ LlmMetrics {
	242	+ tokens_per_sec: 50.0,
	243	+ ttft_ms: ttft,
	244	+ queue_depth: q,
	245	+ kv_pressure_pct: kv,
	246	+ }
	247	+ }
	248	+
	249	+ #[test]
	250	+ fn seeds_min_replicas_on_first_reconcile() {
	251	+ let mut s = Supervisor::new(cfg());
	252	+ // Healthy metrics → desired stays at min (1). No live yet → start one.
	253	+ let actions = s.reconcile(&metrics(10.0, 100.0, 0));
	254	+ assert_eq!(s.desired(), 1);
	255	+ assert_eq!(s.live_count(), 1);
	256	+ assert_eq!(actions.len(), 1);
	257	+ match &actions[0] {
	258	+ SupervisorAction::StartReplica { id, .. } => assert!(id.contains("llama-7b-v1")),
	259	+ _ => panic!("expected StartReplica"),
	260	+ }
	261	+ // Router sees the replica (not yet healthy → not routable).
	262	+ assert_eq!(s.router().replicas().len(), 1);
	263	+ assert_eq!(s.healthy_count(), 0);
	264	+ }
	265	+
	266	+ #[test]
	267	+ fn stays_stable_once_replicas_healthy() {
	268	+ let mut s = Supervisor::new(cfg());
	269	+ let actions = s.reconcile(&metrics(10.0, 100.0, 0));
	270	+ let id = match &actions[0] {
	271	+ SupervisorAction::StartReplica { id, .. } => id.clone(),
	272	+ _ => panic!(),
	273	+ };
	274	+ s.on_replica_started(&id);
	275	+ assert_eq!(s.healthy_count(), 1);
	276	+ // Another healthy reconcile → no new actions.
	277	+ let actions = s.reconcile(&metrics(10.0, 100.0, 0));
	278	+ assert!(actions.is_empty());
	279	+ assert_eq!(s.live_count(), 1);
	280	+ }
	281	+
	282	+ #[test]
	283	+ fn scales_up_under_load() {
	284	+ let mut s = Supervisor::new(cfg());
	285	+ // First establish the baseline (1 replica, mark healthy).
	286	+ let a = s.reconcile(&metrics(10.0, 100.0, 0));
	287	+ let id = match &a[0] {
	288	+ SupervisorAction::StartReplica { id, .. } => id.clone(),
	289	+ _ => panic!(),
	290	+ };
	291	+ s.on_replica_started(&id);
	292	+ // Now saturate: KV pressure high → autoscaler wants more.
	293	+ // Cooldown ticks pass between reconciles.
	294	+ let _ = s.reconcile(&metrics(90.0, 100.0, 0));
	295	+ let _ = s.reconcile(&metrics(90.0, 100.0, 0));
	296	+ let actions = s.reconcile(&metrics(90.0, 100.0, 0));
	297	+ // Either we already scaled up via the cooldown window or this emits starts.
	298	+ assert!(s.desired() >= 2);
	299	+ if !actions.is_empty() {
	300	+ assert!(actions
	301	+ .iter()
	302	+ .all(\|a\| matches!(a, SupervisorAction::StartReplica { .. })));
	303	+ }
	304	+ assert!(s.live_count() >= 2);
	305	+ }
	306	+
	307	+ #[test]
	308	+ fn never_exceeds_max() {
	309	+ let mut s = Supervisor::new(cfg());
	310	+ // Hammer with saturated metrics across many ticks; live must cap at max.
	311	+ for _ in 0..12 {
	312	+ let _ = s.reconcile(&metrics(99.0, 2000.0, 100));
	313	+ }
	314	+ assert_eq!(s.live_count(), 4);
	315	+ assert_eq!(s.desired(), 4);
	316	+ }
	317	+
	318	+ #[test]
	319	+ fn scales_down_under_low_load() {
	320	+ let mut c = cfg();
	321	+ c.max_replicas = 3;
	322	+ let mut s = Supervisor::new(c);
	323	+ // Bring up to 3 replicas.
	324	+ for _ in 0..12 {
	325	+ let _ = s.reconcile(&metrics(99.0, 2000.0, 100));
	326	+ }
	327	+ // Mark them all healthy.
	328	+ for r in s.router().replicas() {
	329	+ s.on_replica_started(&r.id);
	330	+ }
	331	+ assert_eq!(s.live_count(), 3);
	332	+ assert_eq!(s.healthy_count(), 3);
	333	+ // Now underused for many ticks → scale down to min.
	334	+ for _ in 0..12 {
	335	+ let _ = s.reconcile(&metrics(5.0, 50.0, 0));
	336	+ }
	337	+ assert_eq!(s.live_count(), 1);
	338	+ assert_eq!(s.healthy_count(), 1);
	339	+ }
	340	+
	341	+ #[test]
	342	+ fn router_sync_reflects_live_set() {
	343	+ let mut s = Supervisor::new(cfg());
	344	+ let _ = s.reconcile(&metrics(99.0, 2000.0, 100));
	345	+ let _ = s.reconcile(&metrics(99.0, 2000.0, 100));
	346	+ let _ = s.reconcile(&metrics(99.0, 2000.0, 100));
	347	+ // Router replica count must equal the live count (synced each reconcile).
	348	+ assert_eq!(s.router().replicas().len(), s.live_count());
	349	+ // After stopping one via the executor, the router drops it too.
	350	+ let id = s.router().replicas()[0].id.clone();
	351	+ s.on_replica_stopped(&id);
	352	+ assert_eq!(s.router().replicas().len(), s.live_count());
	353	+ assert!(!s.router().replicas().iter().any(\|r\| r.id == id));
	354	+ }
	355	+}
0	356

▼crates/justapi-inference/tests/engine.rs

+177-0

@@ -0,0 +1,177 @@
	1	+//! End-to-end tests for the inference engine foundation (Phase 41).
	2	+//! Runs on CPU with the weight-free `MockModel` — no GPU/weights required.
	3	+
	4	+use justapi_inference::{
	5	+ AcceptanceStats, Engine, EngineDevice, FinishReason, MockModel, Model, ModelError,
	6	+ SamplingParams,
	7	+};
	8	+
	9	+#[test]
	10	+fn mock_tokenize_detokenize_roundtrip() {
	11	+ let text = "hello";
	12	+ let ids = MockModel::tokenize(text);
	13	+ assert_eq!(ids, vec![104, 101, 108, 108, 111]);
	14	+ assert_eq!(MockModel::detokenize(&ids), text);
	15	+}
	16	+
	17	+#[test]
	18	+fn mock_generate_respects_max_tokens() {
	19	+ let model = MockModel::new(26);
	20	+ let params = SamplingParams {
	21	+ max_tokens: 10,
	22	+ ..Default::default()
	23	+ };
	24	+ let count = std::cell::Cell::new(0usize);
	25	+ let reason = model
	26	+ .generate(&[0], &params, &\|t\| {
	27	+ count.set(count.get() + 1);
	28	+ assert!(!t.text.is_empty());
	29	+ true
	30	+ })
	31	+ .unwrap();
	32	+ assert_eq!(count.get(), 10);
	33	+ assert_eq!(reason, FinishReason::Length);
	34	+}
	35	+
	36	+#[test]
	37	+fn mock_generate_stops_on_stop_token() {
	38	+ let model = MockModel::new(26);
	39	+ // Starting from prompt [0], next is 1 ('b'); stop when we produce token 3.
	40	+ let params = SamplingParams {
	41	+ max_tokens: 100,
	42	+ stop_tokens: vec![3],
	43	+ ..Default::default()
	44	+ };
	45	+ let produced: Vec<u32> = Vec::new();
	46	+ let produced = std::sync::Mutex::new(produced);
	47	+ let reason = model
	48	+ .generate(&[0], &params, &\|t\| {
	49	+ produced.lock().unwrap().push(t.id);
	50	+ true
	51	+ })
	52	+ .unwrap();
	53	+ assert_eq!(reason, FinishReason::Stop);
	54	+ assert!(produced.lock().unwrap().contains(&3));
	55	+}
	56	+
	57	+#[test]
	58	+fn device_discover_includes_cpu() {
	59	+ let devices = EngineDevice::discover();
	60	+ assert!(devices.contains(&EngineDevice::Cpu));
	61	+}
	62	+
	63	+#[test]
	64	+fn engine_registers_and_lists_mock() {
	65	+ let engine = Engine::new(EngineDevice::Cpu).unwrap();
	66	+ assert!(engine.list_models().is_empty());
	67	+ engine.register_mock("demo");
	68	+ let models = engine.list_models();
	69	+ assert_eq!(models, vec!["demo".to_string()]);
	70	+}
	71	+
	72	+#[tokio::test]
	73	+async fn engine_generate_streams_tokens() {
	74	+ let engine = Engine::new(EngineDevice::Cpu).unwrap();
	75	+ engine.register_mock("demo");
	76	+ let params = SamplingParams {
	77	+ max_tokens: 8,
	78	+ ..Default::default()
	79	+ };
	80	+ let mut rx = engine.generate("demo", &[0], params).unwrap();
	81	+
	82	+ let mut count = 0;
	83	+ let mut last_reason = None;
	84	+ while let Some(tok) = rx.recv().await {
	85	+ count += 1;
	86	+ last_reason = tok.finish_reason;
	87	+ }
	88	+ assert_eq!(count, 8);
	89	+ // Stream ended by length; no final token carries a finish reason in mock.
	90	+ assert!(last_reason.is_none() \|\| last_reason == Some(FinishReason::Length));
	91	+}
	92	+
	93	+#[test]
	94	+fn engine_generate_unknown_model_errors() {
	95	+ let engine = Engine::new(EngineDevice::Cpu).unwrap();
	96	+ let err = engine.generate("nope", &[0], SamplingParams::default());
	97	+ assert!(matches!(err, Err(ModelError::NotFound(_))));
	98	+}
	99	+
	100	+#[test]
	101	+#[cfg(not(feature = "real"))]
	102	+fn load_without_real_feature_is_explicit_error() {
	103	+ let engine = Engine::new(EngineDevice::Cpu).unwrap();
	104	+ // Without the `real` feature this must fail loudly, not silently no-op.
	105	+ let err = engine.load("m", std::path::Path::new("/tmp/nonexistent-model"));
	106	+ assert!(matches!(err, Err(ModelError::FeatureRequired("real"))));
	107	+}
	108	+
	109	+#[test]
	110	+#[cfg(feature = "real")]
	111	+fn load_with_real_feature_attempts_load() {
	112	+ let engine = Engine::new(EngineDevice::Cpu).unwrap();
	113	+ // With the `real` feature, load attempts to read the directory and fails
	114	+ // with a concrete error (not FeatureRequired).
	115	+ let err = engine.load("m", std::path::Path::new("/tmp/nonexistent-model"));
	116	+ assert!(err.is_err());
	117	+ assert!(!matches!(err, Err(ModelError::FeatureRequired(_))));
	118	+}
	119	+
	120	+#[test]
	121	+fn engine_speculative_decoding_end_to_end() {
	122	+ // Register a speculative model served through the normal Engine path:
	123	+ // target == draft (perfect draft) → acceptance rate 1.0 and gamma+1 tokens
	124	+ // emitted per verify step. Streamed over the same tokio mpsc as plain gen.
	125	+ let engine = Engine::new(EngineDevice::Cpu).unwrap();
	126	+ let target = std::sync::Arc::new(MockModel::new(32));
	127	+ let draft = std::sync::Arc::new(MockModel::new(32));
	128	+ engine.register_speculative("spec", target, draft, 4, 7);
	129	+
	130	+ let params = SamplingParams {
	131	+ max_tokens: 20,
	132	+ temperature: 0.0,
	133	+ ..Default::default()
	134	+ };
	135	+
	136	+ // Collect the speculatively-decoded stream.
	137	+ let mut spec_rx = engine.generate("spec", &[0], params.clone()).unwrap();
	138	+ let mut spec_ids = Vec::new();
	139	+ while let Some(tok) = spec_rx.blocking_recv() {
	140	+ spec_ids.push(tok.id);
	141	+ }
	142	+ assert_eq!(spec_ids.len(), 20);
	143	+
	144	+ // Correctness: must equal plain target decode (speculation is lossless).
	145	+ engine.register("__plain", std::sync::Arc::new(MockModel::new(32)));
	146	+ let mut plain_rx = engine.generate("__plain", &[0], params).unwrap();
	147	+ let mut plain_ids = Vec::new();
	148	+ while let Some(tok) = plain_rx.blocking_recv() {
	149	+ plain_ids.push(tok.id);
	150	+ }
	151	+ assert_eq!(
	152	+ spec_ids, plain_ids,
	153	+ "speculation must not change the output"
	154	+ );
	155	+}
	156	+
	157	+#[test]
	158	+fn acceptance_rate_perfect_draft_is_one() {
	159	+ // Direct verify of AcceptanceStats with a perfect draft.
	160	+ let target = MockModel::new(32);
	161	+ let draft = MockModel::new(32);
	162	+ let params = SamplingParams {
	163	+ max_tokens: 30,
	164	+ temperature: 0.0,
	165	+ ..Default::default()
	166	+ };
	167	+ let out = std::cell::RefCell::new(Vec::new());
	168	+ let (_finish, stats) =
	169	+ justapi_inference::speculative_generate(&target, &draft, &[0], &params, 4, 11, &\|t\| {
	170	+ out.borrow_mut().push(t.id);
	171	+ true
	172	+ })
	173	+ .unwrap();
	174	+ assert!((stats.acceptance_rate() - 1.0).abs() < 1e-9);
	175	+ assert_eq!(stats.tokens_emitted(), out.borrow().len());
	176	+ let _ = AcceptanceStats::default();
	177	+}
0	178